diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml new file mode 100644 index 0000000000000000000000000000000000000000..59c2437a85b44618ac680ba7f277a880a5d3e99a --- /dev/null +++ b/.github/workflows/ci.yml @@ -0,0 +1,75 @@ +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 + + 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 diff --git a/.gitignore b/.gitignore index f99b49e07aea8444d8bc57f2e00889b42a756a91..60d736279b7e9d9abe5da4cb6bdcef746e4b4388 100644 --- a/.gitignore +++ b/.gitignore @@ -4,3 +4,29 @@ scores/ slurm/ .onnx .png + +# Python tooling +*.py[cod] +.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/raw/* +!data/raw/.gitkeep +data/processed/ +outputs/ +checkpoints/ + +# Keep the legacy implementation immutable by convention; do not add generated files there. diff --git a/.python-version b/.python-version new file mode 100644 index 0000000000000000000000000000000000000000..e4fba2183587225f216eeada4c78dfab6b2e65f5 --- /dev/null +++ b/.python-version @@ -0,0 +1 @@ +3.12 diff --git a/AGENTS.md b/AGENTS.md new file mode 100644 index 0000000000000000000000000000000000000000..9182e1d7d3304edd3c50f247a381bcde1ebb66a8 --- /dev/null +++ b/AGENTS.md @@ -0,0 +1,477 @@ +# GNN4Colliders Agent Instructions + +## Project purpose + +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//.** + +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 +data +model +task +trainer +environment +``` + +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. + +--- + +## CLI design + +The intended user-facing interface is a single project CLI with subcommands conceptually similar to: + +```bash +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/ +``` + +### 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 + +Do not hide failed checks or incomplete behavior. diff --git a/CHANGELOG.md b/CHANGELOG.md new file mode 100644 index 0000000000000000000000000000000000000000..87c90d8900d3454ece4837d862f6707b90919d00 --- /dev/null +++ b/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__`. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000000000000000000000000000000000000..d44958299e9105ae46d7ef8861b6a7106f8b6adb --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1,3 @@ +include README.md +include CHANGELOG.md +recursive-include docs *.md diff --git a/README.md b/README.md index e5e004a4ae410d57a4be172387fa1036ff0f142b..fe009588e8caf2e25a8aea36a69292a05489a5b8 100644 --- a/README.md +++ b/README.md @@ -1,358 +1,252 @@ ---- -license: mit -tags: -- arXiv:2412.10665 ---- - -This is a demo is of the approach described in the paper, ["Pretrained Event Classification Model for High Energy Physics Analysis"](https://arxiv.org/abs/2412.10665) -``` -@misc{ho2024pretrained, - title={Pretrained Event Classification Model for High Energy Physics Analysis}, - author={Joshua Ho, Benjamin Ryan Roberts, Shuo Han, Haichen Wang}, - year={2024}, - eprint={2412.10665}, - archivePrefix={arXiv} -} -``` - -## Abstract - -We introduce a foundation model for event classification in high-energy physics, built on a **Graph Neural Network** architecture and trained on **120 million simulated proton-proton collision events** spanning 12 distinct physics processes. The model is *pretrained* to learn a general and robust representation of collision data using challenging multiclass and multilabel classification tasks. - -Its performance is evaluated across five event classification tasks, which include both physics processes used during pretraining and new processes not encountered during pretraining. Fine-tuning the pretrained model significantly improves classification performance, particularly in scenarios with limited training data, demonstrating gains in both accuracy and computational efficiency. - -To investigate the underlying mechanisms behind these performance improvements, we employ a representational similarity evaluation framework based on *Centered Kernel Alignment*. This analysis reveals notable differences in the learned representations of fine-tuned pretrained models compared to baseline models trained from scratch. - -## Introduction - -Machine learning has become a ubiquitous tool in particle physics, employed in a variety of tasks including triggering, simulation, reconstruction, and offline analysis. While its utility spans classification, regression, and generative tasks, the current paradigm of developing machine learning models from scratch for each specific application presents several challenges. This approach not only demands specialized expertise and substantial computing resources but can also result in suboptimal performance due to limited training data. The from-scratch development of models necessitates individual validation studies to ensure that neural networks utilize well-modeled information from training samples, whether derived from Monte Carlo simulations or control samples from experimental data. - -Foundation models offer a promising direction to address these limitations. These models, pre-trained on large, diverse datasets across various tasks, provide robust and general representations of underlying data structures. Notable examples in other fields include GPT-4 [OpenAI et al., 2024](#ref-openai-2024-gpt4) and BERT [Devlin et al., 2018](#ref-devlin-2018-bert) in natural language processing, Stable Diffusion [Rombach et al., 2021](#ref-rombach-2021-latentdiffusion) in image processing, and AlphaFold [Jumper et al., 2021](#ref-jumper-2021-alphafold) in structural biology. The foundation model approach offers several advantages for particle physics applications: reduced computing resources for fine-tuning [Yosinski et al., 2014](#ref-yosinski-2014-transfer) compared to training from scratch, superior performance on specific tasks (particularly with limited training data), and potentially simplified validation procedures as downstream tasks inherit verified representations from the pre-trained model. - -Current literature on pretrained models for particle physics can be categorized based on the data representation they handle. Models operating on particle- or event-level numerical data use features like particle four momenta or jets, leveraging self-supervised or generative methods to learn versatile representations. Detector-focused models operate on high-dimensional responses such as calorimeter deposits or pixel hits, employing geometry-aware techniques for accurate simulation and analysis. Finally, models using textual or code representations apply large language model architectures to integrate domain knowledge, enabling tasks like question answering and code generation. - -Recent studies have begun exploring foundation models tailored to particle physics data, which has a variety of distinct structures and properties across many experiments and data processing stages, including: - -- particle-level & event-level numeric data [Wildridge et al., 2024](#ref-wildridge-2024-bumblebee), [Katel et al., 2024](#ref-katel-2024-jet), [Golling et al., 2024](#ref-golling-2024-maskedset), [Mikuni & Nachman, 2024](#ref-mikuni-2024-omnilearn), [Harris et al., 2024](#ref-harris-2024-resimulation), [Birk et al., 2024](#ref-birk-2024-omnijet), [Vigl et al., 2024](#ref-vigl-2024-finetune), -- detector-level & geometry-aware data [Araz et al., 2024](#ref-araz-2024-pointcloud), [Liu et al., 2023](#ref-liu-2023-gaam), [Hashemi et al., 2024](#ref-hashemi-2024-gen), [Huang et al., 2024](#ref-huang-2024-lmtracking), -- textual or code data [Zhang et al., 2024](#ref-zhang-2024-xiwu). - -This paper presents a foundation model designed specifically for collider event-level data. In modern collider experiments, final-stage analysis processes information from reconstructed objects that either directly correspond to particles in collision final states (such as leptons and photons) or serve as proxies (such as jets and missing transverse energy). While traditional approaches often relied on "high-level" variables calculated from object features, recent trends favor direct input of event objects and their features into neural networks for analysis tasks. A notable example is [ATLAS Collaboration, 2023](#ref-atlas-2023-4top), which established the observation of simultaneous production of four top quarks with the ATLAS experiment by employing a graph neural network (GNN) architecture to process event-level object information. - -We present foundation models that adopt an architecture similar to that used for [ATLAS Collaboration, 2023](#ref-atlas-2023-4top). Our models are pre-trained using either multiclass classification or multi-label learning tasks across 12 distinct physics processes. We evaluate these models through fine-tuning and testing on five classification tasks, including both familiar and novel processes not seen during pre-training. Our analysis benchmarks the models' performance improvements, their scaling behavior with training sample size, and computational efficiency, representing the first prototype of a foundation model operating on collider final-state object data. - -## Data Samples - -To provide a diverse set of physics processes for the pretraining, we use Madgraph@NLO 2.7.3 [Alwall et al., 2014](#ref-alwall-2014hca) to generate proton-proton collision events at next-to-leading order (NLO) in Quantum Chromodynamics (QCD). We generate 12 distinct Standard Model (SM) physics processes, including six major Higgs boson production mechanisms: gluon fusion production \\(ggF\\), vector boson fusion \\(VBF\\), associated production of the Higgs boson with a W boson \\(WH\\) or a Z boson \\(ZH\\), associated production of the Higgs boson with a top-quark pair \\(t\bar{t}H\\), and associated production of the Higgs boson with a single top quark and a forward quark \\(tHq\\). Additionally, we simulate six top quark production processes: single top production, top-quark pair production \\(t\bar{t}\\), top quark pair production in association with a pair of photons \\(t\bar{t}\gamma\gamma\\), associated production of a top-quark pair with a W boson \\(t\bar{t}W\\), simultaneous production of three top quarks \\(t\bar{t}t\\), and simultaneous production of four top quarks \\(t\bar{t}t\bar{t}\\). In these samples, the Higgs boson and top quarks decay inclusively. These 12 Higgs and top quark production processes constitute the pretraining dataset. - -To test the pretrained model, we further generated four processes including three beyond Standard Model (SM) processes: a SM \\(t\bar{t}H\\) production where the Higgs boson decays exclusively to a pair of photons, a \\(t\bar{t}H\\) production with the Higgs boson decaying to a pair of photons, where the top-Yukawa coupling is CP-odd, implemented using the Higgs Characterization model [Artoisenet et al., 2013](#ref-artoisinet-2013puc), the production of a pair of superpartners of the top quark (s-top) using the Minimal Supersymmetric Standard Model (MSSM) [Rosiek, 1990](#ref-rosiek-1990), [Allanach et al., 2009](#ref-allanach-2009), and flavor changing neutral current (FCNC) processes [Degrande et al., 2015](#ref-degrande-2015), [Durieux et al., 2015](#ref-durieux-2015). For the s-top process, we simulate the production of heavier s-top pairs \\(t_2\bar{t_2}\\), where each heavier s-top (mass 582 GeV) decays into a lighter s-top \\(t_1\\) or \\(\bar{t_1}\\), mass 400 GeV) and a Higgs boson. The FCNC process involves \\(t\bar{t}\\) production where one top quark decays to a Higgs boson and a light quark. We generate 10 million events for each process, except for \\(tHq\\) and \\(t\bar{t}t\bar{t}\\), where 5 million events were produced. - -In all simulation samples, the center of mass energy of the proton-proton collision is set to 13 TeV. The Higgs boson, top quarks, and vector bosons are set to decay inclusively (except the \\(t\bar{t}H \rightarrow \gamma\gamma\\) samples), with MadSpin [Artoisenet et al., 2012](#ref-artoisinet-2012st) handling the decays of top quarks and W bosons. The generated events are processed through Pythia 8.235 [Sjostrand et al., 2015](#ref-sjostrand-2015) for parton showering and heavy particle decays, followed by Delphes 3.4.2 [de Favereau et al., 2014](#ref-defavereau-2014) configured to emulate the ATLAS detector [ATLAS Collaboration, 2008](#ref-atlas-2008) for fast detector simulation. - -The detector-level object selection criteria are defined to align with typical experimental conditions. Photons are required to have transverse momentum \\(p_T \geq 20~\mathrm{GeV}\\) and pseudorapidity \\(|\eta| \leq 2.37\\), excluding the electromagnetic calorimeter crack region \\(1.37 < |\eta| < 1.52\\). Electrons must have \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.47\\) (excluding the same crack region), while muons are selected with \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.7\\). Jets are reconstructed using the anti-\\(k_t\\) algorithm [Cacciari et al., 2008](#ref-cacciari-2008gp) with radius parameter \\(\Delta R=0.4\\), where \\(\Delta R\\) is defined as \\(\sqrt{\Delta\eta ^2 + \Delta\phi^2}\\), with \\(\Delta\eta\\) being the difference in pseudorapidity and \\(\Delta\phi\\) the difference in azimuthal angle. Jets must satisfy \\(p_T \geq 25~\mathrm{GeV}\\) and \\(|\eta| \leq 2.5\\). To avoid double-counting, jets are removed if they are within \\(\Delta R < 0.4\\) of a photon or lepton. The identification of jets originating from b-quark decays (b-tagging) is performed by matching jets within \\(\Delta R = 0.4\\) of a b-quark, with efficiency corrections applied to match the performance of the ATLAS experiment's b-tagging algorithm [ATLAS Collaboration, 2019](#ref-atlas-2019bwq). - -## Methods - -### Overview - -We present a methodology for developing and evaluating a foundation model for particle collision event analysis. The approach centers on pretraining a Graph Neural Network (GNN) architecture using a comprehensive dataset that spans multiple physics tasks, enabling the model to learn robust and transferable features. For task-specific applications, we employ a fine-tuning strategy that combines output layer adaptation with carefully calibrated learning rates for updating the pretrained parameters. - -Given the prevalence of classification problems in particle physics data analysis, we evaluate the model's efficacy through a systematic assessment across five binary classification tasks: - -- \\(t\bar{t}H(\rightarrow \gamma\gamma)\\) with CP-even versus CP-odd t-H interaction -- \\(t\bar{t}\\) with FCNC top quark decays versus $tHq$ processes -- \\(t\bar{t}W\\) versus $ttt$ processes -- Stop pair production with Higgs bosons in the decay chain versus \\(t\bar{t}H\\) processes -- \\(WH\\) versus \\(ZH\\) production modes - -Our evaluation metrics encompass classification performance, computational efficiency, and model interpretability. The investigation extends to analyzing the model's scaling behavior with respect to training dataset size, benchmarked against models trained without pretraining. Although we explored transfer learning through parameter freezing of pretrained layers, this approach did not yield performance improvements, leading us to focus our detailed analysis on fine-tuning strategies. - -This methodological framework demonstrates the potential of foundation models to enhance the efficiency of particle physics analyses while improving task-specific performance, offering a promising direction for future high-energy physics research. - ---- - -### GNN Architecture - -We implement a Graph Neural Network (GNN) architecture that naturally accommodates the point-cloud structure of particle physics data, employing the DGL framework with a PyTorch backend [Wang et al., 2019][ref-dgl-2019], [Paszke et al., 2019][ref-pytorch-2019]. A fully connected graph is constructed for each event, with nodes corresponding to reconstructed jets, electrons, muons, photons, and \\(\vec{E}_T^{\text{miss}}\\). The features of each node include the four-momentum \\((p_T, \eta, \phi, E)\\) of the object with a massless assumption (\\(E = p_T \cosh \eta\\)), the b-tagging label (for jets), the charge (for leptons), and an integer labeling the type of object represented by the node. We use a placeholder value of 0 for features which are not defined for every node type such as the b-jet tag, lepton charge, or the pseudorapidity of \\(\vec{E}_T^{\text{miss}}\\). We assign the angular distances (\\(\Delta \eta, \Delta \phi, \Delta R\\)) as edge features and the number of nodes $N$ in the graph as a global feature. We denote the node features \\(\{\vec x_i\}\\), edge features \\(\{\vec y_{ij}\}\\), and global features \\(\{\vec z\}\\). - -The GNN model is based on the graph network architecture described in [Battaglia et al., 2018][ref-graphnets-2018] using simple multilayer perceptron (MLP) feature functions and summation aggregation. The model is comprised of three primary components: an encoder, the graph network, and a decoder. In the encoder, three MLPs embed the nodes, edges, and global features into a latent space of dimension 64. The graph network block, which is designed to facilitate message passing between different domains of the graph, performs an edge update $f_e$, followed by a node update $f_n$, and finally a global update $f_g$, all defined below. The inputs to each update MLP are concatenated. - -$$ -\vec {y'}_{ij} = f_e\left(\{\vec x_k\},\vec y_{ij},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\vec x_j,\vec y_{ij},\vec z\right) -$$ - -$$ -\vec{x'}_{i} = f_n\left(\vec x_i,\{\vec{y'}_{jk}\},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\sum_j\vec{y'}_{ij},\vec z\right) -$$ - -$$ -\vec{z'} = f_g\left(\{\vec{x'}_i\},\{\vec{y'}_{ij}\},\vec z\right) = \mathrm{MLP}\left(\sum_i\vec{x'}_i,\sum_{i,j}\vec{y'}_{ij},\vec z\right) -$$ - -This graph block is iterated four times with the same update MLPs. Finally, the global features are passed through a decoder MLP and a final layer linear to produce the desired model outputs. Each MLP consists of 4 linear layers, each with an output width of 64, with the `ReLU` activation function. The output of the MLP is then passed through a `LayerNorm` layer [Ba et al., 2016][ref-layernorm-2016]. The total number of trainable parameters in this model is about 400,000. - -As a performance benchmark, a baseline GNN model is trained from scratch for each classification task. The initial learning rate is set to \\(10^{-4}\\) with an exponential decay following \\(LR(x) = LR_{\text{initial}}\cdot(0.99)^x\\), where \\(x\\) represents the epoch number. - ---- - -### Pretraining Strategy - -We explore two complementary pretraining approaches to develop robust representations of collision events: (1) multi-class classification, which trains the model to distinguish between different physics processes, and (2) multi-label classification, which predicts the existence and kinematics of heavy particles with prompt decays. The pretraining dataset consists of approximately 120 million events, evenly distributed across 12 distinct physics processes, including all major Higgs boson production mechanisms and top quark processes as described in [Data Samples](#sec-data). This large-scale pretraining effort was conducted on the Perlmutter supercomputer at NERSC. - -#### Multi-class Classification - -For Monte Carlo simulated events, the underlying physics process that generated each event is known precisely, providing natural labels for supervised learning. However, the challenge lies in the complexity of collision events: different physics processes can produce similar kinematics and event topologies, particularly in certain regions of phase space. No single observable can unambiguously identify the underlying process. By training the model to distinguish between 12 different processes simultaneously, we challenge it to learn subtle differences in kinematics and topology that collectively characterize each process. The model is trained using categorical cross entropy as the loss function. The output layer of the multiclass classification model has 832 trainable parameters. - -#### Multi-label Classification - -This approach combines both classification and regression tasks to characterize collision events. For discrete properties like particle presence in specific kinematic regions, we employ classification labels with binary cross-entropy loss. For continuous quantities like particle multiplicities, we use regression labels with mean-squared error loss. This hybrid approach enables the model to learn both categorical and continuous aspects of the physics processes simultaneously. - -We develop a comprehensive set of 41 labels that capture both particle multiplicities and kinematic properties. This approach increases prediction granularity and enhances model interpretability. By training the model to predict event kinematics rather than event identification, we create a task-independent framework that can potentially generalize better to novel scenarios not seen during pretraining. - -The particle multiplicity labels count the number of Higgs bosons (\\(n_{\text{higgs}}\\)), top quarks (\\(n_{\text{tops}}\\)), vector bosons (\\(n_V\\)), \\(W\\) bosons (\\(n_W\\)), and \\(Z\\) bosons (\\(n_Z\\)). The kinematic labels characterize the transverse momentum (\\(p_T\\)), pseudorapidity (\\(\eta\\)), and azimuthal angle (\\(\phi\\)) of Higgs bosons and top quarks through binned classifications. - -For Higgs bosons, $p_T$ is categorized into three ranges: (0, 30) GeV, (30, 200) GeV, and (200, \\(\infty\\)) GeV, with the upper range particularly sensitive to potential BSM effects. Similarly, both leading and subleading top quarks have $p_T$ classifications spanning (0, 30) GeV, (30, 300) GeV, and (300, \\(\infty\\)) GeV. When no particle exists within a specific \\(p_T\\) range, the corresponding label is set to \\([0, 0, 0]\\). For all particles, \\(\eta\\) measurements are divided into 4 bins with boundaries at \\([-1.5, 0, 1.5]\\), while \\(\phi\\) measurements use 4 bins with boundaries at \\([-\frac{\pi}{2}, 0, \frac{\pi}{2}]\\). As with \\(p_T\\), both \\(\eta\\) and \\(\phi\\) labels default to \\([0, 0, 0, 0]\\) in the absence of a particle. This comprehensive labeling schema enables fine-grained learning of kinematic distributions and particle multiplicities, essential for characterizing complex collision events. - -The loss function combines individual losses from all 41 labels through weighted averaging. Binary cross-entropy is applied to classification labels, while mean-squared error is used for regression labels. The model generates predictions for all labels simultaneously, with individual losses calculated according to their respective types. The final loss is computed as an equally-weighted average across all labels, with weights set to 1 to ensure uniform contribution to the optimization process. The output layer of the multilabel model has 2,688 trainable parameters. - -#### Pretraining - -During pre-training, the initial learning rate is \\(10^{-4}\\), and the learning rate decays by 1% each epoch following the power law function \\(LR(x) = 10^{-4}\cdot(0.99)^x\\), where \\(x\\) is the number of epochs. Both pre-trained models reach a plateau in loss by epoch 50, at which point the training is stopped. - ---- -### Fine-tuning Methodology - -For downstream tasks, we adjust the model architecture for fine-tuning by replacing the original output layer (final linear layer) with a newly initialized linear layer while retaining the pre-trained weights for all other layers. This modification allows the model to specialize in the specific downstream task while leveraging the general features learned during pretraining. - -The fine-tuning process begins with distinct learning rate setups for different parts of the model. The newly initialized linear layer is trained with an initial learning rate of \\(10^{-4}\\), matching the rate used for models trained from scratch. Meanwhile, the pre-trained layers are fine-tuned more cautiously with a lower initial learning rate of \\(10^{-5}\\). This approach ensures that the pre-trained layers adapt gradually without losing their general features, while the new layer learns effectively from scratch. Both learning rates decay over time following the same power law function, \\(LR(x) = LR_{initial} \cdot (0.99)^x\\), to promote stable convergence as training progresses. - -We also evaluated a transfer learning setup in which either the decoder MLP or the final linear layer was replaced with a newly initialized component. During this process, all other model parameters remained frozen, leveraging the pre-trained features without further updating them. However, we did not observe performance improvements using the transfer learning setup. Consequently, we focus on reporting results obtained with the fine-tuning approach. - ---- - -### Performance Evaluation - -We assess model performance using two figures of merit: the classification accuracy and the Area Under the Curve (AUC) of the Receiver Operating Characteristic (ROC) curve. The accuracy is defined as the fraction of correctly classified events when applying a threshold of 0.5 to the neural network output score. Both metrics demonstrate consistent trends in our analysis. - -To obtain reliable performance estimates and uncertainties, we employ an ensemble training approach where 5 independent models are trained for each configuration with random weight initialization and random subsets of the training dataset. This enables us to evaluate both the models' sensitivity to initial parameters and to quantify uncertainties in their performance. - -To investigate how model performance scales with training data, we conducted training runs using sample sizes ranging from \\(10^3\\) to \\(10^7\\) events per class (\\(10^3\\), \\(10^4\\), \\(10^5\\), \\(10^6\\), and \\(10^7\\)) for each model setup: the from-scratch baseline and models fine-tuned from multi-class or multi-label pretrained models. For the \\(10^7\\) case, only the initialization was randomized due to dataset size limitations. All models were evaluated on the same testing dataset, consisting of 2 million events per class, which remained separate from the training process. - -| **Name of Task** | **Pretraining Task** | \\(10^3\\) | \\(10^4\\) | \\(10^5\\) | \\(10^6\\) | \\(10^7\\) | -|----------------------|----------------------|--------------------|--------------------|--------------------|--------------------|--------------------| -| **ttH CP Even vs Odd** | Baseline Accuracy | 56.5 ± 1.1 | 62.2 ± 0.1 | 64.3 ± 0.0 | 65.7 ± 0.0 | 66.2 ± 0.0 | -| | Multiclass (%) | +4.8 ± 1.1 | +3.4 ± 0.1 | +1.3 ± 0.0 | +0.2 ± 0.0 | −0.0 ± 0.0 | -| | Multilabel (%) | +2.1 ± 1.2 | +1.9 ± 0.1 | +0.8 ± 0.1 | +0.0 ± 0.0 | −0.1 ± 0.0 | -| **FCNC vs tHq** | Baseline Accuracy | 63.6 ± 0.7 | 67.8 ± 0.4 | 68.4 ± 0.3 | 69.3 ± 0.3 | 67.9 ± 0.0 | -| | Multiclass (%) | +5.8 ± 0.8 | +1.2 ± 0.4 | +1.4 ± 0.3 | +0.5 ± 0.3 | −0.0 ± 0.0 | -| | Multilabel (%) | −5.3 ± 0.8 | −1.3 ± 0.4 | +0.9 ± 0.4 | +0.3 ± 0.3 | +0.4 ± 0.1 | -| **ttW vs ttt** | Baseline Accuracy | 75.8 ± 0.1 | 77.6 ± 0.1 | 78.9 ± 0.0 | 79.8 ± 0.0 | 80.3 ± 0.0 | -| | Multiclass (%) | +3.7 ± 0.1 | +2.7 ± 0.1 | +1.3 ± 0.0 | +0.4 ± 0.0 | +0.0 ± 0.0 | -| | Multilabel (%) | +2.2 ± 0.1 | +1.1 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.1 ± 0.0 | -| **stop vs ttH** | Baseline Accuracy | 83.0 ± 0.2 | 86.3 ± 0.1 | 87.6 ± 0.0 | 88.5 ± 0.0 | 88.8 ± 0.0 | -| | Multiclass (%) | +0.4 ± 0.2 | +1.9 ± 0.1 | +1.0 ± 0.0 | +0.3 ± 0.0 | +0.0 ± 0.0 | -| | Multilabel (%) | +2.8 ± 0.2 | +1.0 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.0 ± 0.0 | -| **WH vs ZH** | Baseline Accuracy | 51.4 ± 0.1 | 53.9 ± 0.1 | 55.8 ± 0.0 | 57.5 ± 0.0 | 58.0 ± 0.0 | -| | Multiclass (%) | +5.2 ± 0.1 | +5.3 ± 0.1 | +3.1 ± 0.0 | +0.6 ± 0.0 | +0.1 ± 0.0 | -| | Multilabel (%) | −1.1 ± 0.1 | −0.9 ± 0.2 | +0.5 ± 0.1 | +0.1 ± 0.0 | −0.1 ± 0.0 | - -> **Table 1**: Accuracy of the traditional model versus the accuracy increase due to fine-tuning from various pretraining tasks. -> The accuracies are averaged over 5 independently trained models with randomly initialized weights and trained on a random subset of the data. One exception is the \\(10^7\\) training where all models use the same dataset due to limitations on our dataset size. The random subsets are allowed to overlap, but this overlap should be very minimal because all models take an independent random subset of \\(10^7\\) events. The testing accuracy is calculated from the same testing set of 2 million events per class across all models for a specific training task. The errors are the propagated errors (root sum of squares) of the standard deviation of accuracies for each model. - -## Results - -### Classification Performance - -Since the observations of AUC and accuracy show similar trends, we focus the presentation of the results using accuracy here for conciseness in Table 1. - -In general, the fine-tuned pretrained model achieves at least the same level of classification performance as the baseline model. Notably, there are significant improvements, particularly when the sample size is small, ranging from \\(10^3\\) to \\(10^4\\) events. In some cases, the accuracy improvements exceed five percentage points, demonstrating that pretrained models provide a strong initial representation that compensates for limited data. The numerical values of the improvements in accuracy may not fully capture the impact on the sensitivity of the measurements for which the neural network classifier is used, and the final sensitivity improvement is likely to be greater. - -As the training sample size grows to \\(10^5\\), \\(10^6\\), and eventually \\(10^7\\) events, the added benefit of pretraining diminishes. With abundant data, models trained from scratch approach or even match the accuracy of fine-tuned pretrained models. This suggests that large datasets enable effective learning from scratch, rendering the advantage of pretraining negligible in such scenarios. - -Although both pretraining approaches offer benefits, multiclass pretraining tends to provide more consistent improvements across tasks, especially in the low-data regime. In contrast, multilabel pretraining can sometimes lead to neutral or even slightly negative effects for certain tasks and data sizes. This highlights the importance of the pretraining task design, as the similarity between pretraining and fine-tuning tasks in the multiclass approach appears to yield better-aligned representations. - -Finally, the spread of accuracy across the five tasks for the baseline model is quite large, offering a robust test of fine-tuning across tasks of varying difficulty. The consistent observation of these trends across tasks confirms the reliability and robustness of the findings. - ---- - -### Model Interpretability - -We aim to understand whether pretrained and baseline models learn the same underlying representations. If the two models exhibit high similarity, a plausible interpretation is that pretraining provides the pretrained model with an advantageous initialization, allowing it to converge to a similar state as the baseline model more efficiently. Conversely, significant differences between the models would indicate that pretraining facilitates the development of a more general and robust latent space, which serves as a foundation for fine-tuning to effectively adapt to the downstream task. To investigate this, we analyzed the representational similarity between a pretrained model fine-tuned for the downstream task and a baseline model trained directly on the downstream task without pretraining. - -We use Centered Kernel Alignment (CKA) [Kornblith et al., 2019][ref-kornblith-2019-cka] to analyze model similarity and interpretability. CKA is a robust metric that quantifies the similarity between the internal representations of neural networks by comparing their feature matrices in a manner that is invariant to scaling, rotation, and alignment. This invariance makes CKA particularly effective for studying relationships between network layers, even across networks of different sizes or those trained from varying initializations. - -The similarity is evaluated using a 64-dimensional latent representation after the decoder stage of the GNN model. This choice allows us to compare the internal states of the models at a fine-grained level and understand how training strategies impact the representations directly used for the output task. - -To provide an intuitive understanding of CKA values, we construct a table of the CKA scores for various transformations performed on a set of dummy data. - -- **A:** randomly initialized matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\)) -- **B:** matrix with shape (1000, 64) constructed via various transformations performed on \\(A\\) -- **Noise:** randomly initialized noise matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\)) - -| Dataset | CKA Score | -|---------|-----------| -| \\(A, B = A\\) | 1.00 | -| \\(A, B =\\) permutation on columns of \\(A\\) | 1.00 | -| \\(A, B = A + \mathrm{Noise}(0.1)\\) | 0.99 | -| \\(A, B = A + \mathrm{Noise}(0.5)\\) | 0.80 | -| \\(A, B = A + \mathrm{Noise}(0.75)\\) | 0.77 | -| \\(A, B = A \cdot \mathrm{Noise}(1)\\) (Linear Transformation) | 0.76 | -| \\(A, B = A + \mathrm{Noise}(1)\\) | 0.69 | -| \\(A, B = A + \mathrm{Noise}(2)\\) | 0.51 | -| \\(A, B = A + \mathrm{Noise}(5)\\) | 0.39 | - -**Table 2:** CKA scores for a dummy dataset \\(A\\) and \\(B\\), where \\(B\\) is created via various transformations performed on \\(A\\). - -As seen in Table 2 and in the definition of the CKA, the CKA score is permutation-invariant. We will use the CKA score to evaluate the similarity between various models and gain insight into the learned representation of detector events in each model (i.e., the information that each model learns). - -We train ensembles of models for each training task to observe how the CKA score changes due to the random initialization of our models. The CKA score between two models is then defined to be: - -\\[ -CKA(A, B) = \frac{1}{n^2} \sum_i^n \sum_j^n CKA(A_i, B_j) -\\] - -where \\(A_i\\) is the representation learned by the \\(i^{\text{th}}\\) model in an ensemble with \\(n\\) total models. The error in CKA is the standard deviation of \\(CKA(A_i, B_j)\\). - -Here we present results for the CKA similarity between the final model in each setup with the final model in the baseline, shown in Table 3. - -| Training Task | Baseline | Multiclass | Multilabel | -|-----------------------|------------------|-----------------|-----------------| -| ttH CP Even vs Odd | 0.94 ± 0.05 | 0.82 ± 0.01 | 0.77 ± 0.06 | -| FCNC vs tHq | 0.96 ± 0.03 | 0.76 ± 0.01 | 0.81 ± 0.01 | -| ttW vs ttt | 0.91 ± 0.08 | 0.75 ± 0.10 | 0.72 ± 0.05 | -| stop vs ttH | 0.87 ± 0.11 | 0.79 ± 0.12 | 0.71 ± 0.08 | -| WH vs ZH | 0.90 ± 0.07 | 0.53 ± 0.03 | 0.44 ± 0.06 | - -**Table 3:** CKA Similarity of the latent representation before the decoder with the baseline model, averaged over 3 models per training setup, and all models trained with the full dataset (\\(10^7\\)). The baseline column is not guaranteed to be 1.0 because of the random initialization of the model. Each baseline model converges to a slightly different representation as seen in the CKA values in that column. - -The baseline models with different initializations exhibit high similarity values, ranging from approximately 0.87 to 0.96, which indicates that independently trained baseline models tend to converge on similar internal representations despite random initialization. Across the considered tasks, models trained as multi-class or multi-label classifiers exhibit noticeably lower CKA similarity scores when compared to the baseline model. For example, in the WH vs ZH task, the baseline model and another baseline trained model have a high similarity of 0.90, whereas the multi-class and multi-label models show significantly reduced similarities (0.53 and 0.44, respectively). This pattern suggests that the representational spaces developed by multi-class or multi-label models differ substantially from those learned by the baseline model that was trained directly on the downstream classification task. - -### Computational Efficiency - -To estimate the computational resources required for each approach, we measured the wall time needed for a model to reach its final performance. For baseline models, this is defined as the wall time from the start of training until the loss of the model plateaus. For the foundation model approach, the estimate includes both the pretraining time and the fine-tuning time, each measured from the start of training until the loss plateaus. This approach ensures a consistent and comprehensive evaluation of the computational demands. - -![The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.](training_time.png) -*Fig. 1: The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.* - -Figure 1 shows the fine-tuning time for the model pretrained with multiclass classification, relative to the time required for the baseline model, as a function of training sample size. In general, the fine-tuning time is significantly shorter than the training time required by the baseline model approach. For smaller training sets, on the order of \\(10^5\\) events, tasks such as FCNC vs. tHq and ttW vs. ttt benefit substantially from the pretrained model’s “head start,” achieving their final performance in only about 1% of the baseline time. For large training datasets, the fine-tuning time relative to the baseline training time becomes larger; however, given that the large training sample typically requires longer training time, fine-tuning still yields much faster training convergence. The ttH CP-even vs. ttH CP-odd task, with a training sample size of \\(10^7\\) events, is an exception where the fine-tuning time exceeds the training time required for the baseline model. This is likely because the processes involved in this task include photon objects in the final states, which are absent from the events used during pretraining. - -To accurately evaluate the total time consumption, it is necessary to include the pretraining time required for the foundation model approach. The pretraining times are as follows: - -- **Multi-class pretraining:** 45.5 GPU hours -- **Multi-label pretraining:** 60.0 GPU hours - -The GPU hours recorded for the multi-label model represent the total time required when training the model in parallel on 16 GPUs. This includes a model synchronization step, which results in higher GPU hours compared to the multi-class pretraining model. - -The foundation model approach becomes increasingly efficient when a large number of tasks are fine-tuned using the same pretrained model, compared to training each task independently from scratch. To illustrate this, we evaluate the computational time required for a scenario where the training sample contains \\(10^7\\) events. For the five tasks tested in this study, the baseline training time (training from scratch) ranges from 1.68 GPU hours (WH vs. ZH) to 5.30 GPU hours (ttW vs. ttt), with an average baseline training time of 2.94 GPU hours. In contrast, the average fine-tuning time for the foundation model approach, relative to the baseline, is 38% of the baseline training time for \\(10^7\\) events. Based on these averages, we estimate that the foundation model approach becomes more computationally efficient than the baseline approach when fine-tuning is performed for more than 41 tasks. - -As a practical example, the ATLAS measurement of Higgs boson couplings using the \\(H \rightarrow \gamma\gamma\\) decay channel [ATLAS Collaboration, 2023][ref-atlas-2023-higg] involved training 42 classifiers for event categorization. This coincides with our estimate, suggesting that the foundation model approach can reduce computational costs even for a single high-energy physics measurement. - -## Conclusions - -We presented an in-depth study of a particle physics foundation model designed to operate on the four-momentum and identification properties of event final-state objects. This model is built on a Graph Neural Network (GNN) architecture and trained on a dataset comprising 120 million simulated proton-proton collision events across 12 distinct physics processes. The pretraining phase explored both multiclass and multilabel classification tasks, providing a robust foundation for downstream applications. Notably, the pretrained models demonstrated significant improvements in event classification performance when fine-tuned, particularly for tasks with limited training samples. - -The foundation model approach also offers substantial computational advantages. By leveraging fine-tuning, this methodology reduces the computational resources required for large-scale applications across multiple tasks. Our estimates indicate that significant resource savings can be achieved even for single particle physics measurements, making this approach both scalable and efficient. - -To better understand the learned representations of the pretrained model and guide future optimization efforts, we employed a representational similarity evaluation framework using Centered Kernel Alignment (CKA). This metric allowed us to investigate the source of the performance gains observed in the foundation model. Our analysis revealed notable differences in the learned representations between the fine-tuned pretrained model and a baseline model trained from scratch. In deep learning, it is well-established that multiple equally valid solutions can exist. Future studies are necessary to determine whether the low similarity in latent representations reflects complementary information uniquely captured by the foundation and baseline models, or if it can simply be attributed to connected local minima in the loss landscape. - -## Acknowledgments - -This work is supported by the U.S. National Science Foundation under the Award No. 2046280, and by U.S. Department of Energy, Office of Science under contract DE-AC02-05CH11231. - -## References - -- **OpenAI et al.** GPT-4 Technical Report. arXiv:2303.08774 (2024). [https://arxiv.org/abs/2303.08774](https://arxiv.org/abs/2303.08774) - -- **Jason Yosinski, Jeff Clune, Yoshua Bengio, Hod Lipson.** How transferable are features in deep neural networks? CoRR abs/1411.1792 (2014). [http://arxiv.org/abs/1411.1792](http://arxiv.org/abs/1411.1792) - -- **Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer.** High-Resolution Image Synthesis with Latent Diffusion Models. CoRR abs/2112.10752 (2021). [https://arxiv.org/abs/2112.10752](https://arxiv.org/abs/2112.10752) - -- **Dustin Podell, Zion English, Kyle Lacey et al.** SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. arXiv:2307.01952 (2023). [https://arxiv.org/abs/2307.01952](https://arxiv.org/abs/2307.01952) - -- **John Jumper, Richard Evans, Alexander Pritzel et al.** Highly accurate protein structure prediction with AlphaFold. Nature 596, 583-589 (2021). [https://doi.org/10.1038/s41586-021-03819-2](https://doi.org/10.1038/s41586-021-03819-2) - -- **Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova.** BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. CoRR abs/1810.04805 (2018). [http://arxiv.org/abs/1810.04805](http://arxiv.org/abs/1810.04805) - -- **ATLAS Collaboration.** Measurement of the properties of Higgs boson production at \\(\sqrt{s} = 13\,\text{TeV}\\) in the \\(H \to \gamma\gamma\\) channel using \\(139\,\text{fb}^{-1}\\) of \\(pp\\) collision data with the ATLAS experiment. JHEP 07 (2023) 088. [arXiv:2207.00348](https://arxiv.org/abs/2207.00348), [https://doi.org/10.1007/JHEP07(2023)088](https://doi.org/10.1007/JHEP07(2023)088) - -- **ATLAS Collaboration.** Observation of four-top-quark production in the multilepton final state with the ATLAS detector. Eur. Phys. J. C 83 (2023) 496. [arXiv:2303.15061](https://arxiv.org/abs/2303.15061), [https://doi.org/10.1140/epjc/s10052-023-11573-0](https://doi.org/10.1140/epjc/s10052-023-11573-0) - -- **Simon Kornblith, Mohammad Norouzi, Honglak Lee, Geoffrey Hinton.** Similarity of Neural Network Representations Revisited. CoRR abs/1905.00414 (2019). [http://arxiv.org/abs/1905.00414](http://arxiv.org/abs/1905.00414) - ---- - - - -- **N. D. Birell, P. C. W. Davies.** Quantum Fields in Curved Space. Cambridge Univ. Press (1982). - -- **R. P. Feynman.** Phys. Rev. 94, 262 (1954). - -- **A. Einstein, Yu. Podolsky, N. Rosen.** Phys. Rev. 47, 777 (1935). - -- **G. P. Berman, Jr., F. M. Izrailev, Jr.** Stability of nonlinear modes. Physica D 88, 445 (1983). - -- **E. B. Davies, L. Parns.** Trapped modes in acoustic waveguides. Q. J. Mech. Appl. Math. 51, 477–492 (1988). - -- **Edward Witten.** hep-th/0106109 (2001). [https://arxiv.org/abs/hep-th/0106109](https://arxiv.org/abs/hep-th/0106109) - ---- - - - -- **E. Beutler.** Williams Hematology, 5th Edition, Chapter 7, pp. 654–662. McGraw-Hill, New York (1994). - -- **Donald E. Knuth.** The Art of Computer Programming vol. 1: Fundamental Algorithms, 2nd Ed., Addison-Wesley (1973). - -- **J. S. Smith, G. W. Johnson.** Philos. Trans. R. Soc. London, Ser. B 777, 1395 (2005). - -- **W. J. Smith, T. J. Johnson, B. G. Miller.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (unpublished, 2010). - -- **V. K. Smith, K. Johnson, M. O. Klein.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (submitted, 2010). - -- **Ulrich Underwood, Ned Net, Paul Pot.** Lower Bounds for Wishful Research Results. Talk at Fanstord University (1988). - -- **M. P. Johnson, K. L. Miller, K. Smith.** Personal communication (Jan-May 2007). - ---- - - - -- **Adam Paszke et al.** PyTorch: An Imperative Style, High-Performance Deep Learning Library. arXiv:1912.01703 (2019). [http://arxiv.org/abs/1912.01703](http://arxiv.org/abs/1912.01703) - -- **Minjie Wang et al.** Deep Graph Library: Towards Efficient and Scalable Deep Learning on Graphs. arXiv:1909.01315 (2019). [http://arxiv.org/abs/1909.01315](http://arxiv.org/abs/1909.01315) - -- **Peter W. Battaglia et al.** Relational inductive biases, deep learning, and graph networks. arXiv:1806.01261 (2018). [http://arxiv.org/abs/1806.01261](http://arxiv.org/abs/1806.01261) - -- **Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton.** Layer Normalization. arXiv:1607.06450 (2016). [https://arxiv.org/abs/1607.06450](https://arxiv.org/abs/1607.06450) - ---- - - - -- **Andrew J. Wildridge et al.** Bumblebee: Foundation Model for Particle Physics Discovery. arXiv:2412.07867 (2024). [https://arxiv.org/abs/2412.07867](https://arxiv.org/abs/2412.07867) - -- **Subash Katel et al.** Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture. arXiv:2412.05333 (2024). [https://arxiv.org/abs/2412.05333](https://arxiv.org/abs/2412.05333) - -- **Jack Y. Araz et al.** Point cloud-based diffusion models for the Electron-Ion Collider. arXiv:2410.22421 (2024). [https://arxiv.org/abs/2410.22421](https://arxiv.org/abs/2410.22421) - -- **Matthew Leigh et al.** Is Tokenization Needed for Masked Particle Modelling? arXiv:2409.12589 (2024). [https://arxiv.org/abs/2409.12589](https://arxiv.org/abs/2409.12589) - -- **Vinicius Mikuni, Benjamin Nachman.** OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks. arXiv:2404.16091 (2024). [https://arxiv.org/abs/2404.16091](https://arxiv.org/abs/2404.16091) - -- **Zhengde Zhang et al.** Xiwu: A Basis Flexible and Learnable LLM for High Energy Physics. arXiv:2404.08001 (2024). [https://arxiv.org/abs/2404.08001](https://arxiv.org/abs/2404.08001) - -- **Philip Harris et al.** Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models. arXiv:2403.07066 (2024). [https://arxiv.org/abs/2403.07066](https://arxiv.org/abs/2403.07066) - -- **Joschka Birk, Anna Hallin, Gregor Kasieczka.** OmniJet-$\alpha$: the first cross-task foundation model for particle physics. Machine Learning: Science and Technology. 5(3), 035031 (Aug 2024). [https://doi.org/10.1088/2632-2153/ad66ad](https://doi.org/10.1088/2632-2153/ad66ad) - -- **Andris Huang et al.** A Language Model for Particle Tracking. arXiv:2402.10239 (2024). [https://arxiv.org/abs/2402.10239](https://arxiv.org/abs/2402.10239) - -- **Tobias Golling et al.** Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models. arXiv:2401.13537 (2024). [https://arxiv.org/abs/2401.13537](https://arxiv.org/abs/2401.13537) - -- **Junze Liu et al.** Generalizing to new geometries with Geometry-Aware Autoregressive Models (GAAMs) for fast calorimeter simulation. Journal of Instrumentation 18(11), P11003 (Nov 2023). [https://doi.org/10.1088/1748-0221/18/11/p11003](https://doi.org/10.1088/1748-0221/18/11/p11003) - -- **Baran Hashemi et al.** Ultra-high-granularity detector simulation with intra-event aware generative adversarial network and self-supervised relational reasoning. Nature Communications 15(1) (June 2024). [https://doi.org/10.1038/s41467-024-49104-4](https://doi.org/10.1038/s41467-024-49104-4) - -- **Matthias Vigl et al.** Finetuning Foundation Models for Joint Analysis Optimization. arXiv:2401.13536 (2024). [https://arxiv.org/abs/2401.13536](https://arxiv.org/abs/2401.13536) - -- **Chen Li, Hao Cai, Xianyang Jiang.** Refine neutrino events reconstruction with BEiT-3. Journal of Instrumentation 19(6), T06003 (Jun 2024). [https://doi.org/10.1088/1748-0221/19/06/t06003](https://doi.org/10.1088/1748-0221/19/06/t06003) \ No newline at end of file +# 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`), with +PyTorch 2.2.2 and the optional ROOT-GNN stack DGL 2.4.0. The canonical setup +is: + +```bash +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. The DGL extra uses the validated CUDA 12.1 wheel +source configured in `pyproject.toml`; a compatible NVIDIA driver is still +required. Do not add site-specific CUDA, Slurm, or filesystem paths to model +or task 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 +`:` 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 +`/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. diff --git a/benchmarks/README.md b/benchmarks/README.md new file mode 100644 index 0000000000000000000000000000000000000000..82ab74777dfa02d3c8ebe824f7b0829808f6d731 --- /dev/null +++ b/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. diff --git a/benchmarks/__init__.py b/benchmarks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d1e440834585da30b3969042ba9f8e4d5f5f893d --- /dev/null +++ b/benchmarks/__init__.py @@ -0,0 +1 @@ +"""Standalone, reproducible performance measurements for GNN4Colliders.""" diff --git a/benchmarks/_common.py b/benchmarks/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..842b9868f12deaa72d568c4504658cb26260f572 --- /dev/null +++ b/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)) diff --git a/benchmarks/benchmark_dataloader.py b/benchmarks/benchmark_dataloader.py new file mode 100644 index 0000000000000000000000000000000000000000..c57eba9128660b18cd21aaa0a6be5da7172c1fd8 --- /dev/null +++ b/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() diff --git a/benchmarks/benchmark_inference.py b/benchmarks/benchmark_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..7db8fdaf569c035219bf9389ed5e6e6c5a82cdaf --- /dev/null +++ b/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() diff --git a/benchmarks/benchmark_preprocessing.py b/benchmarks/benchmark_preprocessing.py new file mode 100644 index 0000000000000000000000000000000000000000..9b3ccb75e2b7aaead52e6c2cfedac5c99bbbc91d --- /dev/null +++ b/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() diff --git a/benchmarks/benchmark_training.py b/benchmarks/benchmark_training.py new file mode 100644 index 0000000000000000000000000000000000000000..51743dd3a4fdf52aa60541b2eebf076e609c3728 --- /dev/null +++ b/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() diff --git a/configs/.gitkeep b/configs/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/data/.gitkeep b/configs/data/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/data/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/environment/.gitkeep b/configs/environment/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/environment/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/model/.gitkeep b/configs/model/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/model/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/model/root_gnn/.gitkeep b/configs/model/root_gnn/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/model/root_gnn/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/task/.gitkeep b/configs/task/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/task/.gitkeep @@ -0,0 +1 @@ + diff --git a/configs/trainer/.gitkeep b/configs/trainer/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/configs/trainer/.gitkeep @@ -0,0 +1 @@ + diff --git a/data/fixtures/.gitkeep b/data/fixtures/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/data/fixtures/.gitkeep @@ -0,0 +1 @@ + diff --git a/data/raw/.gitkeep b/data/raw/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/data/raw/.gitkeep @@ -0,0 +1 @@ + diff --git a/docs/architecture.md b/docs/architecture.md new file mode 100644 index 0000000000000000000000000000000000000000..05654ba157be41b0949534f2eb29f14e86f817ae --- /dev/null +++ b/docs/architecture.md @@ -0,0 +1,375 @@ +# 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, 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. 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. This is an inference boundary only; +the native model continues to consume DGL graphs and ONNX 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_.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_.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. diff --git a/docs/compatibility.md b/docs/compatibility.md new file mode 100644 index 0000000000000000000000000000000000000000..d56c12e798248a17666e4528c5de08003eae4ba2 --- /dev/null +++ b/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. diff --git a/docs/configuration.md b/docs/configuration.md new file mode 100644 index 0000000000000000000000000000000000000000..d5079012d15740a0e65149099f8297eb62621361 --- /dev/null +++ b/docs/configuration.md @@ -0,0 +1,94 @@ +# 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 +`/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=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. diff --git a/docs/end_to_end_validation.md b/docs/end_to_end_validation.md new file mode 100644 index 0000000000000000000000000000000000000000..aa78013ab5a37ee17a2852619bea351deae5a3eb --- /dev/null +++ b/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 . +``` diff --git a/docs/export.md b/docs/export.md new file mode 100644 index 0000000000000000000000000000000000000000..37d9280953376494ba94ab0c7c9cdc1eec81a1b6 --- /dev/null +++ b/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. diff --git a/docs/migration.md b/docs/migration.md new file mode 100644 index 0000000000000000000000000000000000000000..4ee44050221e663da80423950b65207cb49eca1b --- /dev/null +++ b/docs/migration.md @@ -0,0 +1,288 @@ +# 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. + +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_.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. diff --git a/docs/performance.md b/docs/performance.md new file mode 100644 index 0000000000000000000000000000000000000000..08b5e09a90289195186f2491ec6163a09e36f068 --- /dev/null +++ b/docs/performance.md @@ -0,0 +1,20 @@ +# Performance notes + +The scripts in [`../benchmarks/`](../benchmarks/) separate setup, warmup, and +steady-state timings and emit JSON lines with execution metadata. Float32 eager +execution remains the correctness reference path. + +`fully_connected_edges` has a bounded 32-entry cache keyed by node count, +self-loop policy, and device. It preserves source-major ordering and only +reuses topology indices; event-dependent edge features are always recomputed. +The cache is graph-specific and does not alter scientific behavior. + +Training already uses `zero_grad(set_to_none=True)` and inference already uses +`torch.inference_mode()` with detached CPU accumulation. + +Mixed precision, `torch.compile`, custom kernels, aggressive worker defaults, +and cache-format replacement were not retained without target-machine +measurements. The main known bottleneck is the quadratic graph workload +`N * (N - 1)` and associated DGL message passing; size-aware batching and +streaming prediction remain follow-up work because they affect ordering or +output semantics. diff --git a/docs/perlmutter.md b/docs/perlmutter.md new file mode 100644 index 0000000000000000000000000000000000000000..25a6398c7c0fbe6c7dde9b0ff730c41a8c1d0b82 --- /dev/null +++ b/docs/perlmutter.md @@ -0,0 +1,58 @@ +# Perlmutter execution + +The package does not encode Perlmutter paths, modules, or allocation policy. +Create the uv environment in the project location appropriate for your +account, select a site-compatible GPU/driver environment, and use the same +Hydra configuration as on a workstation. + +## Single GPU + +```bash +uv run gnn4colliders train \ + data.cache.path=/path/to/graphs.pt \ + environment=perlmutter \ + trainer.device=cuda \ + trainer.max_epochs=10 +``` + +The equivalent Slurm wrapper is: + +```bash +sbatch scripts/slurm/train_single_gpu.sh \ + data.cache.path=/path/to/graphs.pt trainer.max_epochs=10 +``` + +## Single-node DDP + +```bash +GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_gpu.sh \ + data.cache.path=/path/to/graphs.pt trainer.max_epochs=10 +``` + +The wrapper uses `torchrun`; `batch_size` and `num_workers` are per GPU. +Effective batch size is `data.batch_size * number_of_processes`, and only rank +0 writes the shared checkpoint/config/prediction artifacts. + +## Multi-node DDP + +```bash +GPUS_PER_NODE=4 sbatch scripts/slurm/train_multi_node.sh \ + data.cache.path=/path/to/graphs.pt trainer.max_epochs=10 +``` + +Use the provided script as a template and adapt only allocation/account +settings required by the site. Evaluation and prediction can use +`scripts/slurm/evaluate.sh`; prediction currently gathers moderate-size +results in memory. + +## Checks and common failures + +```bash +uv run python -c "import torch, dgl; print(torch.__version__, dgl.__version__, torch.cuda.is_available())" +uv run gnn4colliders --help +``` + +An unavailable DGL wheel, incompatible driver, missing cache, or mismatched +cache schema should be fixed in the environment/input rather than hidden with +package-level path changes. GPU kernels and distributed execution are not +promised bitwise deterministic. diff --git a/docs/releasing.md b/docs/releasing.md new file mode 100644 index 0000000000000000000000000000000000000000..8e3863dc352e9f71586e08030bae24b756b09e3d --- /dev/null +++ b/docs/releasing.md @@ -0,0 +1,23 @@ +# Releasing + +The project uses a static version defined once in +`src/gnn4colliders/__init__.py`; setuptools reads that value for package +metadata. The current package is not published to PyPI. + +Before a release: + +1. Update `__version__`, `CHANGELOG.md`, and this release checklist. +2. Run `uv sync --dev` (and `--extra root-gnn` when ROOT-GNN validation is + available), then run lint and tests. +3. Run `bash scripts/dev/check_release.sh` to build and inspect the sdist and + wheel and smoke-test a wheel installation outside the checkout. +4. Review the generated artifacts and `git diff`, then create a version tag + according to the repository's release policy. + +The release-validation script never publishes artifacts. Publishing, if +adopted later, must use repository-managed credentials or trusted publishing. + +The direct DGL wheel source is retained in `pyproject.toml` because the +validated ROOT-GNN environment requires the CUDA 12.1 DGL wheel. Core imports +do not import DGL or ONNX; install the corresponding extras for those +workflows. diff --git a/docs/testing.md b/docs/testing.md new file mode 100644 index 0000000000000000000000000000000000000000..09353fa89bf67962ef5f5d781e9423ebb40389f2 --- /dev/null +++ b/docs/testing.md @@ -0,0 +1,29 @@ +# Testing + +The suite is layered by dependency and purpose: + +```bash +# Fast shared-package tests +uv run pytest tests/unit + +# Full CPU suite, including integration and parity tests +uv run pytest + +# Required ROOT-GNN parity gate (must fail if DGL is unavailable) +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest + +# Optional layers +uv run pytest -m distributed -v +uv run pytest -m onnx -v +uv run pytest -m gpu -v +GNN4COLLIDERS_ROOT_FIXTURE=/path/to/reduced.root uv run pytest -m real_data -v +``` + +Unit tests use deterministic, small tensors and generated ROOT files. DGL, +ONNX, CUDA, distributed execution, and reduced real-data fixtures remain +optional layers. Tests should assert public contracts and scientific +invariants rather than private call sequences. New regression tests should +use `tmp_path`, explicit seeds, and justified numerical tolerances. + +Release validation additionally builds a wheel and runs the CLI help command +in a clean environment; it is not part of the ordinary pytest suite. diff --git a/examples/README.md b/examples/README.md new file mode 100644 index 0000000000000000000000000000000000000000..b53cb0f871e9ec4602009d9c54200bfa568ce09b --- /dev/null +++ b/examples/README.md @@ -0,0 +1,18 @@ +# Examples + +The supported experiment interface is the root CLI with Hydra configuration; +examples do not define a second application framework. The repository's +reference groups cover multiclass pretraining, binary fine-tuning, resume, +evaluation, and NPZ prediction. Start with the commands in the top-level +[`README.md`](../README.md), then copy a config group into a local config file +when a dataset needs more than command-line overrides. + +For a complete CPU smoke workflow using generated temporary ROOT data: + +```bash +uv run python scripts/dev/smoke_end_to_end.py +``` + +Values such as input ROOT paths, cache locations, and checkpoint paths are +deliberately local placeholders. Production datasets and generated outputs do +not belong in this directory. diff --git a/LICENSE b/legacy/LICENSE similarity index 100% rename from LICENSE rename to legacy/LICENSE diff --git a/legacy/README.md b/legacy/README.md new file mode 100644 index 0000000000000000000000000000000000000000..e5e004a4ae410d57a4be172387fa1036ff0f142b --- /dev/null +++ b/legacy/README.md @@ -0,0 +1,358 @@ +--- +license: mit +tags: +- arXiv:2412.10665 +--- + +This is a demo is of the approach described in the paper, ["Pretrained Event Classification Model for High Energy Physics Analysis"](https://arxiv.org/abs/2412.10665) +``` +@misc{ho2024pretrained, + title={Pretrained Event Classification Model for High Energy Physics Analysis}, + author={Joshua Ho, Benjamin Ryan Roberts, Shuo Han, Haichen Wang}, + year={2024}, + eprint={2412.10665}, + archivePrefix={arXiv} +} +``` + +## Abstract + +We introduce a foundation model for event classification in high-energy physics, built on a **Graph Neural Network** architecture and trained on **120 million simulated proton-proton collision events** spanning 12 distinct physics processes. The model is *pretrained* to learn a general and robust representation of collision data using challenging multiclass and multilabel classification tasks. + +Its performance is evaluated across five event classification tasks, which include both physics processes used during pretraining and new processes not encountered during pretraining. Fine-tuning the pretrained model significantly improves classification performance, particularly in scenarios with limited training data, demonstrating gains in both accuracy and computational efficiency. + +To investigate the underlying mechanisms behind these performance improvements, we employ a representational similarity evaluation framework based on *Centered Kernel Alignment*. This analysis reveals notable differences in the learned representations of fine-tuned pretrained models compared to baseline models trained from scratch. + +## Introduction + +Machine learning has become a ubiquitous tool in particle physics, employed in a variety of tasks including triggering, simulation, reconstruction, and offline analysis. While its utility spans classification, regression, and generative tasks, the current paradigm of developing machine learning models from scratch for each specific application presents several challenges. This approach not only demands specialized expertise and substantial computing resources but can also result in suboptimal performance due to limited training data. The from-scratch development of models necessitates individual validation studies to ensure that neural networks utilize well-modeled information from training samples, whether derived from Monte Carlo simulations or control samples from experimental data. + +Foundation models offer a promising direction to address these limitations. These models, pre-trained on large, diverse datasets across various tasks, provide robust and general representations of underlying data structures. Notable examples in other fields include GPT-4 [OpenAI et al., 2024](#ref-openai-2024-gpt4) and BERT [Devlin et al., 2018](#ref-devlin-2018-bert) in natural language processing, Stable Diffusion [Rombach et al., 2021](#ref-rombach-2021-latentdiffusion) in image processing, and AlphaFold [Jumper et al., 2021](#ref-jumper-2021-alphafold) in structural biology. The foundation model approach offers several advantages for particle physics applications: reduced computing resources for fine-tuning [Yosinski et al., 2014](#ref-yosinski-2014-transfer) compared to training from scratch, superior performance on specific tasks (particularly with limited training data), and potentially simplified validation procedures as downstream tasks inherit verified representations from the pre-trained model. + +Current literature on pretrained models for particle physics can be categorized based on the data representation they handle. Models operating on particle- or event-level numerical data use features like particle four momenta or jets, leveraging self-supervised or generative methods to learn versatile representations. Detector-focused models operate on high-dimensional responses such as calorimeter deposits or pixel hits, employing geometry-aware techniques for accurate simulation and analysis. Finally, models using textual or code representations apply large language model architectures to integrate domain knowledge, enabling tasks like question answering and code generation. + +Recent studies have begun exploring foundation models tailored to particle physics data, which has a variety of distinct structures and properties across many experiments and data processing stages, including: + +- particle-level & event-level numeric data [Wildridge et al., 2024](#ref-wildridge-2024-bumblebee), [Katel et al., 2024](#ref-katel-2024-jet), [Golling et al., 2024](#ref-golling-2024-maskedset), [Mikuni & Nachman, 2024](#ref-mikuni-2024-omnilearn), [Harris et al., 2024](#ref-harris-2024-resimulation), [Birk et al., 2024](#ref-birk-2024-omnijet), [Vigl et al., 2024](#ref-vigl-2024-finetune), +- detector-level & geometry-aware data [Araz et al., 2024](#ref-araz-2024-pointcloud), [Liu et al., 2023](#ref-liu-2023-gaam), [Hashemi et al., 2024](#ref-hashemi-2024-gen), [Huang et al., 2024](#ref-huang-2024-lmtracking), +- textual or code data [Zhang et al., 2024](#ref-zhang-2024-xiwu). + +This paper presents a foundation model designed specifically for collider event-level data. In modern collider experiments, final-stage analysis processes information from reconstructed objects that either directly correspond to particles in collision final states (such as leptons and photons) or serve as proxies (such as jets and missing transverse energy). While traditional approaches often relied on "high-level" variables calculated from object features, recent trends favor direct input of event objects and their features into neural networks for analysis tasks. A notable example is [ATLAS Collaboration, 2023](#ref-atlas-2023-4top), which established the observation of simultaneous production of four top quarks with the ATLAS experiment by employing a graph neural network (GNN) architecture to process event-level object information. + +We present foundation models that adopt an architecture similar to that used for [ATLAS Collaboration, 2023](#ref-atlas-2023-4top). Our models are pre-trained using either multiclass classification or multi-label learning tasks across 12 distinct physics processes. We evaluate these models through fine-tuning and testing on five classification tasks, including both familiar and novel processes not seen during pre-training. Our analysis benchmarks the models' performance improvements, their scaling behavior with training sample size, and computational efficiency, representing the first prototype of a foundation model operating on collider final-state object data. + +## Data Samples + +To provide a diverse set of physics processes for the pretraining, we use Madgraph@NLO 2.7.3 [Alwall et al., 2014](#ref-alwall-2014hca) to generate proton-proton collision events at next-to-leading order (NLO) in Quantum Chromodynamics (QCD). We generate 12 distinct Standard Model (SM) physics processes, including six major Higgs boson production mechanisms: gluon fusion production \\(ggF\\), vector boson fusion \\(VBF\\), associated production of the Higgs boson with a W boson \\(WH\\) or a Z boson \\(ZH\\), associated production of the Higgs boson with a top-quark pair \\(t\bar{t}H\\), and associated production of the Higgs boson with a single top quark and a forward quark \\(tHq\\). Additionally, we simulate six top quark production processes: single top production, top-quark pair production \\(t\bar{t}\\), top quark pair production in association with a pair of photons \\(t\bar{t}\gamma\gamma\\), associated production of a top-quark pair with a W boson \\(t\bar{t}W\\), simultaneous production of three top quarks \\(t\bar{t}t\\), and simultaneous production of four top quarks \\(t\bar{t}t\bar{t}\\). In these samples, the Higgs boson and top quarks decay inclusively. These 12 Higgs and top quark production processes constitute the pretraining dataset. + +To test the pretrained model, we further generated four processes including three beyond Standard Model (SM) processes: a SM \\(t\bar{t}H\\) production where the Higgs boson decays exclusively to a pair of photons, a \\(t\bar{t}H\\) production with the Higgs boson decaying to a pair of photons, where the top-Yukawa coupling is CP-odd, implemented using the Higgs Characterization model [Artoisenet et al., 2013](#ref-artoisinet-2013puc), the production of a pair of superpartners of the top quark (s-top) using the Minimal Supersymmetric Standard Model (MSSM) [Rosiek, 1990](#ref-rosiek-1990), [Allanach et al., 2009](#ref-allanach-2009), and flavor changing neutral current (FCNC) processes [Degrande et al., 2015](#ref-degrande-2015), [Durieux et al., 2015](#ref-durieux-2015). For the s-top process, we simulate the production of heavier s-top pairs \\(t_2\bar{t_2}\\), where each heavier s-top (mass 582 GeV) decays into a lighter s-top \\(t_1\\) or \\(\bar{t_1}\\), mass 400 GeV) and a Higgs boson. The FCNC process involves \\(t\bar{t}\\) production where one top quark decays to a Higgs boson and a light quark. We generate 10 million events for each process, except for \\(tHq\\) and \\(t\bar{t}t\bar{t}\\), where 5 million events were produced. + +In all simulation samples, the center of mass energy of the proton-proton collision is set to 13 TeV. The Higgs boson, top quarks, and vector bosons are set to decay inclusively (except the \\(t\bar{t}H \rightarrow \gamma\gamma\\) samples), with MadSpin [Artoisenet et al., 2012](#ref-artoisinet-2012st) handling the decays of top quarks and W bosons. The generated events are processed through Pythia 8.235 [Sjostrand et al., 2015](#ref-sjostrand-2015) for parton showering and heavy particle decays, followed by Delphes 3.4.2 [de Favereau et al., 2014](#ref-defavereau-2014) configured to emulate the ATLAS detector [ATLAS Collaboration, 2008](#ref-atlas-2008) for fast detector simulation. + +The detector-level object selection criteria are defined to align with typical experimental conditions. Photons are required to have transverse momentum \\(p_T \geq 20~\mathrm{GeV}\\) and pseudorapidity \\(|\eta| \leq 2.37\\), excluding the electromagnetic calorimeter crack region \\(1.37 < |\eta| < 1.52\\). Electrons must have \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.47\\) (excluding the same crack region), while muons are selected with \\(p_T \geq 10~\mathrm{GeV}\\) and \\(|\eta| \leq 2.7\\). Jets are reconstructed using the anti-\\(k_t\\) algorithm [Cacciari et al., 2008](#ref-cacciari-2008gp) with radius parameter \\(\Delta R=0.4\\), where \\(\Delta R\\) is defined as \\(\sqrt{\Delta\eta ^2 + \Delta\phi^2}\\), with \\(\Delta\eta\\) being the difference in pseudorapidity and \\(\Delta\phi\\) the difference in azimuthal angle. Jets must satisfy \\(p_T \geq 25~\mathrm{GeV}\\) and \\(|\eta| \leq 2.5\\). To avoid double-counting, jets are removed if they are within \\(\Delta R < 0.4\\) of a photon or lepton. The identification of jets originating from b-quark decays (b-tagging) is performed by matching jets within \\(\Delta R = 0.4\\) of a b-quark, with efficiency corrections applied to match the performance of the ATLAS experiment's b-tagging algorithm [ATLAS Collaboration, 2019](#ref-atlas-2019bwq). + +## Methods + +### Overview + +We present a methodology for developing and evaluating a foundation model for particle collision event analysis. The approach centers on pretraining a Graph Neural Network (GNN) architecture using a comprehensive dataset that spans multiple physics tasks, enabling the model to learn robust and transferable features. For task-specific applications, we employ a fine-tuning strategy that combines output layer adaptation with carefully calibrated learning rates for updating the pretrained parameters. + +Given the prevalence of classification problems in particle physics data analysis, we evaluate the model's efficacy through a systematic assessment across five binary classification tasks: + +- \\(t\bar{t}H(\rightarrow \gamma\gamma)\\) with CP-even versus CP-odd t-H interaction +- \\(t\bar{t}\\) with FCNC top quark decays versus $tHq$ processes +- \\(t\bar{t}W\\) versus $ttt$ processes +- Stop pair production with Higgs bosons in the decay chain versus \\(t\bar{t}H\\) processes +- \\(WH\\) versus \\(ZH\\) production modes + +Our evaluation metrics encompass classification performance, computational efficiency, and model interpretability. The investigation extends to analyzing the model's scaling behavior with respect to training dataset size, benchmarked against models trained without pretraining. Although we explored transfer learning through parameter freezing of pretrained layers, this approach did not yield performance improvements, leading us to focus our detailed analysis on fine-tuning strategies. + +This methodological framework demonstrates the potential of foundation models to enhance the efficiency of particle physics analyses while improving task-specific performance, offering a promising direction for future high-energy physics research. + +--- + +### GNN Architecture + +We implement a Graph Neural Network (GNN) architecture that naturally accommodates the point-cloud structure of particle physics data, employing the DGL framework with a PyTorch backend [Wang et al., 2019][ref-dgl-2019], [Paszke et al., 2019][ref-pytorch-2019]. A fully connected graph is constructed for each event, with nodes corresponding to reconstructed jets, electrons, muons, photons, and \\(\vec{E}_T^{\text{miss}}\\). The features of each node include the four-momentum \\((p_T, \eta, \phi, E)\\) of the object with a massless assumption (\\(E = p_T \cosh \eta\\)), the b-tagging label (for jets), the charge (for leptons), and an integer labeling the type of object represented by the node. We use a placeholder value of 0 for features which are not defined for every node type such as the b-jet tag, lepton charge, or the pseudorapidity of \\(\vec{E}_T^{\text{miss}}\\). We assign the angular distances (\\(\Delta \eta, \Delta \phi, \Delta R\\)) as edge features and the number of nodes $N$ in the graph as a global feature. We denote the node features \\(\{\vec x_i\}\\), edge features \\(\{\vec y_{ij}\}\\), and global features \\(\{\vec z\}\\). + +The GNN model is based on the graph network architecture described in [Battaglia et al., 2018][ref-graphnets-2018] using simple multilayer perceptron (MLP) feature functions and summation aggregation. The model is comprised of three primary components: an encoder, the graph network, and a decoder. In the encoder, three MLPs embed the nodes, edges, and global features into a latent space of dimension 64. The graph network block, which is designed to facilitate message passing between different domains of the graph, performs an edge update $f_e$, followed by a node update $f_n$, and finally a global update $f_g$, all defined below. The inputs to each update MLP are concatenated. + +$$ +\vec {y'}_{ij} = f_e\left(\{\vec x_k\},\vec y_{ij},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\vec x_j,\vec y_{ij},\vec z\right) +$$ + +$$ +\vec{x'}_{i} = f_n\left(\vec x_i,\{\vec{y'}_{jk}\},\vec z\right) = \mathrm{MLP}\left(\vec x_i,\sum_j\vec{y'}_{ij},\vec z\right) +$$ + +$$ +\vec{z'} = f_g\left(\{\vec{x'}_i\},\{\vec{y'}_{ij}\},\vec z\right) = \mathrm{MLP}\left(\sum_i\vec{x'}_i,\sum_{i,j}\vec{y'}_{ij},\vec z\right) +$$ + +This graph block is iterated four times with the same update MLPs. Finally, the global features are passed through a decoder MLP and a final layer linear to produce the desired model outputs. Each MLP consists of 4 linear layers, each with an output width of 64, with the `ReLU` activation function. The output of the MLP is then passed through a `LayerNorm` layer [Ba et al., 2016][ref-layernorm-2016]. The total number of trainable parameters in this model is about 400,000. + +As a performance benchmark, a baseline GNN model is trained from scratch for each classification task. The initial learning rate is set to \\(10^{-4}\\) with an exponential decay following \\(LR(x) = LR_{\text{initial}}\cdot(0.99)^x\\), where \\(x\\) represents the epoch number. + +--- + +### Pretraining Strategy + +We explore two complementary pretraining approaches to develop robust representations of collision events: (1) multi-class classification, which trains the model to distinguish between different physics processes, and (2) multi-label classification, which predicts the existence and kinematics of heavy particles with prompt decays. The pretraining dataset consists of approximately 120 million events, evenly distributed across 12 distinct physics processes, including all major Higgs boson production mechanisms and top quark processes as described in [Data Samples](#sec-data). This large-scale pretraining effort was conducted on the Perlmutter supercomputer at NERSC. + +#### Multi-class Classification + +For Monte Carlo simulated events, the underlying physics process that generated each event is known precisely, providing natural labels for supervised learning. However, the challenge lies in the complexity of collision events: different physics processes can produce similar kinematics and event topologies, particularly in certain regions of phase space. No single observable can unambiguously identify the underlying process. By training the model to distinguish between 12 different processes simultaneously, we challenge it to learn subtle differences in kinematics and topology that collectively characterize each process. The model is trained using categorical cross entropy as the loss function. The output layer of the multiclass classification model has 832 trainable parameters. + +#### Multi-label Classification + +This approach combines both classification and regression tasks to characterize collision events. For discrete properties like particle presence in specific kinematic regions, we employ classification labels with binary cross-entropy loss. For continuous quantities like particle multiplicities, we use regression labels with mean-squared error loss. This hybrid approach enables the model to learn both categorical and continuous aspects of the physics processes simultaneously. + +We develop a comprehensive set of 41 labels that capture both particle multiplicities and kinematic properties. This approach increases prediction granularity and enhances model interpretability. By training the model to predict event kinematics rather than event identification, we create a task-independent framework that can potentially generalize better to novel scenarios not seen during pretraining. + +The particle multiplicity labels count the number of Higgs bosons (\\(n_{\text{higgs}}\\)), top quarks (\\(n_{\text{tops}}\\)), vector bosons (\\(n_V\\)), \\(W\\) bosons (\\(n_W\\)), and \\(Z\\) bosons (\\(n_Z\\)). The kinematic labels characterize the transverse momentum (\\(p_T\\)), pseudorapidity (\\(\eta\\)), and azimuthal angle (\\(\phi\\)) of Higgs bosons and top quarks through binned classifications. + +For Higgs bosons, $p_T$ is categorized into three ranges: (0, 30) GeV, (30, 200) GeV, and (200, \\(\infty\\)) GeV, with the upper range particularly sensitive to potential BSM effects. Similarly, both leading and subleading top quarks have $p_T$ classifications spanning (0, 30) GeV, (30, 300) GeV, and (300, \\(\infty\\)) GeV. When no particle exists within a specific \\(p_T\\) range, the corresponding label is set to \\([0, 0, 0]\\). For all particles, \\(\eta\\) measurements are divided into 4 bins with boundaries at \\([-1.5, 0, 1.5]\\), while \\(\phi\\) measurements use 4 bins with boundaries at \\([-\frac{\pi}{2}, 0, \frac{\pi}{2}]\\). As with \\(p_T\\), both \\(\eta\\) and \\(\phi\\) labels default to \\([0, 0, 0, 0]\\) in the absence of a particle. This comprehensive labeling schema enables fine-grained learning of kinematic distributions and particle multiplicities, essential for characterizing complex collision events. + +The loss function combines individual losses from all 41 labels through weighted averaging. Binary cross-entropy is applied to classification labels, while mean-squared error is used for regression labels. The model generates predictions for all labels simultaneously, with individual losses calculated according to their respective types. The final loss is computed as an equally-weighted average across all labels, with weights set to 1 to ensure uniform contribution to the optimization process. The output layer of the multilabel model has 2,688 trainable parameters. + +#### Pretraining + +During pre-training, the initial learning rate is \\(10^{-4}\\), and the learning rate decays by 1% each epoch following the power law function \\(LR(x) = 10^{-4}\cdot(0.99)^x\\), where \\(x\\) is the number of epochs. Both pre-trained models reach a plateau in loss by epoch 50, at which point the training is stopped. + +--- +### Fine-tuning Methodology + +For downstream tasks, we adjust the model architecture for fine-tuning by replacing the original output layer (final linear layer) with a newly initialized linear layer while retaining the pre-trained weights for all other layers. This modification allows the model to specialize in the specific downstream task while leveraging the general features learned during pretraining. + +The fine-tuning process begins with distinct learning rate setups for different parts of the model. The newly initialized linear layer is trained with an initial learning rate of \\(10^{-4}\\), matching the rate used for models trained from scratch. Meanwhile, the pre-trained layers are fine-tuned more cautiously with a lower initial learning rate of \\(10^{-5}\\). This approach ensures that the pre-trained layers adapt gradually without losing their general features, while the new layer learns effectively from scratch. Both learning rates decay over time following the same power law function, \\(LR(x) = LR_{initial} \cdot (0.99)^x\\), to promote stable convergence as training progresses. + +We also evaluated a transfer learning setup in which either the decoder MLP or the final linear layer was replaced with a newly initialized component. During this process, all other model parameters remained frozen, leveraging the pre-trained features without further updating them. However, we did not observe performance improvements using the transfer learning setup. Consequently, we focus on reporting results obtained with the fine-tuning approach. + +--- + +### Performance Evaluation + +We assess model performance using two figures of merit: the classification accuracy and the Area Under the Curve (AUC) of the Receiver Operating Characteristic (ROC) curve. The accuracy is defined as the fraction of correctly classified events when applying a threshold of 0.5 to the neural network output score. Both metrics demonstrate consistent trends in our analysis. + +To obtain reliable performance estimates and uncertainties, we employ an ensemble training approach where 5 independent models are trained for each configuration with random weight initialization and random subsets of the training dataset. This enables us to evaluate both the models' sensitivity to initial parameters and to quantify uncertainties in their performance. + +To investigate how model performance scales with training data, we conducted training runs using sample sizes ranging from \\(10^3\\) to \\(10^7\\) events per class (\\(10^3\\), \\(10^4\\), \\(10^5\\), \\(10^6\\), and \\(10^7\\)) for each model setup: the from-scratch baseline and models fine-tuned from multi-class or multi-label pretrained models. For the \\(10^7\\) case, only the initialization was randomized due to dataset size limitations. All models were evaluated on the same testing dataset, consisting of 2 million events per class, which remained separate from the training process. + +| **Name of Task** | **Pretraining Task** | \\(10^3\\) | \\(10^4\\) | \\(10^5\\) | \\(10^6\\) | \\(10^7\\) | +|----------------------|----------------------|--------------------|--------------------|--------------------|--------------------|--------------------| +| **ttH CP Even vs Odd** | Baseline Accuracy | 56.5 ± 1.1 | 62.2 ± 0.1 | 64.3 ± 0.0 | 65.7 ± 0.0 | 66.2 ± 0.0 | +| | Multiclass (%) | +4.8 ± 1.1 | +3.4 ± 0.1 | +1.3 ± 0.0 | +0.2 ± 0.0 | −0.0 ± 0.0 | +| | Multilabel (%) | +2.1 ± 1.2 | +1.9 ± 0.1 | +0.8 ± 0.1 | +0.0 ± 0.0 | −0.1 ± 0.0 | +| **FCNC vs tHq** | Baseline Accuracy | 63.6 ± 0.7 | 67.8 ± 0.4 | 68.4 ± 0.3 | 69.3 ± 0.3 | 67.9 ± 0.0 | +| | Multiclass (%) | +5.8 ± 0.8 | +1.2 ± 0.4 | +1.4 ± 0.3 | +0.5 ± 0.3 | −0.0 ± 0.0 | +| | Multilabel (%) | −5.3 ± 0.8 | −1.3 ± 0.4 | +0.9 ± 0.4 | +0.3 ± 0.3 | +0.4 ± 0.1 | +| **ttW vs ttt** | Baseline Accuracy | 75.8 ± 0.1 | 77.6 ± 0.1 | 78.9 ± 0.0 | 79.8 ± 0.0 | 80.3 ± 0.0 | +| | Multiclass (%) | +3.7 ± 0.1 | +2.7 ± 0.1 | +1.3 ± 0.0 | +0.4 ± 0.0 | +0.0 ± 0.0 | +| | Multilabel (%) | +2.2 ± 0.1 | +1.1 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.1 ± 0.0 | +| **stop vs ttH** | Baseline Accuracy | 83.0 ± 0.2 | 86.3 ± 0.1 | 87.6 ± 0.0 | 88.5 ± 0.0 | 88.8 ± 0.0 | +| | Multiclass (%) | +0.4 ± 0.2 | +1.9 ± 0.1 | +1.0 ± 0.0 | +0.3 ± 0.0 | +0.0 ± 0.0 | +| | Multilabel (%) | +2.8 ± 0.2 | +1.0 ± 0.1 | +0.5 ± 0.0 | +0.0 ± 0.0 | −0.0 ± 0.0 | +| **WH vs ZH** | Baseline Accuracy | 51.4 ± 0.1 | 53.9 ± 0.1 | 55.8 ± 0.0 | 57.5 ± 0.0 | 58.0 ± 0.0 | +| | Multiclass (%) | +5.2 ± 0.1 | +5.3 ± 0.1 | +3.1 ± 0.0 | +0.6 ± 0.0 | +0.1 ± 0.0 | +| | Multilabel (%) | −1.1 ± 0.1 | −0.9 ± 0.2 | +0.5 ± 0.1 | +0.1 ± 0.0 | −0.1 ± 0.0 | + +> **Table 1**: Accuracy of the traditional model versus the accuracy increase due to fine-tuning from various pretraining tasks. +> The accuracies are averaged over 5 independently trained models with randomly initialized weights and trained on a random subset of the data. One exception is the \\(10^7\\) training where all models use the same dataset due to limitations on our dataset size. The random subsets are allowed to overlap, but this overlap should be very minimal because all models take an independent random subset of \\(10^7\\) events. The testing accuracy is calculated from the same testing set of 2 million events per class across all models for a specific training task. The errors are the propagated errors (root sum of squares) of the standard deviation of accuracies for each model. + +## Results + +### Classification Performance + +Since the observations of AUC and accuracy show similar trends, we focus the presentation of the results using accuracy here for conciseness in Table 1. + +In general, the fine-tuned pretrained model achieves at least the same level of classification performance as the baseline model. Notably, there are significant improvements, particularly when the sample size is small, ranging from \\(10^3\\) to \\(10^4\\) events. In some cases, the accuracy improvements exceed five percentage points, demonstrating that pretrained models provide a strong initial representation that compensates for limited data. The numerical values of the improvements in accuracy may not fully capture the impact on the sensitivity of the measurements for which the neural network classifier is used, and the final sensitivity improvement is likely to be greater. + +As the training sample size grows to \\(10^5\\), \\(10^6\\), and eventually \\(10^7\\) events, the added benefit of pretraining diminishes. With abundant data, models trained from scratch approach or even match the accuracy of fine-tuned pretrained models. This suggests that large datasets enable effective learning from scratch, rendering the advantage of pretraining negligible in such scenarios. + +Although both pretraining approaches offer benefits, multiclass pretraining tends to provide more consistent improvements across tasks, especially in the low-data regime. In contrast, multilabel pretraining can sometimes lead to neutral or even slightly negative effects for certain tasks and data sizes. This highlights the importance of the pretraining task design, as the similarity between pretraining and fine-tuning tasks in the multiclass approach appears to yield better-aligned representations. + +Finally, the spread of accuracy across the five tasks for the baseline model is quite large, offering a robust test of fine-tuning across tasks of varying difficulty. The consistent observation of these trends across tasks confirms the reliability and robustness of the findings. + +--- + +### Model Interpretability + +We aim to understand whether pretrained and baseline models learn the same underlying representations. If the two models exhibit high similarity, a plausible interpretation is that pretraining provides the pretrained model with an advantageous initialization, allowing it to converge to a similar state as the baseline model more efficiently. Conversely, significant differences between the models would indicate that pretraining facilitates the development of a more general and robust latent space, which serves as a foundation for fine-tuning to effectively adapt to the downstream task. To investigate this, we analyzed the representational similarity between a pretrained model fine-tuned for the downstream task and a baseline model trained directly on the downstream task without pretraining. + +We use Centered Kernel Alignment (CKA) [Kornblith et al., 2019][ref-kornblith-2019-cka] to analyze model similarity and interpretability. CKA is a robust metric that quantifies the similarity between the internal representations of neural networks by comparing their feature matrices in a manner that is invariant to scaling, rotation, and alignment. This invariance makes CKA particularly effective for studying relationships between network layers, even across networks of different sizes or those trained from varying initializations. + +The similarity is evaluated using a 64-dimensional latent representation after the decoder stage of the GNN model. This choice allows us to compare the internal states of the models at a fine-grained level and understand how training strategies impact the representations directly used for the output task. + +To provide an intuitive understanding of CKA values, we construct a table of the CKA scores for various transformations performed on a set of dummy data. + +- **A:** randomly initialized matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\)) +- **B:** matrix with shape (1000, 64) constructed via various transformations performed on \\(A\\) +- **Noise:** randomly initialized noise matrix with shape (1000, 64), following a normal distribution (\\(\sigma = 1, \mu = 0\\)) + +| Dataset | CKA Score | +|---------|-----------| +| \\(A, B = A\\) | 1.00 | +| \\(A, B =\\) permutation on columns of \\(A\\) | 1.00 | +| \\(A, B = A + \mathrm{Noise}(0.1)\\) | 0.99 | +| \\(A, B = A + \mathrm{Noise}(0.5)\\) | 0.80 | +| \\(A, B = A + \mathrm{Noise}(0.75)\\) | 0.77 | +| \\(A, B = A \cdot \mathrm{Noise}(1)\\) (Linear Transformation) | 0.76 | +| \\(A, B = A + \mathrm{Noise}(1)\\) | 0.69 | +| \\(A, B = A + \mathrm{Noise}(2)\\) | 0.51 | +| \\(A, B = A + \mathrm{Noise}(5)\\) | 0.39 | + +**Table 2:** CKA scores for a dummy dataset \\(A\\) and \\(B\\), where \\(B\\) is created via various transformations performed on \\(A\\). + +As seen in Table 2 and in the definition of the CKA, the CKA score is permutation-invariant. We will use the CKA score to evaluate the similarity between various models and gain insight into the learned representation of detector events in each model (i.e., the information that each model learns). + +We train ensembles of models for each training task to observe how the CKA score changes due to the random initialization of our models. The CKA score between two models is then defined to be: + +\\[ +CKA(A, B) = \frac{1}{n^2} \sum_i^n \sum_j^n CKA(A_i, B_j) +\\] + +where \\(A_i\\) is the representation learned by the \\(i^{\text{th}}\\) model in an ensemble with \\(n\\) total models. The error in CKA is the standard deviation of \\(CKA(A_i, B_j)\\). + +Here we present results for the CKA similarity between the final model in each setup with the final model in the baseline, shown in Table 3. + +| Training Task | Baseline | Multiclass | Multilabel | +|-----------------------|------------------|-----------------|-----------------| +| ttH CP Even vs Odd | 0.94 ± 0.05 | 0.82 ± 0.01 | 0.77 ± 0.06 | +| FCNC vs tHq | 0.96 ± 0.03 | 0.76 ± 0.01 | 0.81 ± 0.01 | +| ttW vs ttt | 0.91 ± 0.08 | 0.75 ± 0.10 | 0.72 ± 0.05 | +| stop vs ttH | 0.87 ± 0.11 | 0.79 ± 0.12 | 0.71 ± 0.08 | +| WH vs ZH | 0.90 ± 0.07 | 0.53 ± 0.03 | 0.44 ± 0.06 | + +**Table 3:** CKA Similarity of the latent representation before the decoder with the baseline model, averaged over 3 models per training setup, and all models trained with the full dataset (\\(10^7\\)). The baseline column is not guaranteed to be 1.0 because of the random initialization of the model. Each baseline model converges to a slightly different representation as seen in the CKA values in that column. + +The baseline models with different initializations exhibit high similarity values, ranging from approximately 0.87 to 0.96, which indicates that independently trained baseline models tend to converge on similar internal representations despite random initialization. Across the considered tasks, models trained as multi-class or multi-label classifiers exhibit noticeably lower CKA similarity scores when compared to the baseline model. For example, in the WH vs ZH task, the baseline model and another baseline trained model have a high similarity of 0.90, whereas the multi-class and multi-label models show significantly reduced similarities (0.53 and 0.44, respectively). This pattern suggests that the representational spaces developed by multi-class or multi-label models differ substantially from those learned by the baseline model that was trained directly on the downstream classification task. + +### Computational Efficiency + +To estimate the computational resources required for each approach, we measured the wall time needed for a model to reach its final performance. For baseline models, this is defined as the wall time from the start of training until the loss of the model plateaus. For the foundation model approach, the estimate includes both the pretraining time and the fine-tuning time, each measured from the start of training until the loss plateaus. This approach ensures a consistent and comprehensive evaluation of the computational demands. + +![The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.](training_time.png) +*Fig. 1: The ratio of the fine-tuning time required to achieve 99% of the baseline model's final classification accuracy to the total time spent training the baseline model.* + +Figure 1 shows the fine-tuning time for the model pretrained with multiclass classification, relative to the time required for the baseline model, as a function of training sample size. In general, the fine-tuning time is significantly shorter than the training time required by the baseline model approach. For smaller training sets, on the order of \\(10^5\\) events, tasks such as FCNC vs. tHq and ttW vs. ttt benefit substantially from the pretrained model’s “head start,” achieving their final performance in only about 1% of the baseline time. For large training datasets, the fine-tuning time relative to the baseline training time becomes larger; however, given that the large training sample typically requires longer training time, fine-tuning still yields much faster training convergence. The ttH CP-even vs. ttH CP-odd task, with a training sample size of \\(10^7\\) events, is an exception where the fine-tuning time exceeds the training time required for the baseline model. This is likely because the processes involved in this task include photon objects in the final states, which are absent from the events used during pretraining. + +To accurately evaluate the total time consumption, it is necessary to include the pretraining time required for the foundation model approach. The pretraining times are as follows: + +- **Multi-class pretraining:** 45.5 GPU hours +- **Multi-label pretraining:** 60.0 GPU hours + +The GPU hours recorded for the multi-label model represent the total time required when training the model in parallel on 16 GPUs. This includes a model synchronization step, which results in higher GPU hours compared to the multi-class pretraining model. + +The foundation model approach becomes increasingly efficient when a large number of tasks are fine-tuned using the same pretrained model, compared to training each task independently from scratch. To illustrate this, we evaluate the computational time required for a scenario where the training sample contains \\(10^7\\) events. For the five tasks tested in this study, the baseline training time (training from scratch) ranges from 1.68 GPU hours (WH vs. ZH) to 5.30 GPU hours (ttW vs. ttt), with an average baseline training time of 2.94 GPU hours. In contrast, the average fine-tuning time for the foundation model approach, relative to the baseline, is 38% of the baseline training time for \\(10^7\\) events. Based on these averages, we estimate that the foundation model approach becomes more computationally efficient than the baseline approach when fine-tuning is performed for more than 41 tasks. + +As a practical example, the ATLAS measurement of Higgs boson couplings using the \\(H \rightarrow \gamma\gamma\\) decay channel [ATLAS Collaboration, 2023][ref-atlas-2023-higg] involved training 42 classifiers for event categorization. This coincides with our estimate, suggesting that the foundation model approach can reduce computational costs even for a single high-energy physics measurement. + +## Conclusions + +We presented an in-depth study of a particle physics foundation model designed to operate on the four-momentum and identification properties of event final-state objects. This model is built on a Graph Neural Network (GNN) architecture and trained on a dataset comprising 120 million simulated proton-proton collision events across 12 distinct physics processes. The pretraining phase explored both multiclass and multilabel classification tasks, providing a robust foundation for downstream applications. Notably, the pretrained models demonstrated significant improvements in event classification performance when fine-tuned, particularly for tasks with limited training samples. + +The foundation model approach also offers substantial computational advantages. By leveraging fine-tuning, this methodology reduces the computational resources required for large-scale applications across multiple tasks. Our estimates indicate that significant resource savings can be achieved even for single particle physics measurements, making this approach both scalable and efficient. + +To better understand the learned representations of the pretrained model and guide future optimization efforts, we employed a representational similarity evaluation framework using Centered Kernel Alignment (CKA). This metric allowed us to investigate the source of the performance gains observed in the foundation model. Our analysis revealed notable differences in the learned representations between the fine-tuned pretrained model and a baseline model trained from scratch. In deep learning, it is well-established that multiple equally valid solutions can exist. Future studies are necessary to determine whether the low similarity in latent representations reflects complementary information uniquely captured by the foundation and baseline models, or if it can simply be attributed to connected local minima in the loss landscape. + +## Acknowledgments + +This work is supported by the U.S. National Science Foundation under the Award No. 2046280, and by U.S. Department of Energy, Office of Science under contract DE-AC02-05CH11231. + +## References + +- **OpenAI et al.** GPT-4 Technical Report. arXiv:2303.08774 (2024). [https://arxiv.org/abs/2303.08774](https://arxiv.org/abs/2303.08774) + +- **Jason Yosinski, Jeff Clune, Yoshua Bengio, Hod Lipson.** How transferable are features in deep neural networks? CoRR abs/1411.1792 (2014). [http://arxiv.org/abs/1411.1792](http://arxiv.org/abs/1411.1792) + +- **Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, Björn Ommer.** High-Resolution Image Synthesis with Latent Diffusion Models. CoRR abs/2112.10752 (2021). [https://arxiv.org/abs/2112.10752](https://arxiv.org/abs/2112.10752) + +- **Dustin Podell, Zion English, Kyle Lacey et al.** SDXL: Improving Latent Diffusion Models for High-Resolution Image Synthesis. arXiv:2307.01952 (2023). [https://arxiv.org/abs/2307.01952](https://arxiv.org/abs/2307.01952) + +- **John Jumper, Richard Evans, Alexander Pritzel et al.** Highly accurate protein structure prediction with AlphaFold. Nature 596, 583-589 (2021). [https://doi.org/10.1038/s41586-021-03819-2](https://doi.org/10.1038/s41586-021-03819-2) + +- **Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova.** BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding. CoRR abs/1810.04805 (2018). [http://arxiv.org/abs/1810.04805](http://arxiv.org/abs/1810.04805) + +- **ATLAS Collaboration.** Measurement of the properties of Higgs boson production at \\(\sqrt{s} = 13\,\text{TeV}\\) in the \\(H \to \gamma\gamma\\) channel using \\(139\,\text{fb}^{-1}\\) of \\(pp\\) collision data with the ATLAS experiment. JHEP 07 (2023) 088. [arXiv:2207.00348](https://arxiv.org/abs/2207.00348), [https://doi.org/10.1007/JHEP07(2023)088](https://doi.org/10.1007/JHEP07(2023)088) + +- **ATLAS Collaboration.** Observation of four-top-quark production in the multilepton final state with the ATLAS detector. Eur. Phys. J. C 83 (2023) 496. [arXiv:2303.15061](https://arxiv.org/abs/2303.15061), [https://doi.org/10.1140/epjc/s10052-023-11573-0](https://doi.org/10.1140/epjc/s10052-023-11573-0) + +- **Simon Kornblith, Mohammad Norouzi, Honglak Lee, Geoffrey Hinton.** Similarity of Neural Network Representations Revisited. CoRR abs/1905.00414 (2019). [http://arxiv.org/abs/1905.00414](http://arxiv.org/abs/1905.00414) + +--- + + + +- **N. D. Birell, P. C. W. Davies.** Quantum Fields in Curved Space. Cambridge Univ. Press (1982). + +- **R. P. Feynman.** Phys. Rev. 94, 262 (1954). + +- **A. Einstein, Yu. Podolsky, N. Rosen.** Phys. Rev. 47, 777 (1935). + +- **G. P. Berman, Jr., F. M. Izrailev, Jr.** Stability of nonlinear modes. Physica D 88, 445 (1983). + +- **E. B. Davies, L. Parns.** Trapped modes in acoustic waveguides. Q. J. Mech. Appl. Math. 51, 477–492 (1988). + +- **Edward Witten.** hep-th/0106109 (2001). [https://arxiv.org/abs/hep-th/0106109](https://arxiv.org/abs/hep-th/0106109) + +--- + + + +- **E. Beutler.** Williams Hematology, 5th Edition, Chapter 7, pp. 654–662. McGraw-Hill, New York (1994). + +- **Donald E. Knuth.** The Art of Computer Programming vol. 1: Fundamental Algorithms, 2nd Ed., Addison-Wesley (1973). + +- **J. S. Smith, G. W. Johnson.** Philos. Trans. R. Soc. London, Ser. B 777, 1395 (2005). + +- **W. J. Smith, T. J. Johnson, B. G. Miller.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (unpublished, 2010). + +- **V. K. Smith, K. Johnson, M. O. Klein.** Surface chemistry and preferential crystal orientation on a silicon surface. J. Appl. Phys. (submitted, 2010). + +- **Ulrich Underwood, Ned Net, Paul Pot.** Lower Bounds for Wishful Research Results. Talk at Fanstord University (1988). + +- **M. P. Johnson, K. L. Miller, K. Smith.** Personal communication (Jan-May 2007). + +--- + + + +- **Adam Paszke et al.** PyTorch: An Imperative Style, High-Performance Deep Learning Library. arXiv:1912.01703 (2019). [http://arxiv.org/abs/1912.01703](http://arxiv.org/abs/1912.01703) + +- **Minjie Wang et al.** Deep Graph Library: Towards Efficient and Scalable Deep Learning on Graphs. arXiv:1909.01315 (2019). [http://arxiv.org/abs/1909.01315](http://arxiv.org/abs/1909.01315) + +- **Peter W. Battaglia et al.** Relational inductive biases, deep learning, and graph networks. arXiv:1806.01261 (2018). [http://arxiv.org/abs/1806.01261](http://arxiv.org/abs/1806.01261) + +- **Jimmy Lei Ba, Jamie Ryan Kiros, Geoffrey E. Hinton.** Layer Normalization. arXiv:1607.06450 (2016). [https://arxiv.org/abs/1607.06450](https://arxiv.org/abs/1607.06450) + +--- + + + +- **Andrew J. Wildridge et al.** Bumblebee: Foundation Model for Particle Physics Discovery. arXiv:2412.07867 (2024). [https://arxiv.org/abs/2412.07867](https://arxiv.org/abs/2412.07867) + +- **Subash Katel et al.** Learning Symmetry-Independent Jet Representations via Jet-Based Joint Embedding Predictive Architecture. arXiv:2412.05333 (2024). [https://arxiv.org/abs/2412.05333](https://arxiv.org/abs/2412.05333) + +- **Jack Y. Araz et al.** Point cloud-based diffusion models for the Electron-Ion Collider. arXiv:2410.22421 (2024). [https://arxiv.org/abs/2410.22421](https://arxiv.org/abs/2410.22421) + +- **Matthew Leigh et al.** Is Tokenization Needed for Masked Particle Modelling? arXiv:2409.12589 (2024). [https://arxiv.org/abs/2409.12589](https://arxiv.org/abs/2409.12589) + +- **Vinicius Mikuni, Benjamin Nachman.** OmniLearn: A Method to Simultaneously Facilitate All Jet Physics Tasks. arXiv:2404.16091 (2024). [https://arxiv.org/abs/2404.16091](https://arxiv.org/abs/2404.16091) + +- **Zhengde Zhang et al.** Xiwu: A Basis Flexible and Learnable LLM for High Energy Physics. arXiv:2404.08001 (2024). [https://arxiv.org/abs/2404.08001](https://arxiv.org/abs/2404.08001) + +- **Philip Harris et al.** Re-Simulation-based Self-Supervised Learning for Pre-Training Foundation Models. arXiv:2403.07066 (2024). [https://arxiv.org/abs/2403.07066](https://arxiv.org/abs/2403.07066) + +- **Joschka Birk, Anna Hallin, Gregor Kasieczka.** OmniJet-$\alpha$: the first cross-task foundation model for particle physics. Machine Learning: Science and Technology. 5(3), 035031 (Aug 2024). [https://doi.org/10.1088/2632-2153/ad66ad](https://doi.org/10.1088/2632-2153/ad66ad) + +- **Andris Huang et al.** A Language Model for Particle Tracking. arXiv:2402.10239 (2024). [https://arxiv.org/abs/2402.10239](https://arxiv.org/abs/2402.10239) + +- **Tobias Golling et al.** Masked Particle Modeling on Sets: Towards Self-Supervised High Energy Physics Foundation Models. arXiv:2401.13537 (2024). [https://arxiv.org/abs/2401.13537](https://arxiv.org/abs/2401.13537) + +- **Junze Liu et al.** Generalizing to new geometries with Geometry-Aware Autoregressive Models (GAAMs) for fast calorimeter simulation. Journal of Instrumentation 18(11), P11003 (Nov 2023). [https://doi.org/10.1088/1748-0221/18/11/p11003](https://doi.org/10.1088/1748-0221/18/11/p11003) + +- **Baran Hashemi et al.** Ultra-high-granularity detector simulation with intra-event aware generative adversarial network and self-supervised relational reasoning. Nature Communications 15(1) (June 2024). [https://doi.org/10.1038/s41467-024-49104-4](https://doi.org/10.1038/s41467-024-49104-4) + +- **Matthias Vigl et al.** Finetuning Foundation Models for Joint Analysis Optimization. arXiv:2401.13536 (2024). [https://arxiv.org/abs/2401.13536](https://arxiv.org/abs/2401.13536) + +- **Chen Li, Hao Cai, Xianyang Jiang.** Refine neutrino events reconstruction with BEiT-3. Journal of Instrumentation 19(6), T06003 (Jun 2024). [https://doi.org/10.1088/1748-0221/19/06/t06003](https://doi.org/10.1088/1748-0221/19/06/t06003) \ No newline at end of file diff --git a/physicsnemo/configs/config.yaml b/legacy/physicsnemo/configs/config.yaml similarity index 100% rename from physicsnemo/configs/config.yaml rename to legacy/physicsnemo/configs/config.yaml diff --git a/physicsnemo/configs/config_stats_all.yaml b/legacy/physicsnemo/configs/config_stats_all.yaml similarity index 100% rename from physicsnemo/configs/config_stats_all.yaml rename to legacy/physicsnemo/configs/config_stats_all.yaml diff --git a/physicsnemo/configs/tHjb_CP_0_vs_45.yaml b/legacy/physicsnemo/configs/tHjb_CP_0_vs_45.yaml similarity index 100% rename from physicsnemo/configs/tHjb_CP_0_vs_45.yaml rename to legacy/physicsnemo/configs/tHjb_CP_0_vs_45.yaml diff --git a/physicsnemo/configs/tHjb_CP_0_vs_90.yaml b/legacy/physicsnemo/configs/tHjb_CP_0_vs_90.yaml similarity index 100% rename from physicsnemo/configs/tHjb_CP_0_vs_90.yaml rename to legacy/physicsnemo/configs/tHjb_CP_0_vs_90.yaml diff --git a/physicsnemo/configs/tHjb_CP_0_vs_90_edge_network.yaml b/legacy/physicsnemo/configs/tHjb_CP_0_vs_90_edge_network.yaml similarity index 100% rename from physicsnemo/configs/tHjb_CP_0_vs_90_edge_network.yaml rename to legacy/physicsnemo/configs/tHjb_CP_0_vs_90_edge_network.yaml diff --git a/physicsnemo/configs/tHjb_CP_0_vs_90_globals.yaml b/legacy/physicsnemo/configs/tHjb_CP_0_vs_90_globals.yaml similarity index 100% rename from physicsnemo/configs/tHjb_CP_0_vs_90_globals.yaml rename to legacy/physicsnemo/configs/tHjb_CP_0_vs_90_globals.yaml diff --git a/physicsnemo/dataset/Dataset.py b/legacy/physicsnemo/dataset/Dataset.py similarity index 100% rename from physicsnemo/dataset/Dataset.py rename to legacy/physicsnemo/dataset/Dataset.py diff --git a/physicsnemo/dataset/GraphBuilder.py b/legacy/physicsnemo/dataset/GraphBuilder.py similarity index 100% rename from physicsnemo/dataset/GraphBuilder.py rename to legacy/physicsnemo/dataset/GraphBuilder.py diff --git a/physicsnemo/dataset/Graphs.py b/legacy/physicsnemo/dataset/Graphs.py similarity index 100% rename from physicsnemo/dataset/Graphs.py rename to legacy/physicsnemo/dataset/Graphs.py diff --git a/physicsnemo/dataset/Normalization.py b/legacy/physicsnemo/dataset/Normalization.py similarity index 100% rename from physicsnemo/dataset/Normalization.py rename to legacy/physicsnemo/dataset/Normalization.py diff --git a/physicsnemo/metrics.py b/legacy/physicsnemo/metrics.py similarity index 100% rename from physicsnemo/metrics.py rename to legacy/physicsnemo/metrics.py diff --git a/physicsnemo/models/Edge_Network.py b/legacy/physicsnemo/models/Edge_Network.py similarity index 100% rename from physicsnemo/models/Edge_Network.py rename to legacy/physicsnemo/models/Edge_Network.py diff --git a/physicsnemo/models/MeshGraphNet.py b/legacy/physicsnemo/models/MeshGraphNet.py similarity index 100% rename from physicsnemo/models/MeshGraphNet.py rename to legacy/physicsnemo/models/MeshGraphNet.py diff --git a/physicsnemo/models/utils.py b/legacy/physicsnemo/models/utils.py similarity index 100% rename from physicsnemo/models/utils.py rename to legacy/physicsnemo/models/utils.py diff --git a/physicsnemo/setup/Dockerfile b/legacy/physicsnemo/setup/Dockerfile similarity index 100% rename from physicsnemo/setup/Dockerfile rename to legacy/physicsnemo/setup/Dockerfile diff --git a/physicsnemo/setup/build_image.sh b/legacy/physicsnemo/setup/build_image.sh similarity index 100% rename from physicsnemo/setup/build_image.sh rename to legacy/physicsnemo/setup/build_image.sh diff --git a/physicsnemo/train.py b/legacy/physicsnemo/train.py similarity index 100% rename from physicsnemo/train.py rename to legacy/physicsnemo/train.py diff --git a/physicsnemo/utils.py b/legacy/physicsnemo/utils.py similarity index 100% rename from physicsnemo/utils.py rename to legacy/physicsnemo/utils.py diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-data-preparation/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-data-preparation/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-data-preparation/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-data-preparation/SKILL.md diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-env-setup/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-env-setup/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-env-setup/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-env-setup/SKILL.md diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-inference/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-inference/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-inference/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-inference/SKILL.md diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-plotting/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-plotting/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-plotting/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-plotting/SKILL.md diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-training/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-training/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-training/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-training/SKILL.md diff --git a/root_gnn_dgl/.codex/skills/root-gnn-dgl-workflow/SKILL.md b/legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-workflow/SKILL.md similarity index 100% rename from root_gnn_dgl/.codex/skills/root-gnn-dgl-workflow/SKILL.md rename to legacy/root_gnn_dgl/.codex/skills/root-gnn-dgl-workflow/SKILL.md diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/config.yaml diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_0.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_0.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_0.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_0.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_1.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_1.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_1.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_1.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_10.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_10.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_10.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_10.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_11.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_11.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_11.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_11.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_12.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_12.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_12.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_12.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_13.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_13.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_13.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_13.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_14.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_14.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_14.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_14.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_15.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_15.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_15.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_15.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_16.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_16.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_16.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_16.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_17.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_17.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_17.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_17.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_18.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_18.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_18.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_18.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_19.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_19.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_19.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_19.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_2.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_2.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_2.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_2.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_20.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_20.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_20.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_20.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_21.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_21.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_21.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_21.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_22.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_22.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_22.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_22.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_23.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_23.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_23.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_23.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_24.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_24.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_24.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_24.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_25.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_25.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_25.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_25.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_26.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_26.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_26.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_26.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_27.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_27.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_27.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_27.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_28.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_28.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_28.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_28.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_29.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_29.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_29.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_29.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_3.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_3.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_3.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_3.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_30.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_30.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_30.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_30.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_31.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_31.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_31.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_31.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_32.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_32.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_32.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_32.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_33.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_33.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_33.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_33.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_34.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_34.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_34.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_34.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_35.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_35.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_35.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_35.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_36.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_36.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_36.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_36.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_37.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_37.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_37.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_37.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_38.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_38.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_38.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_38.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_39.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_39.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_39.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_39.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_4.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_4.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_4.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_4.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_40.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_40.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_40.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_40.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_41.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_41.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_41.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_41.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_42.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_42.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_42.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_42.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_43.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_43.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_43.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_43.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_44.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_44.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_44.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_44.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_45.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_45.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_45.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_45.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_46.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_46.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_46.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_46.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_47.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_47.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_47.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_47.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_48.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_48.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_48.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_48.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_49.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_49.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_49.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_49.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_5.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_5.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_5.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_5.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_50.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_50.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_50.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_50.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_51.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_51.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_51.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_51.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_52.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_52.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_52.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_52.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_53.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_53.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_53.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_53.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_54.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_54.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_54.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_54.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_55.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_55.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_55.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_55.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_56.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_56.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_56.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_56.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_57.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_57.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_57.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_57.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_58.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_58.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_58.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_58.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_59.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_59.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_59.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_59.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_6.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_6.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_6.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_6.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_60.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_60.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_60.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_60.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_61.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_61.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_61.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_61.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_62.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_62.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_62.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_62.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_63.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_63.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_63.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_63.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_64.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_64.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_64.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_64.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_65.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_65.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_65.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_65.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_66.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_66.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_66.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_66.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_67.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_67.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_67.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_67.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_68.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_68.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_68.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_68.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_69.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_69.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_69.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_69.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_7.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_7.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_7.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_7.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_70.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_70.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_70.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_70.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_8.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_8.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_8.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_8.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_9.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_9.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_9.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_9.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/training.log b/legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/training.log similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/training.log rename to legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/training.log diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/config.yaml diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_0.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_0.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_0.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_0.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_1.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_1.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_1.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_1.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_10.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_10.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_10.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_10.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_11.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_11.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_11.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_11.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_12.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_12.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_12.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_12.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_13.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_13.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_13.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_13.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_14.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_14.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_14.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_14.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_15.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_15.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_15.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_15.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_16.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_16.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_16.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_16.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_17.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_17.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_17.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_17.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_18.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_18.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_18.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_18.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_19.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_19.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_19.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_19.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_2.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_2.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_2.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_2.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_20.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_20.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_20.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_20.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_21.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_21.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_21.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_21.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_22.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_22.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_22.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_22.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_23.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_23.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_23.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_23.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_24.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_24.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_24.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_24.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_25.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_25.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_25.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_25.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_26.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_26.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_26.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_26.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_27.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_27.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_27.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_27.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_28.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_28.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_28.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_28.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_29.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_29.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_29.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_29.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_3.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_3.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_3.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_3.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_30.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_30.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_30.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_30.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_31.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_31.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_31.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_31.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_32.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_32.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_32.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_32.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_33.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_33.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_33.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_33.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_34.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_34.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_34.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_34.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_35.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_35.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_35.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_35.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_36.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_36.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_36.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_36.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_37.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_37.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_37.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_37.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_38.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_38.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_38.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_38.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_39.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_39.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_39.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_39.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_4.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_4.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_4.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_4.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_40.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_40.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_40.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_40.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_41.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_41.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_41.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_41.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_42.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_42.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_42.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_42.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_43.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_43.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_43.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_43.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_44.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_5.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_5.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_5.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_5.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_6.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_6.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_6.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_6.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_7.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_7.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_7.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_7.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_8.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_8.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_8.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_8.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_9.pt b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_9.pt similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_9.pt rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/model_epoch_9.pt diff --git a/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/training.log b/legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/training.log similarity index 100% rename from root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/training.log rename to legacy/root_gnn_dgl/Pretrained_GNN/multilabel_pretrained_model_41/training.log diff --git a/root_gnn_dgl/README.md b/legacy/root_gnn_dgl/README.md similarity index 100% rename from root_gnn_dgl/README.md rename to legacy/root_gnn_dgl/README.md diff --git a/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_baseline.yaml b/legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_baseline.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_baseline.yaml rename to legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_baseline.yaml diff --git a/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multiclass_12_process.yaml b/legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multiclass_12_process.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multiclass_12_process.yaml rename to legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multiclass_12_process.yaml diff --git a/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml b/legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml rename to legacy/root_gnn_dgl/configs/delphes/FCNC_vs_tHjb_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml diff --git a/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_baseline.yaml b/legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_baseline.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_baseline.yaml rename to legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_baseline.yaml diff --git a/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multiclass_12_process.yaml b/legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multiclass_12_process.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multiclass_12_process.yaml rename to legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multiclass_12_process.yaml diff --git a/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml b/legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml rename to legacy/root_gnn_dgl/configs/delphes/WH_vs_ZH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml diff --git a/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_baseline.yaml b/legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_baseline.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_baseline.yaml rename to legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_baseline.yaml diff --git a/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multiclass_12_process.yaml b/legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multiclass_12_process.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multiclass_12_process.yaml rename to legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multiclass_12_process.yaml diff --git a/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml b/legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml rename to legacy/root_gnn_dgl/configs/delphes/stop_vs_ttH_inc_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml diff --git a/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_baseline.yaml b/legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_baseline.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_baseline.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_baseline.yaml diff --git a/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multiclass_12_process.yaml b/legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multiclass_12_process.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multiclass_12_process.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multiclass_12_process.yaml diff --git a/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml b/legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttH_CP_even_vs_odd_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml diff --git a/root_gnn_dgl/configs/delphes/ttW_vs_ttt_baseline.yaml b/legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_baseline.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttW_vs_ttt_baseline.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_baseline.yaml diff --git a/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multiclass_12_process.yaml b/legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multiclass_12_process.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multiclass_12_process.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multiclass_12_process.yaml diff --git a/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml b/legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml similarity index 100% rename from root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml rename to legacy/root_gnn_dgl/configs/delphes/ttW_vs_ttt_finetuning_multilabel_41_higgs_tops_all_kinematics.yaml diff --git a/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml b/legacy/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml rename to legacy/root_gnn_dgl/configs/stats_100K/finetuning_ttH_CP_even_vs_odd.yaml diff --git a/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml b/legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml rename to legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml diff --git a/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml b/legacy/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml rename to legacy/root_gnn_dgl/configs/stats_100K/ttH_CP_even_vs_odd.yaml diff --git a/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml b/legacy/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml rename to legacy/root_gnn_dgl/configs/stats_all/finetuning_ttH_CP_even_vs_odd.yaml diff --git a/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml b/legacy/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml rename to legacy/root_gnn_dgl/configs/stats_all/pretraining_multiclass.yaml diff --git a/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml b/legacy/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml similarity index 100% rename from root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml rename to legacy/root_gnn_dgl/configs/stats_all/ttH_CP_even_vs_odd.yaml diff --git a/root_gnn_dgl/jobs/cpu.sh b/legacy/root_gnn_dgl/jobs/cpu.sh similarity index 100% rename from root_gnn_dgl/jobs/cpu.sh rename to legacy/root_gnn_dgl/jobs/cpu.sh diff --git a/root_gnn_dgl/jobs/inference/run_inference.py b/legacy/root_gnn_dgl/jobs/inference/run_inference.py similarity index 100% rename from root_gnn_dgl/jobs/inference/run_inference.py rename to legacy/root_gnn_dgl/jobs/inference/run_inference.py diff --git a/root_gnn_dgl/jobs/interactive.sh b/legacy/root_gnn_dgl/jobs/interactive.sh similarity index 100% rename from root_gnn_dgl/jobs/interactive.sh rename to legacy/root_gnn_dgl/jobs/interactive.sh diff --git a/root_gnn_dgl/jobs/prep_data/parallel_prep.py b/legacy/root_gnn_dgl/jobs/prep_data/parallel_prep.py similarity index 100% rename from root_gnn_dgl/jobs/prep_data/parallel_prep.py rename to legacy/root_gnn_dgl/jobs/prep_data/parallel_prep.py diff --git a/root_gnn_dgl/jobs/prep_data/prep_data.sh b/legacy/root_gnn_dgl/jobs/prep_data/prep_data.sh similarity index 100% rename from root_gnn_dgl/jobs/prep_data/prep_data.sh rename to legacy/root_gnn_dgl/jobs/prep_data/prep_data.sh diff --git a/root_gnn_dgl/jobs/prep_data/run_processing.py b/legacy/root_gnn_dgl/jobs/prep_data/run_processing.py similarity index 100% rename from root_gnn_dgl/jobs/prep_data/run_processing.py rename to legacy/root_gnn_dgl/jobs/prep_data/run_processing.py diff --git a/root_gnn_dgl/jobs/salloc.sh b/legacy/root_gnn_dgl/jobs/salloc.sh similarity index 100% rename from root_gnn_dgl/jobs/salloc.sh rename to legacy/root_gnn_dgl/jobs/salloc.sh diff --git a/root_gnn_dgl/jobs/training/conda/run_job.sh b/legacy/root_gnn_dgl/jobs/training/conda/run_job.sh similarity index 100% rename from root_gnn_dgl/jobs/training/conda/run_job.sh rename to legacy/root_gnn_dgl/jobs/training/conda/run_job.sh diff --git a/root_gnn_dgl/jobs/training/conda/submit.sh b/legacy/root_gnn_dgl/jobs/training/conda/submit.sh similarity index 100% rename from root_gnn_dgl/jobs/training/conda/submit.sh rename to legacy/root_gnn_dgl/jobs/training/conda/submit.sh diff --git a/root_gnn_dgl/jobs/training/podman/run_job.sh b/legacy/root_gnn_dgl/jobs/training/podman/run_job.sh similarity index 100% rename from root_gnn_dgl/jobs/training/podman/run_job.sh rename to legacy/root_gnn_dgl/jobs/training/podman/run_job.sh diff --git a/root_gnn_dgl/jobs/training/podman/run_job_image.sh b/legacy/root_gnn_dgl/jobs/training/podman/run_job_image.sh similarity index 100% rename from root_gnn_dgl/jobs/training/podman/run_job_image.sh rename to legacy/root_gnn_dgl/jobs/training/podman/run_job_image.sh diff --git a/root_gnn_dgl/jobs/training/podman/submit.sh b/legacy/root_gnn_dgl/jobs/training/podman/submit.sh similarity index 100% rename from root_gnn_dgl/jobs/training/podman/submit.sh rename to legacy/root_gnn_dgl/jobs/training/podman/submit.sh diff --git a/root_gnn_dgl/jobs/training/run_parallel_trainings.py b/legacy/root_gnn_dgl/jobs/training/run_parallel_trainings.py similarity index 100% rename from root_gnn_dgl/jobs/training/run_parallel_trainings.py rename to legacy/root_gnn_dgl/jobs/training/run_parallel_trainings.py diff --git a/root_gnn_dgl/models/GCN.py b/legacy/root_gnn_dgl/models/GCN.py similarity index 100% rename from root_gnn_dgl/models/GCN.py rename to legacy/root_gnn_dgl/models/GCN.py diff --git a/root_gnn_dgl/models/loss.py b/legacy/root_gnn_dgl/models/loss.py similarity index 100% rename from root_gnn_dgl/models/loss.py rename to legacy/root_gnn_dgl/models/loss.py diff --git a/root_gnn_dgl/profile.sh b/legacy/root_gnn_dgl/profile.sh similarity index 100% rename from root_gnn_dgl/profile.sh rename to legacy/root_gnn_dgl/profile.sh diff --git a/root_gnn_dgl/root_gnn_base/batched_dataset.py b/legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/batched_dataset.py rename to legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py diff --git a/root_gnn_dgl/root_gnn_base/custom_scheduler.py b/legacy/root_gnn_dgl/root_gnn_base/custom_scheduler.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/custom_scheduler.py rename to legacy/root_gnn_dgl/root_gnn_base/custom_scheduler.py diff --git a/root_gnn_dgl/root_gnn_base/dataset.py b/legacy/root_gnn_dgl/root_gnn_base/dataset.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/dataset.py rename to legacy/root_gnn_dgl/root_gnn_base/dataset.py diff --git a/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py b/legacy/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/photon_ID_dataset.py rename to legacy/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py diff --git a/root_gnn_dgl/root_gnn_base/similarity.py b/legacy/root_gnn_dgl/root_gnn_base/similarity.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/similarity.py rename to legacy/root_gnn_dgl/root_gnn_base/similarity.py diff --git a/root_gnn_dgl/root_gnn_base/uproot_dataset.py b/legacy/root_gnn_dgl/root_gnn_base/uproot_dataset.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/uproot_dataset.py rename to legacy/root_gnn_dgl/root_gnn_base/uproot_dataset.py diff --git a/root_gnn_dgl/root_gnn_base/utils.py b/legacy/root_gnn_dgl/root_gnn_base/utils.py similarity index 100% rename from root_gnn_dgl/root_gnn_base/utils.py rename to legacy/root_gnn_dgl/root_gnn_base/utils.py diff --git a/root_gnn_dgl/run_demo.sh b/legacy/root_gnn_dgl/run_demo.sh similarity index 100% rename from root_gnn_dgl/run_demo.sh rename to legacy/root_gnn_dgl/run_demo.sh diff --git a/root_gnn_dgl/scripts/check_dataset_files.py b/legacy/root_gnn_dgl/scripts/check_dataset_files.py similarity index 100% rename from root_gnn_dgl/scripts/check_dataset_files.py rename to legacy/root_gnn_dgl/scripts/check_dataset_files.py diff --git a/root_gnn_dgl/scripts/export_onnx.py b/legacy/root_gnn_dgl/scripts/export_onnx.py similarity index 100% rename from root_gnn_dgl/scripts/export_onnx.py rename to legacy/root_gnn_dgl/scripts/export_onnx.py diff --git a/root_gnn_dgl/scripts/find_free_port.py b/legacy/root_gnn_dgl/scripts/find_free_port.py similarity index 100% rename from root_gnn_dgl/scripts/find_free_port.py rename to legacy/root_gnn_dgl/scripts/find_free_port.py diff --git a/root_gnn_dgl/scripts/inference.py b/legacy/root_gnn_dgl/scripts/inference.py similarity index 100% rename from root_gnn_dgl/scripts/inference.py rename to legacy/root_gnn_dgl/scripts/inference.py diff --git a/root_gnn_dgl/scripts/plot_config_distributions.py b/legacy/root_gnn_dgl/scripts/plot_config_distributions.py similarity index 100% rename from root_gnn_dgl/scripts/plot_config_distributions.py rename to legacy/root_gnn_dgl/scripts/plot_config_distributions.py diff --git a/root_gnn_dgl/scripts/prep_data.py b/legacy/root_gnn_dgl/scripts/prep_data.py similarity index 100% rename from root_gnn_dgl/scripts/prep_data.py rename to legacy/root_gnn_dgl/scripts/prep_data.py diff --git a/root_gnn_dgl/scripts/selections.py b/legacy/root_gnn_dgl/scripts/selections.py similarity index 100% rename from root_gnn_dgl/scripts/selections.py rename to legacy/root_gnn_dgl/scripts/selections.py diff --git a/root_gnn_dgl/scripts/training_script.py b/legacy/root_gnn_dgl/scripts/training_script.py similarity index 100% rename from root_gnn_dgl/scripts/training_script.py rename to legacy/root_gnn_dgl/scripts/training_script.py diff --git a/root_gnn_dgl/setup/Dockerfile b/legacy/root_gnn_dgl/setup/Dockerfile similarity index 100% rename from root_gnn_dgl/setup/Dockerfile rename to legacy/root_gnn_dgl/setup/Dockerfile diff --git a/root_gnn_dgl/setup/build_image.sh b/legacy/root_gnn_dgl/setup/build_image.sh similarity index 100% rename from root_gnn_dgl/setup/build_image.sh rename to legacy/root_gnn_dgl/setup/build_image.sh diff --git a/root_gnn_dgl/setup/download_data.sh b/legacy/root_gnn_dgl/setup/download_data.sh similarity index 100% rename from root_gnn_dgl/setup/download_data.sh rename to legacy/root_gnn_dgl/setup/download_data.sh diff --git a/root_gnn_dgl/setup/environment.yml b/legacy/root_gnn_dgl/setup/environment.yml similarity index 100% rename from root_gnn_dgl/setup/environment.yml rename to legacy/root_gnn_dgl/setup/environment.yml diff --git a/root_gnn_dgl/setup/launch_image.sh b/legacy/root_gnn_dgl/setup/launch_image.sh similarity index 100% rename from root_gnn_dgl/setup/launch_image.sh rename to legacy/root_gnn_dgl/setup/launch_image.sh diff --git a/root_gnn_dgl/setup/test_setup.py b/legacy/root_gnn_dgl/setup/test_setup.py similarity index 100% rename from root_gnn_dgl/setup/test_setup.py rename to legacy/root_gnn_dgl/setup/test_setup.py diff --git a/training_time.png b/legacy/training_time.png similarity index 100% rename from training_time.png rename to legacy/training_time.png diff --git a/notebooks/.gitkeep b/notebooks/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/notebooks/.gitkeep @@ -0,0 +1 @@ + diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000000000000000000000000000000000000..467610f9ef8a0dce55a5a44de9efa4da445892c7 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,98 @@ +[build-system] +requires = ["setuptools>=68"] +build-backend = "setuptools.build_meta" + +[project] +name = "gnn4colliders" +dynamic = ["version"] +description = "Graph neural network tools for collider event analysis." +readme = "README.md" +requires-python = ">=3.12,<3.13" +dependencies = [ + "awkward>=2.0", + "hydra-core>=1.3", + "numpy>=1.24,<2", + "scikit-learn>=1.3", + "torch==2.2.2", + "uproot>=5.0", +] + +[project.optional-dependencies] +root-gnn = [ + "dgl==2.4.0+cu121", +] +onnx = [ + "onnx>=1.16,<2", + "onnxruntime>=1.18,<2", +] + +[project.scripts] +gnn4colliders = "gnn4colliders.cli:main" + +[dependency-groups] +dev = [ + "matplotlib>=3.8,<4", + "pytest>=8", + "pytest-cov>=5", + "ruff>=0.6", + "twine>=6,<7", + "onnx>=1.16,<2", + "onnxruntime>=1.18,<2", +] + +[tool.setuptools.packages.find] +where = ["src"] + +[tool.setuptools.dynamic] +version = { attr = "gnn4colliders.__version__" } + +[tool.setuptools.package-data] +"gnn4colliders.configs" = ["*.yaml", "*/*.yaml", "*/*/*.yaml"] + +[tool.uv.sources] +torch = { index = "pytorch-cu121" } +dgl = { index = "dgl" } + +[tool.uv] +environments = [ + "sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'", +] + +[[tool.uv.index]] +name = "pytorch-cu121" +url = "https://download.pytorch.org/whl/cu121" +explicit = true + +[[tool.uv.index]] +name = "dgl" +url = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html" +format = "flat" +explicit = true + + +[tool.pytest.ini_options] +testpaths = ["tests"] +addopts = "-ra" +markers = [ + "integration: small end-to-end tests using local fixtures", + "distributed: tests requiring torch distributed execution", + "legacy_env: tests requiring the historical runtime environment", + "gpu: tests requiring CUDA", + "parity: comparisons with the frozen legacy implementation", + "onnx: tests requiring ONNX export/runtime dependencies", + "real_data: tests requiring an optional reduced ROOT fixture", + "slow: tests intentionally excluded from the fast development loop", +] + +[tool.ruff] +line-length = 88 +target-version = "py312" +src = ["src"] +exclude = ["legacy", "tasks"] + +[tool.ruff.lint] +select = ["E", "F", "I"] +exclude = ["legacy", "tasks"] + +[tool.ruff.format] +exclude = ["legacy"] diff --git a/scripts/.gitkeep b/scripts/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/scripts/.gitkeep @@ -0,0 +1 @@ + diff --git a/scripts/dev/.gitkeep b/scripts/dev/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/scripts/dev/.gitkeep @@ -0,0 +1 @@ + diff --git a/scripts/dev/check_release.sh b/scripts/dev/check_release.sh new file mode 100755 index 0000000000000000000000000000000000000000..961d8cbade068d13388a0cef9e1e3fdaf13c8f98 --- /dev/null +++ b/scripts/dev/check_release.sh @@ -0,0 +1,29 @@ +#!/usr/bin/env bash +set -euo pipefail + +repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)" +cd "$repo_root" + +rm -rf dist +uv run ruff check . +uv run ruff format --check . +uv run pytest +uv build + +uv run python -m twine check dist/* + +wheel="$(find dist -maxdepth 1 -type f -name '*.whl' -print -quit)" +test_root="$(mktemp -d)" +trap 'rm -rf "$test_root"' EXIT +python -m venv "$test_root/venv" +"$test_root/venv/bin/python" -m pip install --quiet "$wheel" +( + cd "$test_root" + "$test_root/venv/bin/python" -c \ + 'import gnn4colliders; assert gnn4colliders.__version__ == "0.1.0"' + "$test_root/venv/bin/gnn4colliders" --help + "$test_root/venv/bin/gnn4colliders" train --help + "$test_root/venv/bin/gnn4colliders" evaluate --help + "$test_root/venv/bin/gnn4colliders" predict --help + "$test_root/venv/bin/gnn4colliders" export --help +) diff --git a/scripts/dev/shrink_root_sample.py b/scripts/dev/shrink_root_sample.py new file mode 100644 index 0000000000000000000000000000000000000000..24d1a03bc7aa7e470720cf4fe2dfab861f40ea7d --- /dev/null +++ b/scripts/dev/shrink_root_sample.py @@ -0,0 +1,57 @@ +"""Create a small ROOT fixture from a Delphes sample.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import uproot + +BRANCHES = ( + "jet_pt", + "jet_eta", + "jet_phi", + "jet_btag", + "ph_pt", + "ph_eta", + "ph_phi", + "ele_pt", + "ele_eta", + "ele_phi", + "ele_charge", + "mu_pt", + "mu_eta", + "mu_phi", + "mu_charge", + "MET_met", + "MET_phi", + "weight", + "Number", +) + + +def shrink_sample(source: Path, target: Path, entries: int) -> None: + """Copy the active branches and first ``entries`` events to ``target``.""" + if entries < 1: + raise ValueError("entries must be positive") + target.parent.mkdir(parents=True, exist_ok=True) + with uproot.open(source) as source_file: + arrays = source_file["output"].arrays( + BRANCHES, entry_start=0, entry_stop=entries, library="ak" + ) + with uproot.recreate(target) as target_file: + target_file["output"] = arrays + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("source", type=Path) + parser.add_argument("target", type=Path) + parser.add_argument("--entries", type=int, default=64) + args = parser.parse_args() + shrink_sample(args.source, args.target, args.entries) + print(f"Wrote {args.entries} events to {args.target}") + + +if __name__ == "__main__": + main() diff --git a/scripts/dev/smoke_end_to_end.py b/scripts/dev/smoke_end_to_end.py new file mode 100644 index 0000000000000000000000000000000000000000..a97654cb82823103ae3b677a0eae8b723241594b --- /dev/null +++ b/scripts/dev/smoke_end_to_end.py @@ -0,0 +1,131 @@ +"""Run the documented new-stack workflow on a tiny temporary ROOT sample. + +This is an integration smoke test, not a physics example. It requires the +validated ``root-gnn`` extra and uses only the public CLI for preparation, +training, evaluation, and prediction. +""" + +from __future__ import annotations + +import subprocess +import sys +import tempfile +from pathlib import Path + +import awkward as ak +import numpy as np +import uproot + + +def _run(root: Path, *arguments: str) -> None: + command = [sys.executable, "-m", "gnn4colliders.cli", *arguments] + subprocess.run(command, cwd=root, check=True) + + +def _write_root(path: Path) -> None: + events = 5 + with uproot.recreate(path) as output: + output["Events"] = { + "jet_pt": ak.Array([[40.0, 25.0]] * events), + "jet_eta": ak.Array([[-1.0, 1.0]] * events), + "jet_phi": ak.Array([[-2.5, 2.5]] * events), + "eventNumber": np.arange(events, dtype=np.int64), + "weight": np.ones(events, dtype=np.float32), + } + + +def _prepare(root_file: Path, cache: Path) -> None: + _run( + root_file.parent, + "prepare", + f"data.files=[{root_file}]", + "data.tree_name=Events", + f"data.cache.path={cache}", + '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]", + "data.fold_var=eventNumber", + "data.weight_var=weight", + ) + + +def _train(root: Path, cache: Path, output: Path, *extra: str) -> Path: + _run( + root, + "train", + f"data.cache.path={cache}", + "trainer.max_epochs=1", + "trainer.device=cpu", + "data.batch_size=1", + "model.hid_size=8", + "model.n_layers=1", + "model.n_proc_steps=1", + f"environment.output_root={output}", + *extra, + ) + return output / "checkpoints" / "epoch_0000.pt" + + +def main() -> None: + try: + import dgl # noqa: F401 + except ImportError as error: # pragma: no cover - environment-dependent + raise SystemExit( + "install the root-gnn extra before running this smoke test" + ) from error + + with tempfile.TemporaryDirectory(prefix="gnn4colliders-smoke-") as directory: + root = Path(directory) + root_file = root / "events.root" + cache = root / "graphs.pt" + target_cache = root / "target.pt" + _write_root(root_file) + _prepare(root_file, cache) + pretrained = _train(root, cache, root / "pretrain") + _prepare(root_file, target_cache) + fine_tuned = _train( + root, + target_cache, + root / "finetune", + "model=root_gnn/fine_tuned_edge_network", + "task=binary_classification", + f"checkpoint.pretrained={pretrained}", + "model.freeze_backbone=true", + ) + _run( + root, + "evaluate", + f"data.cache.path={target_cache}", + "inference.split=test", + f"inference.checkpoint={fine_tuned}", + "model=root_gnn/fine_tuned_edge_network", + "task=binary_classification", + f"checkpoint.pretrained={pretrained}", + "model.hid_size=8", + "model.n_layers=1", + "model.n_proc_steps=1", + "trainer.device=cpu", + ) + prediction = root / "predictions.npz" + _run( + root, + "predict", + f"data.cache.path={target_cache}", + "inference.split=test", + f"inference.checkpoint={fine_tuned}", + "model=root_gnn/fine_tuned_edge_network", + "task=binary_classification", + f"checkpoint.pretrained={pretrained}", + "model.hid_size=8", + "model.n_layers=1", + "model.n_proc_steps=1", + f"inference.output={prediction}", + "trainer.device=cpu", + ) + if not prediction.is_file(): + raise RuntimeError("smoke workflow did not produce predictions.npz") + print(f"smoke workflow succeeded in {root}") + + +if __name__ == "__main__": + main() diff --git a/src/gnn4colliders/__init__.py b/src/gnn4colliders/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..c68fe192f771ebebb70dfab6a85ded1f0783cfca --- /dev/null +++ b/src/gnn4colliders/__init__.py @@ -0,0 +1,5 @@ +"""Public package metadata for GNN4Colliders.""" + +__version__ = "0.1.0" + +__all__ = ["__version__"] diff --git a/src/gnn4colliders/cli/__init__.py b/src/gnn4colliders/cli/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..a0e5f1b98cc6dc371cb1dc562439de4635b30cca --- /dev/null +++ b/src/gnn4colliders/cli/__init__.py @@ -0,0 +1,91 @@ +"""Single, thin command-line entry point for GNN4Colliders.""" + +from __future__ import annotations + +import logging +import os +import sys +from importlib import resources + +from hydra import compose, initialize_config_dir +from omegaconf import OmegaConf + +from gnn4colliders.config.application import ( + export_onnx, + predict_or_evaluate, + prepare, + train, +) + +_COMMANDS = {"prepare", "train", "evaluate", "predict", "export"} + + +def _is_main_process() -> bool: + return int(os.environ.get("RANK", "0")) == 0 + + +def _help(command: str | None = None) -> None: + if command: + print(f"Usage: gnn4colliders {command} [key=value ...]") + print("Compose Hydra configuration with semantic overrides.") + else: + print( + "Usage: gnn4colliders " + "[key=value ...]" + ) + print("Use gnn4colliders --help for command help.") + + +def _config(overrides: list[str]): + config_resource = resources.files("gnn4colliders.configs") + with resources.as_file(config_resource) as config_dir: + with initialize_config_dir(version_base=None, config_dir=str(config_dir)): + return compose(config_name="config", overrides=overrides) + + +def main(argv: list[str] | None = None) -> int: + args = list(sys.argv[1:] if argv is None else argv) + if not args or args[0] in {"-h", "--help"}: + _help() + return 0 + command = args.pop(0) + if command not in _COMMANDS: + print( + f"unknown command {command!r}; choose from {', '.join(sorted(_COMMANDS))}", + file=sys.stderr, + ) + return 2 + if "--help" in args or "-h" in args: + _help(command) + return 0 + try: + config = _config(args) + logging.basicConfig( + level=getattr(logging, str(config.logging.level).upper(), logging.INFO) + ) + if command == "prepare": + print(f"prepared cache: {prepare(config)}") + elif command == "train": + if _is_main_process(): + print(f"saved checkpoint: {train(config)}") + elif command == "export": + if _is_main_process(): + print(f"wrote ONNX model: {export_onnx(config)}") + elif command == "evaluate": + result = predict_or_evaluate(config, evaluate=True) + if _is_main_process(): + print(OmegaConf.to_yaml(result.metrics)) + else: + result = predict_or_evaluate(config) + if _is_main_process(): + print( + f"wrote predictions: {config.inference.output} " + f"({len(result.sample_ids)} events)" + ) + return 0 + except (ValueError, FileNotFoundError, KeyError, ImportError) as error: + print(f"gnn4colliders: {error}", file=sys.stderr) + return 2 + + +__all__ = ["main"] diff --git a/src/gnn4colliders/cli/__main__.py b/src/gnn4colliders/cli/__main__.py new file mode 100644 index 0000000000000000000000000000000000000000..aa71a1aa399dfff76fb0f9fffaae25ab64334227 --- /dev/null +++ b/src/gnn4colliders/cli/__main__.py @@ -0,0 +1,4 @@ +from . import main + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/gnn4colliders/compat/__init__.py b/src/gnn4colliders/compat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..edf69f860a117b899d147e538c89db6fc110998c --- /dev/null +++ b/src/gnn4colliders/compat/__init__.py @@ -0,0 +1,17 @@ +"""Explicit adapters for artifacts from the frozen legacy implementation.""" + +from .checkpoint import ( + LEGACY_CHECKPOINT_SCHEMA_VERSION, + load_legacy_checkpoint, + map_legacy_edge_network_state_dict, + normalize_legacy_state_dict_keys, +) +from .metadata import event_metadata_from_legacy_tracking + +__all__ = [ + "LEGACY_CHECKPOINT_SCHEMA_VERSION", + "event_metadata_from_legacy_tracking", + "load_legacy_checkpoint", + "map_legacy_edge_network_state_dict", + "normalize_legacy_state_dict_keys", +] diff --git a/src/gnn4colliders/compat/checkpoint.py b/src/gnn4colliders/compat/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..d1d8920e33584ac20252b5b54fb1fc338ec024a2 --- /dev/null +++ b/src/gnn4colliders/compat/checkpoint.py @@ -0,0 +1,72 @@ +"""Checkpoint adapters for the supported historical ROOT-GNN artifacts.""" + +from __future__ import annotations + +from collections.abc import Mapping +from pathlib import Path +from typing import Any + +import torch + +LEGACY_CHECKPOINT_SCHEMA_VERSION = 0 + + +def normalize_legacy_state_dict_keys(state: Mapping[str, Any]) -> dict[str, Any]: + """Remove the supported DDP and compilation wrappers deterministically.""" + return { + key.removeprefix("module.").removeprefix("_orig_mod."): value + for key, value in state.items() + } + + +def map_legacy_edge_network_state_dict( + state: Mapping[str, Any], +) -> dict[str, Any]: + """Map historical ``classify`` keys to the modern ``classifier`` name.""" + normalized = normalize_legacy_state_dict_keys(state) + return { + key.replace(".classify.", ".classifier.").replace( + "classify.", "classifier." + ): value + for key, value in normalized.items() + } + + +def load_legacy_checkpoint( + source: str | Path | Mapping[str, Any], *, map_location: Any = "cpu" +) -> dict[str, Any]: + """Adapt an active ``model_epoch_N.pt`` payload without rewriting it.""" + payload = ( + torch.load(source, map_location=map_location, weights_only=False) + if not isinstance(source, Mapping) + else dict(source) + ) + if not isinstance(payload, Mapping) or "model_state_dict" not in payload: + raise ValueError("not a supported legacy ROOT-GNN checkpoint") + early = payload.get("early_stop") + if isinstance(early, Mapping): + early = { + "patience": early.get("patience", 15), + "min_delta": early.get("threshold", 1e-8), + "mode": early.get("mode", "min"), + "best": early.get("current_best", float("inf")), + "num_bad_epochs": early.get("count", 0), + "should_stop": early.get("should_stop", False), + } + epoch = int(payload.get("epoch", -1)) + return { + "schema_version": LEGACY_CHECKPOINT_SCHEMA_VERSION, + "legacy": True, + "epoch": epoch, + "global_step": 0, + "model_state_dict": normalize_legacy_state_dict_keys( + payload["model_state_dict"] + ), + "optimizer_state_dict": payload.get("optimizer_state_dict"), + "scheduler_state_dict": None, + "early_stopping_state": early, + "trainer_state": {"epoch": epoch, "global_step": 0}, + "model_config": None, + "task_config": None, + "metadata": {"legacy_format": True}, + } diff --git a/src/gnn4colliders/compat/metadata.py b/src/gnn4colliders/compat/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..b9baaabf89ccd7e38111988faf83fc09e797d832 --- /dev/null +++ b/src/gnn4colliders/compat/metadata.py @@ -0,0 +1,32 @@ +"""Adapters for the historical positional event metadata layout.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from typing import Any + +from gnn4colliders.data.metadata import EventMetadata + + +def event_metadata_from_legacy_tracking( + tracking_row: Sequence[Any], + *, + sample_id: str, + extra: Mapping[str, Any] | None = None, +) -> EventMetadata: + """Convert the only supported legacy layout into named metadata. + + The historical contract is strictly ``tracking[0] = fold`` and + ``tracking[1] = weight``. Additional historical columns are not + interpreted or propagated. + """ + if len(tracking_row) < 2: + raise ValueError( + "legacy tracking must contain at least two values: fold and weight" + ) + return EventMetadata( + fold=int(tracking_row[0]), + weight=float(tracking_row[1]), + sample_id=sample_id, + extra=dict(extra or {}), + ) diff --git a/src/gnn4colliders/config/__init__.py b/src/gnn4colliders/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..54b9e3070fe23551d8f3f68454a070a7e4c0a2f4 --- /dev/null +++ b/src/gnn4colliders/config/__init__.py @@ -0,0 +1,6 @@ +"""Semantic configuration factories used by the command-line applications.""" + +from .factories import build_model, build_task, build_trainer +from .validation import validate_config + +__all__ = ["build_model", "build_task", "build_trainer", "validate_config"] diff --git a/src/gnn4colliders/config/application.py b/src/gnn4colliders/config/application.py new file mode 100644 index 0000000000000000000000000000000000000000..8e80ee2dab86f6396c6e61b66fc724b757279c1e --- /dev/null +++ b/src/gnn4colliders/config/application.py @@ -0,0 +1,282 @@ +"""Application-level orchestration kept separate from the CLI parser.""" + +from __future__ import annotations + +import logging +import os +from pathlib import Path +from typing import Any + +import torch +from omegaconf import DictConfig, OmegaConf + +from gnn4colliders.data import ( + DistributedGraphDataLoader, + GraphDataLoader, + GraphDataset, + GraphSample, + GraphSampleCache, + RootEventDataset, + SplitDefinition, + select_split, +) +from gnn4colliders.distributed import ( + DistributedContext, + barrier, + finalize, + initialize, +) +from gnn4colliders.features import build_node_features +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.inference import Predictor, write_npz +from gnn4colliders.training import ( + CheckpointManager, + restore_training_state, + seed_everything, +) + +from .factories import build_model, build_task, build_trainer +from .validation import validate_config + +logger = logging.getLogger(__name__) + + +def plain(config: DictConfig | dict[str, Any]) -> dict[str, Any]: + return OmegaConf.to_container(config, resolve=True) # type: ignore[return-value] + + +def save_resolved(config: DictConfig, output_root: str | Path) -> Path: + target = Path(output_root) / "resolved_config.yaml" + target.parent.mkdir(parents=True, exist_ok=True) + OmegaConf.save(config, target, resolve=True) + return target + + +def load_samples(config: DictConfig) -> GraphDataset: + data = config.data + cache_path = data.cache.path + if not cache_path: + raise ValueError("data.cache.path is required for train/evaluate/predict") + cache = GraphSampleCache(cache_path) + if not cache.exists(): + raise FileNotFoundError(f"graph sample cache does not exist: {cache_path}") + return GraphDataset(cache.load()) + + +def loaders( + config: DictConfig, context: DistributedContext | None = None +) -> dict[str, GraphDataLoader]: + samples = load_samples(config) + split = SplitDefinition( + train_folds=frozenset(config.data.splits.train_folds), + validation_folds=frozenset(config.data.splits.validation_folds), + test_folds=frozenset(config.data.splits.test_folds), + ) + result = {} + for name in ("train", "validation", "test"): + selected = GraphDataset(select_split(samples.samples, split, name)) + loader_type = ( + DistributedGraphDataLoader + if context is not None and context.enabled + else GraphDataLoader + ) + result[name] = ( + loader_type( + selected, + int(config.data.batch_size), + context, + shuffle=bool(config.data.shuffle) if name == "train" else False, + seed=int(config.data.seed), + ) + if loader_type is DistributedGraphDataLoader + else loader_type( + selected, + int(config.data.batch_size), + shuffle=bool(config.data.shuffle) if name == "train" else False, + seed=int(config.data.seed), + ) + ) + return result + + +def prepare(config: DictConfig) -> Path: + if bool(config.distributed.enabled) or int(os.environ.get("WORLD_SIZE", "1")) > 1: + raise ValueError( + "prepare does not support distributed launch; run it once on rank 0" + ) + data = config.data + if not data.files: + raise ValueError("prepare requires data.files") + feature_branches = plain(data.get("feature_branches")) + object_types = plain(data.get("object_types")) + scales = plain(data.get("scales")) + if feature_branches is None or object_types is None or scales is None: + raise ValueError( + "prepare requires data.feature_branches, object_types, and scales" + ) + source = RootEventDataset( + data.files, + tree_name=str(data.tree_name), + label=data.get("label", 1), + feature_branches=feature_branches, + global_features=plain(data.get("global_features", [])), + fold_var=str(data.get("fold_var", "eventNumber")), + weight_var=data.get("weight_var"), + ) + samples: list[GraphSample] = [] + for event in source: + features, _ = build_node_features( + event.objects, feature_branches, object_types, scales + ) + graph = build_dgl_graph(features) + globals_ = torch.as_tensor(event.global_features, dtype=torch.float32) + samples.append( + GraphSample( + graph, + torch.as_tensor(event.label), + globals_ if globals_.numel() else None, + event.event_metadata, + ) + ) + path = Path(data.cache.path) + GraphSampleCache(path).save(samples) + return path + + +def _model_and_task(config: DictConfig, loader: GraphDataLoader): + first = next(iter(loader), None) + if first is None: + raise ValueError("selected data split is empty") + model_config = plain(config.model) + pretrained = config.checkpoint.pretrained + if pretrained: + model_config["pretrained_checkpoint"] = pretrained + model = build_model( + model_config, sample_graph=first.graph, sample_global=first.global_features + ) + return model, build_task(plain(config.task)) + + +def train(config: DictConfig) -> Path: + validate_config(plain(config)) + context = initialize( + enabled=True if config.distributed.enabled else None, + backend=config.distributed.backend, + device=config.trainer.device, + ) + try: + seed_everything(int(config.trainer.seed) + context.rank) + run_root = Path(config.environment.output_root) + if context.is_main_process: + save_resolved(config, run_root) + barrier(context) + split_loaders = loaders(config, context) + model, task = _model_and_task(config, split_loaders["train"]) + trainer = build_trainer( + plain(config.trainer), model=model, task=task, distributed_context=context + ) + if config.checkpoint.resume: + payload = CheckpointManager.load( + config.checkpoint.resume, map_location="cpu" + ) + restore_training_state( + payload, + model=trainer.model, + trainer=trainer, + optimizer=trainer.optimizer, + scheduler=trainer.scheduler, + early_stopping=trainer.early_stopping, + ) + validation_loader = ( + split_loaders["validation"] if len(split_loaders["validation"]) else None + ) + trainer.fit( + split_loaders["train"], + validation_loader, + epochs=int(config.trainer.max_epochs), + ) + if not context.is_main_process: + barrier(context) + return ( + Path(config.checkpoint.directory) + / f"epoch_{trainer.state.epoch:04d}.pt" + ) + manager = CheckpointManager(config.checkpoint.directory) + path = manager.save( + model=trainer.model, + trainer_state=trainer.state, + optimizer=trainer.optimizer, + scheduler=trainer.scheduler, + early_stopping=trainer.early_stopping, + task_config=plain(config.task), + model_config=plain(config.model), + ) + barrier(context) + return path + finally: + finalize(context) + + +def predict_or_evaluate(config: DictConfig, *, evaluate: bool = False) -> Any: + validate_config(plain(config)) + context = initialize( + enabled=True if config.distributed.enabled else None, + backend=config.distributed.backend, + device=config.trainer.device, + ) + try: + split_loaders = loaders(config, context) + model, task = _model_and_task(config, split_loaders[config.inference.split]) + payload = CheckpointManager.load( + config.inference.checkpoint, map_location="cpu" + ) + from gnn4colliders.training import load_model_weights + + load_model_weights(model, payload) + predictor = Predictor( + model, + task, + device=str(config.trainer.device), + distributed_context=context, + ) + result = ( + predictor.evaluate(split_loaders[config.inference.split]) + if evaluate + else predictor.predict(split_loaders[config.inference.split]) + ) + if not evaluate and context.is_main_process: + output = Path(config.inference.output) + if str(config.inference.format).lower() != "npz": + raise ValueError("only npz inference output is currently supported") + write_npz(result, output) + barrier(context) + return result + finally: + finalize(context) + + +def export_onnx(config: DictConfig) -> Path: + """Export a checkpoint using the first prepared graph batch.""" + if bool(config.distributed.enabled) or int(os.environ.get("WORLD_SIZE", "1")) > 1: + raise ValueError("ONNX export must run on a single process") + checkpoint = config.export.checkpoint + if not checkpoint: + raise ValueError("export.checkpoint is required") + if str(config.export.format).lower() != "onnx": + raise ValueError("only export.format=onnx is supported") + split_loaders = loaders(config) + split = str(config.export.split) + if split not in split_loaders: + raise ValueError(f"unsupported export split {split!r}") + example_batch = next(iter(split_loaders[split]), None) + if example_batch is None: + raise ValueError(f"export split {split!r} is empty") + from gnn4colliders.export import export_checkpoint_to_onnx + + return export_checkpoint_to_onnx( + checkpoint, + config.export.output, + example_batch=example_batch, + opset=int(config.export.opset), + overwrite=bool(config.export.overwrite), + ) diff --git a/src/gnn4colliders/config/factories.py b/src/gnn4colliders/config/factories.py new file mode 100644 index 0000000000000000000000000000000000000000..a678844453bc741015343185a3d61fb6d4b8c370 --- /dev/null +++ b/src/gnn4colliders/config/factories.py @@ -0,0 +1,145 @@ +"""Small allow-listed factories for configured application components.""" + +from __future__ import annotations + +from collections.abc import Mapping +from typing import Any + +import torch + +from gnn4colliders.distributed import DistributedContext +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork +from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask +from gnn4colliders.training import ( + EarlyStopping, + Trainer, + build_optimizer, + build_scheduler, +) + + +def _get(config: Mapping[str, Any], name: str, default: Any = None) -> Any: + value = config.get(name, default) + return value + + +def build_model( + config: Mapping[str, Any], + *, + sample_graph: Any, + sample_global: torch.Tensor | None = None, +) -> torch.nn.Module: + """Build one of the supported models from semantic configuration.""" + family = str(_get(config, "family", "root_gnn")).lower() + name = str(_get(config, "name", "")).lower() + if not name: + name = ( + "fine_tuned_edge_network" + if "fine" in str(_get(config, "class", "")).lower() + else "edge_network" + ) + if family != "root_gnn": + raise ValueError(f"unsupported model family {family!r}") + kwargs = { + "hid_size": int(_get(config, "hid_size", _get(config, "hidden_dim", 128))), + "out_size": int(_get(config, "out_size", 1)), + "n_layers": int(_get(config, "n_layers", 2)), + "n_proc_steps": int( + _get(config, "n_proc_steps", _get(config, "processing_steps", 4)) + ), + "dropout": float(_get(config, "dropout", 0.0)), + } + if name in {"edge_network", "edge"}: + return EdgeNetwork(sample_graph, sample_global, **kwargs) + if name in {"fine_tuned_edge_network", "finetuned_edge_network", "fine_tuned"}: + checkpoint = _get(config, "pretrained_checkpoint") + if not checkpoint: + raise ValueError( + "fine-tuned model requires model.pretrained_checkpoint or " + "checkpoint.pretrained" + ) + from gnn4colliders.training import CheckpointManager, load_model_weights + + payload = CheckpointManager.load(checkpoint, map_location="cpu") + base_config = dict(payload.get("model_config") or {}) + base_config.update( + {key: value for key, value in kwargs.items() if key != "out_size"} + ) + backbone = EdgeNetwork( + sample_graph, + sample_global, + hid_size=int(base_config.get("hid_size", kwargs["hid_size"])), + out_size=int(base_config.get("out_size", 1)), + n_layers=int(base_config.get("n_layers", kwargs["n_layers"])), + n_proc_steps=int(base_config.get("n_proc_steps", kwargs["n_proc_steps"])), + dropout=float(base_config.get("dropout", kwargs["dropout"])), + ) + load_model_weights(backbone, payload) + return FineTunedEdgeNetwork( + backbone, + kwargs["out_size"], + freeze_backbone=bool(_get(config, "freeze_backbone", False)), + ) + raise ValueError(f"unsupported root_gnn model {name!r}") + + +def build_task(config: Mapping[str, Any]) -> Any: + task_type = str(_get(config, "type", "multiclass_classification")).lower() + kwargs = { + "absolute_weights": bool( + _get( + config, "absolute_weights", _get(config, "use_absolute_weights", False) + ) + ) + } + if task_type in {"binary", "binary_classification"}: + return BinaryClassificationTask( + threshold=float(_get(config, "threshold", 0.5)), **kwargs + ) + if task_type in {"multiclass", "multiclass_classification"}: + return MulticlassClassificationTask(**kwargs) + raise ValueError(f"unsupported task type {task_type!r}") + + +def build_trainer( + config: Mapping[str, Any], + *, + model: torch.nn.Module, + task: Any, + distributed_context: DistributedContext | None = None, +) -> Trainer: + optimizer_config = _get(config, "optimizer", {}) or {} + optimizer = build_optimizer(model, **dict(optimizer_config)) + scheduler_config = _get(config, "scheduler", {}) or {} + scheduler = None + scheduler_step = "epoch" + if _get(scheduler_config, "name") not in (None, "", "none"): + scheduler_step = str(_get(scheduler_config, "step", "epoch")) + scheduler_kwargs = { + key: value + for key, value in dict(scheduler_config).items() + if key not in {"name", "step"} + } + scheduler = build_scheduler( + optimizer, name=str(_get(scheduler_config, "name")), **scheduler_kwargs + ) + early_config = _get(config, "early_stopping", {}) or {} + early = None + if bool(_get(early_config, "enabled", False)): + early = EarlyStopping( + **{ + key: value + for key, value in dict(early_config).items() + if key != "enabled" + } + ) + return Trainer( + model, + task, + optimizer, + scheduler, + device=str(_get(config, "device", "cpu")), + early_stopping=early, + scheduler_step=scheduler_step, + distributed_context=distributed_context, + ) diff --git a/src/gnn4colliders/config/validation.py b/src/gnn4colliders/config/validation.py new file mode 100644 index 0000000000000000000000000000000000000000..1a68845211b04b5cdd0b9d2d453b6adaf7131bdc --- /dev/null +++ b/src/gnn4colliders/config/validation.py @@ -0,0 +1,41 @@ +"""Early validation for cross-group configuration relationships.""" + +from __future__ import annotations + +from collections.abc import Mapping + + +def validate_config(config: Mapping[str, object]) -> None: + data = config.get("data", {}) or {} + trainer = config.get("trainer", {}) or {} + model = config.get("model", {}) or {} + task = config.get("task", {}) or {} + if int(data.get("batch_size", 1)) < 1: + raise ValueError("data.batch_size must be positive") + if int(trainer.get("max_epochs", 1)) < 1: + raise ValueError("trainer.max_epochs must be positive") + task_type = str(task.get("type", "")).lower() + out_size = int(model.get("out_size", 1)) + if "binary" in task_type and out_size != 1: + raise ValueError("binary classification requires model.out_size=1") + if "multi" in task_type and out_size != int(task.get("num_classes", out_size)): + raise ValueError("multiclass model.out_size must equal task.num_classes") + if str(model.get("name", "")).lower() in {"fine_tuned_edge_network", "fine_tuned"}: + checkpoint = model.get("pretrained_checkpoint") or ( + config.get("checkpoint", {}) or {} + ).get("pretrained") + if not checkpoint: + raise ValueError("fine-tuned models require checkpoint.pretrained") + checkpoint = config.get("checkpoint", {}) or {} + if checkpoint.get("resume") and checkpoint.get("pretrained"): + raise ValueError( + "checkpoint.resume and checkpoint.pretrained are distinct " + "workflows and cannot both be set" + ) + splits = data.get("splits", {}) or {} + groups = [ + set(splits.get(name, [])) + for name in ("train_folds", "validation_folds", "test_folds") + ] + if any(groups[i] & groups[j] for i in range(3) for j in range(i + 1, 3)): + raise ValueError("train, validation, and test folds must be disjoint") diff --git a/src/gnn4colliders/configs/__init__.py b/src/gnn4colliders/configs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..f66eabc4004d36e8c29edf9e73efc5798ef64fe8 --- /dev/null +++ b/src/gnn4colliders/configs/__init__.py @@ -0,0 +1 @@ +"""Canonical Hydra configuration tree shipped with the package.""" diff --git a/src/gnn4colliders/configs/checkpoint/default.yaml b/src/gnn4colliders/configs/checkpoint/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..55a2d008b0e6292b558e97fcdd54c9045382b785 --- /dev/null +++ b/src/gnn4colliders/configs/checkpoint/default.yaml @@ -0,0 +1,6 @@ +directory: ${environment.output_root}/checkpoints +resume: null +pretrained: null +save_every_epochs: 1 +save_best: false +save_last: true diff --git a/src/gnn4colliders/configs/config.yaml b/src/gnn4colliders/configs/config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..326f81714621768885e9ca2ef7082c67469b31be --- /dev/null +++ b/src/gnn4colliders/configs/config.yaml @@ -0,0 +1,20 @@ +defaults: + - data: local_test + - model: root_gnn/edge_network + - task: pretraining_multiclass + - trainer: default + - checkpoint: default + - inference: default + - environment: local + - distributed: single + - export: onnx + - _self_ +experiment: + name: pretraining_multiclass +logging: + level: INFO +hydra: + run: + dir: . + job: + chdir: false diff --git a/src/gnn4colliders/configs/data/delphes.yaml b/src/gnn4colliders/configs/data/delphes.yaml new file mode 100644 index 0000000000000000000000000000000000000000..af3ffd807cdb77dc97071ead23f4ae7fd34e606e --- /dev/null +++ b/src/gnn4colliders/configs/data/delphes.yaml @@ -0,0 +1,4 @@ +defaults: + - /data/local_test + - _self_ +tree_name: Events diff --git a/src/gnn4colliders/configs/data/local_test.yaml b/src/gnn4colliders/configs/data/local_test.yaml new file mode 100644 index 0000000000000000000000000000000000000000..0e22ba4ad8ea7ca333b55f8d3fd112d8a6e1d35f --- /dev/null +++ b/src/gnn4colliders/configs/data/local_test.yaml @@ -0,0 +1,18 @@ +files: [] +tree_name: Events +feature_branches: null +object_types: null +scales: null +global_features: [] +fold_var: eventNumber +weight_var: null +batch_size: 32 +num_workers: 0 +shuffle: true +seed: 42 +cache: + path: null +splits: + train_folds: [0, 1, 2] + validation_folds: [3] + test_folds: [4] diff --git a/src/gnn4colliders/configs/distributed/ddp.yaml b/src/gnn4colliders/configs/distributed/ddp.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6db7df4c3c51218b786804e3078da6dcdba8a946 --- /dev/null +++ b/src/gnn4colliders/configs/distributed/ddp.yaml @@ -0,0 +1,2 @@ +enabled: true +backend: null diff --git a/src/gnn4colliders/configs/distributed/single.yaml b/src/gnn4colliders/configs/distributed/single.yaml new file mode 100644 index 0000000000000000000000000000000000000000..05857ba25a60aadd28d5ab39456541612840160b --- /dev/null +++ b/src/gnn4colliders/configs/distributed/single.yaml @@ -0,0 +1,2 @@ +enabled: false +backend: null diff --git a/src/gnn4colliders/configs/environment/local.yaml b/src/gnn4colliders/configs/environment/local.yaml new file mode 100644 index 0000000000000000000000000000000000000000..405ecdf7f2baf7aea70418529d6f2450d6895e91 --- /dev/null +++ b/src/gnn4colliders/configs/environment/local.yaml @@ -0,0 +1,3 @@ +device: cpu +data_root: null +output_root: outputs/${experiment.name} diff --git a/src/gnn4colliders/configs/environment/perlmutter.yaml b/src/gnn4colliders/configs/environment/perlmutter.yaml new file mode 100644 index 0000000000000000000000000000000000000000..2b7df3bc957fc88212ea6d3177575b2a98a8290c --- /dev/null +++ b/src/gnn4colliders/configs/environment/perlmutter.yaml @@ -0,0 +1,3 @@ +device: cuda +data_root: null +output_root: outputs/${experiment.name} diff --git a/src/gnn4colliders/configs/export/onnx.yaml b/src/gnn4colliders/configs/export/onnx.yaml new file mode 100644 index 0000000000000000000000000000000000000000..bd727c7c5611cc170ede2b12db784b76dc62c631 --- /dev/null +++ b/src/gnn4colliders/configs/export/onnx.yaml @@ -0,0 +1,6 @@ +format: onnx +checkpoint: null +output: ${environment.output_root}/export/model.onnx +opset: 17 +split: train +overwrite: false diff --git a/src/gnn4colliders/configs/inference/default.yaml b/src/gnn4colliders/configs/inference/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..7c9f60d184618b625f8c140efaabc42ff1dfbbb7 --- /dev/null +++ b/src/gnn4colliders/configs/inference/default.yaml @@ -0,0 +1,4 @@ +checkpoint: null +output: ${environment.output_root}/predictions/predictions.npz +format: npz +split: test diff --git a/src/gnn4colliders/configs/model/root_gnn/edge_network.yaml b/src/gnn4colliders/configs/model/root_gnn/edge_network.yaml new file mode 100644 index 0000000000000000000000000000000000000000..17b359d2241a9db07d45b58e8027383d28f0c5d3 --- /dev/null +++ b/src/gnn4colliders/configs/model/root_gnn/edge_network.yaml @@ -0,0 +1,7 @@ +family: root_gnn +name: edge_network +hid_size: 128 +out_size: 12 +n_layers: 2 +n_proc_steps: 4 +dropout: 0.1 diff --git a/src/gnn4colliders/configs/model/root_gnn/fine_tuned_edge_network.yaml b/src/gnn4colliders/configs/model/root_gnn/fine_tuned_edge_network.yaml new file mode 100644 index 0000000000000000000000000000000000000000..3962b6cc1f15a1c6dc74b8d90924e9af3c250bc1 --- /dev/null +++ b/src/gnn4colliders/configs/model/root_gnn/fine_tuned_edge_network.yaml @@ -0,0 +1,9 @@ +family: root_gnn +name: fine_tuned_edge_network +hid_size: 128 +out_size: 1 +n_layers: 2 +n_proc_steps: 4 +dropout: 0.1 +freeze_backbone: true +pretrained_checkpoint: null diff --git a/src/gnn4colliders/configs/task/binary_classification.yaml b/src/gnn4colliders/configs/task/binary_classification.yaml new file mode 100644 index 0000000000000000000000000000000000000000..c6fe18647d796a49c4fea5435aa56a2597f5b1c4 --- /dev/null +++ b/src/gnn4colliders/configs/task/binary_classification.yaml @@ -0,0 +1,3 @@ +type: binary_classification +threshold: 0.5 +absolute_weights: false diff --git a/src/gnn4colliders/configs/task/pretraining_multiclass.yaml b/src/gnn4colliders/configs/task/pretraining_multiclass.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fa7455b782aaf919454f48724f1a2844c2740227 --- /dev/null +++ b/src/gnn4colliders/configs/task/pretraining_multiclass.yaml @@ -0,0 +1,3 @@ +type: multiclass_classification +num_classes: 12 +absolute_weights: false diff --git a/src/gnn4colliders/configs/task/tth_cp_finetune.yaml b/src/gnn4colliders/configs/task/tth_cp_finetune.yaml new file mode 100644 index 0000000000000000000000000000000000000000..4e91d39020e9de5bfcebeac045de1364888b27f8 --- /dev/null +++ b/src/gnn4colliders/configs/task/tth_cp_finetune.yaml @@ -0,0 +1,3 @@ +defaults: + - binary_classification + - _self_ diff --git a/src/gnn4colliders/configs/trainer/debug.yaml b/src/gnn4colliders/configs/trainer/debug.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fc2ad046b0c3aeeabb34402e50495cee91c8fbe2 --- /dev/null +++ b/src/gnn4colliders/configs/trainer/debug.yaml @@ -0,0 +1,4 @@ +defaults: + - default + - _self_ +max_epochs: 1 diff --git a/src/gnn4colliders/configs/trainer/default.yaml b/src/gnn4colliders/configs/trainer/default.yaml new file mode 100644 index 0000000000000000000000000000000000000000..fe5cc5f08e7bafb8e4f9301d69327e9ebf523072 --- /dev/null +++ b/src/gnn4colliders/configs/trainer/default.yaml @@ -0,0 +1,15 @@ +max_epochs: 100 +device: ${environment.device} +seed: 42 +optimizer: + name: adam + learning_rate: 0.001 + weight_decay: 0.0 +scheduler: + name: null +early_stopping: + enabled: false + monitor: loss + mode: min + patience: 10 + min_delta: 1.0e-8 diff --git a/src/gnn4colliders/data/__init__.py b/src/gnn4colliders/data/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..65f48774e57698388ee3f17a25421702a1b5ea5e --- /dev/null +++ b/src/gnn4colliders/data/__init__.py @@ -0,0 +1,46 @@ +"""Generic data access and dataset infrastructure.""" + +from .cache import CacheMetadata, GraphSampleCache +from .dataset import RootEventDataset +from .graph_dataset import ( + DistributedGraphDataLoader, + GraphBatch, + GraphDataLoader, + GraphDataset, + GraphSample, + batch_graph_samples, +) +from .metadata import ( + CACHE_SCHEMA_VERSION, + FEATURE_SCHEMA_VERSION, + GRAPH_SCHEMA_VERSION, + BatchMetadata, + EventMetadata, +) +from .root_io import branch_names_from_specs, read_tree, tree_num_entries +from .sample import EventSample +from .splits import SplitDefinition, select_folds, select_split + +__all__ = [ + "EventSample", + "EventMetadata", + "BatchMetadata", + "GraphSample", + "GraphBatch", + "GraphDataset", + "GraphDataLoader", + "DistributedGraphDataLoader", + "batch_graph_samples", + "CacheMetadata", + "GraphSampleCache", + "SplitDefinition", + "select_folds", + "select_split", + "FEATURE_SCHEMA_VERSION", + "GRAPH_SCHEMA_VERSION", + "CACHE_SCHEMA_VERSION", + "RootEventDataset", + "branch_names_from_specs", + "read_tree", + "tree_num_entries", +] diff --git a/src/gnn4colliders/data/cache.py b/src/gnn4colliders/data/cache.py new file mode 100644 index 0000000000000000000000000000000000000000..98d266dd18c208e250aa14adf8351aca585e5aed --- /dev/null +++ b/src/gnn4colliders/data/cache.py @@ -0,0 +1,69 @@ +"""Versioned cache for processed graph samples.""" + +from __future__ import annotations + +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Any + +import torch + +from .graph_dataset import GraphSample +from .metadata import CACHE_SCHEMA_VERSION, FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION + + +@dataclass(frozen=True) +class CacheMetadata: + feature_schema_version: int = FEATURE_SCHEMA_VERSION + graph_schema_version: int = GRAPH_SCHEMA_VERSION + cache_schema_version: int = CACHE_SCHEMA_VERSION + source: str | None = None + preprocessing: dict[str, Any] = field(default_factory=dict) + + +class GraphSampleCache: + """Save/load graph samples with schema metadata checked before use.""" + + def __init__(self, path: str | Path, metadata: CacheMetadata | None = None) -> None: + self.path = Path(path) + self.metadata = metadata or CacheMetadata() + + def save(self, samples: list[GraphSample] | tuple[GraphSample, ...]) -> None: + payload = { + "cache_metadata": asdict(self.metadata), + "samples": tuple(samples), + } + self.path.parent.mkdir(parents=True, exist_ok=True) + torch.save(payload, self.path) + + def load(self) -> tuple[GraphSample, ...]: + payload = torch.load(self.path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict): + raise ValueError("cache payload must be a mapping") + actual = payload.get("cache_metadata", {}) + if not isinstance(actual, dict): + raise ValueError("cache metadata is missing or malformed") + expected = asdict(self.metadata) + for key in ( + "feature_schema_version", + "graph_schema_version", + "cache_schema_version", + ): + if actual.get(key) != expected[key]: + raise ValueError( + f"cache schema mismatch for {key}: " + f"found {actual.get(key)!r}, expected {expected[key]!r}" + ) + if actual.get("preprocessing", {}) != expected["preprocessing"]: + raise ValueError("cache preprocessing fingerprint mismatch") + if "samples" not in payload: + raise ValueError("cache samples are missing") + samples = payload["samples"] + if not isinstance(samples, (list, tuple)): + raise ValueError("cache samples are malformed") + if any(not isinstance(sample, GraphSample) for sample in samples): + raise ValueError("cache contains invalid graph samples") + return tuple(samples) + + def exists(self) -> bool: + return self.path.is_file() diff --git a/src/gnn4colliders/data/dataset.py b/src/gnn4colliders/data/dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..1d9da8c32a9ed2bbc2a38064ccac70670865007a --- /dev/null +++ b/src/gnn4colliders/data/dataset.py @@ -0,0 +1,147 @@ +"""Dataset abstraction over one or more ROOT files.""" + +from __future__ import annotations + +from collections.abc import Sequence +from pathlib import Path +from typing import Any + +import numpy as np + +from .metadata import EventMetadata +from .root_io import branch_names_from_specs, read_tree, tree_num_entries +from .sample import EventSample + + +def _scalar(value: Any) -> Any: + """Convert an Awkward/NumPy scalar while retaining jagged values.""" + + array = np.asarray(value) + return array.item() if array.ndim == 0 else value + + +def _as_float(value: Any) -> float: + array = np.asarray(value) + if array.ndim != 0: + raise ValueError(f"tracking/global value is not scalar: {value!r}") + return float(array) + + +def _file_labels(label: Any, count: int) -> list[Any]: + if isinstance(label, Sequence) and not isinstance(label, (str, bytes)): + if len(label) != count: + raise ValueError("one label is required for each input ROOT file") + return list(label) + return [label] * count + + +class RootEventDataset: + """Read selected ROOT branches and expose deterministic event indexing. + + Fold and weight are read as named metadata. Files and events are + concatenated in the order supplied by the caller. + """ + + def __init__( + self, + files: str | Path | Sequence[str | Path], + *, + tree_name: str = "nominal_Loose", + label: Any = 1, + feature_branches: Sequence[Any] = (), + tracking_info: Sequence[Any] = (), + global_features: Sequence[Any] = (), + fold_var: str = "eventNumber", + weight_var: str | None = None, + ) -> None: + if isinstance(files, (str, Path)): + files = [files] + self.files = tuple(Path(path) for path in files) + if not self.files: + raise ValueError("at least one ROOT file is required") + self.tree_name = tree_name + self.labels = tuple(_file_labels(label, len(self.files))) + self.feature_branches = feature_branches + self.global_specs = tuple(global_features) + # ``tracking_info`` is accepted for old callers, but positional + # tracking is not part of EventSample. The historical first two + # entries are represented by the named fold/weight fields below. + del tracking_info + self.metadata_specs = (fold_var, 1 if weight_var is None else weight_var) + self._branches = branch_names_from_specs( + feature_branches, + metadata_specs=self.metadata_specs, + global_features=self.global_specs, + ) + for file_label in self.labels: + if isinstance(file_label, str) and file_label not in self._branches: + self._branches = (*self._branches, file_label) + self._offsets: tuple[int, ...] = self._make_offsets() + + def _make_offsets(self) -> tuple[int, ...]: + offsets = [0] + for path in self.files: + offsets.append(offsets[-1] + tree_num_entries(path, self.tree_name)) + return tuple(offsets) + + @property + def branches(self) -> tuple[str, ...]: + """The exact branch set requested from Uproot.""" + + return self._branches + + def __len__(self) -> int: + return self._offsets[-1] + + def _locate(self, index: int) -> tuple[int, int]: + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(f"event index {index} is out of range") + file_index = int(np.searchsorted(self._offsets, index, side="right") - 1) + return file_index, index - self._offsets[file_index] + + def __getitem__(self, index: int) -> EventSample: + if not isinstance(index, (int, np.integer)): + raise TypeError("event index must be an integer") + global_index = int(index) + if global_index < 0: + global_index += len(self) + file_index, local_index = self._locate(global_index) + arrays = read_tree(self.files[file_index], self.tree_name, self._branches) + event = {branch: arrays[branch][local_index] for branch in self._branches} + metadata_values = np.asarray( + [ + _as_float(event[value]) if isinstance(value, str) else float(value) + for value in self.metadata_specs + ], + dtype=np.float32, + ) + globals_ = np.asarray( + [ + _as_float(event[value]) if isinstance(value, str) else float(value) + for value in self.global_specs + ], + dtype=np.float32, + ) + label = ( + event[self.labels[file_index]] + if isinstance(self.labels[file_index], str) + else self.labels[file_index] + ) + return EventSample( + objects={branch: _scalar(value) for branch, value in event.items()}, + label=_scalar(label), + global_features=globals_, + event_index=global_index, + metadata=EventMetadata( + fold=int(metadata_values[0]), + weight=float(metadata_values[1]), + sample_id=f"{self.files[file_index]}:{self.tree_name}:{local_index}", + extra={ + "source_file": str(self.files[file_index]), + "tree_name": self.tree_name, + "entry_index": local_index, + }, + ), + ) diff --git a/src/gnn4colliders/data/graph_dataset.py b/src/gnn4colliders/data/graph_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..fba36308be726a22c99c2eb9446cc6b782e93880 --- /dev/null +++ b/src/gnn4colliders/data/graph_dataset.py @@ -0,0 +1,202 @@ +"""Graph samples and deterministic batching for the ROOT-GNN path.""" + +from __future__ import annotations + +from collections.abc import Iterator, Sequence +from dataclasses import dataclass +from typing import Any + +import torch + +from .metadata import BatchMetadata, EventMetadata + + +@dataclass(frozen=True) +class GraphSample: + graph: Any + label: torch.Tensor + global_features: torch.Tensor | None + metadata: EventMetadata + + +@dataclass(frozen=True) +class GraphBatch: + graph: Any + labels: torch.Tensor + global_features: torch.Tensor | None + metadata: BatchMetadata + + def to(self, device: torch.device | str) -> "GraphBatch": + """Return a copy with tensor and graph data placed on ``device``. + + Sample identifiers and other Python metadata intentionally remain + ordinary CPU-side values. + """ + + graph = self.graph.to(device) if hasattr(self.graph, "to") else self.graph + metadata = BatchMetadata( + fold=self.metadata.fold.to(device), + weight=self.metadata.weight.to(device), + sample_id=self.metadata.sample_id, + extra=self.metadata.extra, + ) + return GraphBatch( + graph=graph, + labels=self.labels.to(device), + global_features=( + None + if self.global_features is None + else self.global_features.to(device) + ), + metadata=metadata, + ) + + +def batch_graph_samples(samples: Sequence[GraphSample]) -> GraphBatch: + """Batch samples in the supplied order using DGL's native batching.""" + + if not samples: + raise ValueError("cannot batch an empty sequence") + try: + import dgl + except ImportError as error: # pragma: no cover - optional dependency + raise ImportError( + "batch_graph_samples requires the 'root-gnn' extra" + ) from error + graphs = dgl.batch([sample.graph for sample in samples]) + labels = torch.stack([torch.as_tensor(sample.label) for sample in samples]) + globals_ = [sample.global_features for sample in samples] + if any(value is None for value in globals_): + if not all(value is None for value in globals_): + raise ValueError("global_features must be present for every sample or none") + batched_globals = None + else: + batched_globals = torch.stack( + [value for value in globals_ if value is not None] + ) + return GraphBatch( + graph=graphs, + labels=labels, + global_features=batched_globals, + metadata=BatchMetadata.from_events([sample.metadata for sample in samples]), + ) + + +class GraphDataLoader: + """Deterministic, dependency-light loader over :class:`GraphDataset`. + + Shuffling uses a local Torch generator, so constructing or iterating a + loader does not mutate process-wide random state. + """ + + def __init__( + self, + dataset: Sequence[GraphSample], + batch_size: int, + *, + shuffle: bool = False, + drop_last: bool = False, + seed: int = 12345, + ) -> None: + if batch_size < 1: + raise ValueError("batch_size must be positive") + self.dataset = dataset + self.batch_size = batch_size + self.shuffle = shuffle + self.drop_last = drop_last + self.seed = seed + self.epoch = 0 + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def __len__(self) -> int: + count, remainder = divmod(len(self.dataset), self.batch_size) + return count if self.drop_last or remainder == 0 else count + 1 + + def __iter__(self) -> Iterator[GraphBatch]: + indices = torch.arange(len(self.dataset), dtype=torch.long) + if self.shuffle: + generator = torch.Generator() + generator.manual_seed(self.seed + self.epoch) + indices = indices[torch.randperm(len(indices), generator=generator)] + limit = len(indices) - (len(indices) % self.batch_size if self.drop_last else 0) + for start in range(0, limit, self.batch_size): + selected = [ + self.dataset[int(index)] + for index in indices[start : start + self.batch_size] + ] + yield batch_graph_samples(selected) + + +class DistributedGraphDataLoader(GraphDataLoader): + """Shard individual graph samples before batching. + + Training sharding pads to equal rank lengths; evaluation sharding does not + pad, preventing duplicate events from biasing metrics or predictions. + ``batch_size`` is a per-process value. + """ + + def __init__( + self, + dataset, + batch_size, + context, + *, + shuffle=False, + drop_last=False, + seed=12345, + ): + super().__init__( + dataset, batch_size, shuffle=shuffle, drop_last=drop_last, seed=seed + ) + from gnn4colliders.distributed import DistributedIndices + + self.distributed_indices = DistributedIndices( + len(dataset), + context, + shuffle=shuffle, + drop_last=drop_last, + seed=seed, + pad=shuffle, + ) + + def set_epoch(self, epoch: int) -> None: + super().set_epoch(epoch) + self.distributed_indices.set_epoch(epoch) + + def __len__(self) -> int: + count = len(self.distributed_indices.indices()) + return ( + count // self.batch_size + if self.drop_last + else (count + self.batch_size - 1) // self.batch_size + ) + + def __iter__(self) -> Iterator[GraphBatch]: + indices = self.distributed_indices.indices() + limit = len(indices) - (len(indices) % self.batch_size if self.drop_last else 0) + for start in range(0, limit, self.batch_size): + yield batch_graph_samples( + [ + self.dataset[index] + for index in indices[start : start + self.batch_size] + ] + ) + + +class GraphDataset(Sequence[GraphSample]): + """Small immutable-semantics dataset over already processed graph samples.""" + + def __init__(self, samples: Sequence[GraphSample]) -> None: + self._samples = tuple(samples) + + def __len__(self) -> int: + return len(self._samples) + + def __getitem__(self, index: int) -> GraphSample: + return self._samples[index] + + @property + def samples(self) -> tuple[GraphSample, ...]: + return self._samples diff --git a/src/gnn4colliders/data/metadata.py b/src/gnn4colliders/data/metadata.py new file mode 100644 index 0000000000000000000000000000000000000000..2eb813170b5f797a54903d67a1673feb9df45556 --- /dev/null +++ b/src/gnn4colliders/data/metadata.py @@ -0,0 +1,61 @@ +"""Named metadata and schema versions for processed collider data.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +import torch + +FEATURE_SCHEMA_VERSION = 1 +GRAPH_SCHEMA_VERSION = 1 +CACHE_SCHEMA_VERSION = 1 + + +@dataclass(frozen=True) +class EventMetadata: + """Metadata belonging to one event, independent of its representation.""" + + fold: int + weight: float + sample_id: str + extra: Mapping[str, Any] = field(default_factory=dict) + + @classmethod + def from_legacy_tracking( + cls, + tracking_row: Sequence[Any], + *, + sample_id: str, + extra: Mapping[str, Any] | None = None, + ) -> "EventMetadata": + """Use the explicit compatibility adapter for old tracking rows.""" + from gnn4colliders.compat.metadata import event_metadata_from_legacy_tracking + + return event_metadata_from_legacy_tracking( + tracking_row, sample_id=sample_id, extra=extra + ) + + +@dataclass(frozen=True) +class BatchMetadata: + """Columnar metadata for a deterministic ordered batch of events.""" + + fold: torch.Tensor + weight: torch.Tensor + sample_id: tuple[str, ...] + extra: Mapping[str, tuple[Any, ...]] = field(default_factory=dict) + + @classmethod + def from_events(cls, events: Sequence[EventMetadata]) -> "BatchMetadata": + keys = sorted({key for event in events for key in event.extra}) + extra = {key: tuple(event.extra.get(key) for event in events) for key in keys} + return cls( + fold=torch.tensor([event.fold for event in events], dtype=torch.long), + weight=torch.tensor( + [event.weight for event in events], dtype=torch.float32 + ), + sample_id=tuple(event.sample_id for event in events), + extra=extra, + ) diff --git a/src/gnn4colliders/data/root_io.py b/src/gnn4colliders/data/root_io.py new file mode 100644 index 0000000000000000000000000000000000000000..2d92788487f53390194a0bf921af7cbe7e4bf91a --- /dev/null +++ b/src/gnn4colliders/data/root_io.py @@ -0,0 +1,72 @@ +"""Small Uproot boundary for reading the active ROOT event data.""" + +from __future__ import annotations + +from collections.abc import Iterable, Sequence +from pathlib import Path +from typing import Any + +import awkward as ak +import uproot + + +def validate_branches(available: Iterable[str], requested: Sequence[str]) -> None: + """Raise a useful error when requested branches are absent.""" + + available_set = set(available) + missing = sorted(set(requested).difference(available_set)) + if missing: + raise KeyError(f"tree is missing requested branches: {missing}") + + +def read_tree( + path: str | Path, + tree_name: str, + branches: Sequence[str], +) -> ak.Array: + """Read selected branches from one ROOT tree as an Awkward record array.""" + + requested = tuple(dict.fromkeys(branches)) + with uproot.open(path) as root_file: + if tree_name not in root_file: + raise KeyError(f"tree {tree_name!r} was not found in {path}") + tree = root_file[tree_name] + validate_branches(tree.keys(), requested) + return tree.arrays(list(requested), library="ak") + + +def tree_num_entries(path: str | Path, tree_name: str) -> int: + """Return the number of entries without reading branch data.""" + + with uproot.open(path) as root_file: + if tree_name not in root_file: + raise KeyError(f"tree {tree_name!r} was not found in {path}") + return int(root_file[tree_name].num_entries) + + +def branch_names_from_specs( + feature_branches: Sequence[Any] = (), + *, + metadata_specs: Sequence[Any] = (), + tracking_info: Sequence[Any] = (), + global_features: Sequence[Any] = (), +) -> tuple[str, ...]: + """Collect real branch names from feature, tracking, and global specs.""" + + names: list[str] = [] + + def visit(value: Any) -> None: + if isinstance(value, str): + if value not in {"CALC_E", "NODE_TYPE"} and value not in names: + names.append(value) + elif isinstance(value, (list, tuple)): + for item in value: + visit(item) + + visit(feature_branches) + visit(metadata_specs) + # Compatibility alias for old preparation callers; metadata is named in + # the production API and this value is not retained in samples. + visit(tracking_info) + visit(global_features) + return tuple(names) diff --git a/src/gnn4colliders/data/sample.py b/src/gnn4colliders/data/sample.py new file mode 100644 index 0000000000000000000000000000000000000000..bce15bbaf82d9c9b7ad3ebeb839d19acdc7e3b91 --- /dev/null +++ b/src/gnn4colliders/data/sample.py @@ -0,0 +1,54 @@ +"""Architecture-neutral representation of one collider event.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass +from typing import Any + +import numpy as np + +from .metadata import EventMetadata + + +@dataclass(frozen=True) +class EventSample: + """One event before conversion into a model-specific representation. + + ``objects`` is a read-only mapping of selected ROOT branch names to the + event's scalar or jagged values. Keeping branch values here lets the + shared feature layer consume the sample without introducing DGL or model + classes into the data package. + """ + + objects: Mapping[str, Any] + label: Any + global_features: np.ndarray + event_index: int + metadata: EventMetadata | None = None + + @property + def branches(self) -> Mapping[str, Any]: + """Alias that makes the raw branch-level representation explicit.""" + + return self.objects + + @property + def event_metadata(self) -> EventMetadata: + """Return the required named metadata for this event.""" + if self.metadata is None: + raise ValueError("sample has no named metadata") + return self.metadata + + @property + def tracking(self): + """Compatibility view for callers still expecting a tracking row. + + New code must use :attr:`event_metadata`; this view is deliberately + created only when an old caller asks for it and is never propagated by + the new data pipeline. + """ + import numpy as np + + metadata = self.event_metadata + return np.asarray([metadata.fold, metadata.weight], dtype=np.float32) diff --git a/src/gnn4colliders/data/splits.py b/src/gnn4colliders/data/splits.py new file mode 100644 index 0000000000000000000000000000000000000000..0336ed7138cf3d27d7763448ae484c5c1f688992 --- /dev/null +++ b/src/gnn4colliders/data/splits.py @@ -0,0 +1,49 @@ +"""Representation-independent fold selection and split definitions.""" + +from __future__ import annotations + +from collections.abc import Collection, Sequence +from dataclasses import dataclass +from typing import TypeVar + +T = TypeVar("T") + + +@dataclass(frozen=True) +class SplitDefinition: + train_folds: frozenset[int] = frozenset() + validation_folds: frozenset[int] = frozenset() + test_folds: frozenset[int] = frozenset() + + def __post_init__(self) -> None: + groups = (self.train_folds, self.validation_folds, self.test_folds) + if any(fold < 0 for group in groups for fold in group): + raise ValueError("fold identifiers must be non-negative") + for left, right in ( + (groups[0], groups[1]), + (groups[0], groups[2]), + (groups[1], groups[2]), + ): + if left & right: + raise ValueError("train, validation, and test folds must be disjoint") + + +def select_folds(samples: Sequence[T], folds: Collection[int]) -> list[T]: + """Return samples in input order whose named fold is selected.""" + + selected = frozenset(folds) + return [sample for sample in samples if sample.metadata.fold in selected] + + +def select_split(samples: Sequence[T], split: SplitDefinition, name: str) -> list[T]: + """Select one named split without changing the source sequence.""" + + try: + folds = { + "train": split.train_folds, + "validation": split.validation_folds, + "test": split.test_folds, + }[name] + except KeyError as error: + raise ValueError(f"unknown split {name!r}") from error + return select_folds(samples, folds) diff --git a/src/gnn4colliders/distributed/__init__.py b/src/gnn4colliders/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..92de02fc5f50980e7840bc7fd5c7a920cce1d8aa --- /dev/null +++ b/src/gnn4colliders/distributed/__init__.py @@ -0,0 +1,19 @@ +"""Minimal PyTorch DDP support for GNN4Colliders applications.""" + +from .collectives import barrier, broadcast_bool, gather_objects, gather_tensor +from .context import DistributedContext, finalize, initialize, process_group +from .model import prepare_model +from .sampler import DistributedIndices + +__all__ = [ + "DistributedContext", + "DistributedIndices", + "barrier", + "broadcast_bool", + "finalize", + "gather_objects", + "gather_tensor", + "initialize", + "prepare_model", + "process_group", +] diff --git a/src/gnn4colliders/distributed/collectives.py b/src/gnn4colliders/distributed/collectives.py new file mode 100644 index 0000000000000000000000000000000000000000..6b862b3e2246db1002855f151d37fb3315b5026e --- /dev/null +++ b/src/gnn4colliders/distributed/collectives.py @@ -0,0 +1,52 @@ +"""Small tensor/object collectives used by training and inference.""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.distributed as dist + +from .context import DistributedContext + + +def broadcast_bool(value: bool, context: DistributedContext, *, src: int = 0) -> bool: + if not context.enabled: + return value + tensor = torch.tensor([int(value)], device=context.device) + dist.broadcast(tensor, src=src) + return bool(tensor.item()) + + +def gather_tensor(tensor: torch.Tensor, context: DistributedContext) -> torch.Tensor: + """All-gather a first-dimension-sharded tensor, including unequal shards.""" + + if not context.enabled: + return tensor + local = tensor.contiguous() + size = torch.tensor([local.shape[0]], device=local.device, dtype=torch.long) + sizes = [torch.zeros_like(size) for _ in range(context.world_size)] + dist.all_gather(sizes, size) + max_size = max(int(item.item()) for item in sizes) + padded = torch.zeros( + (max_size, *local.shape[1:]), dtype=local.dtype, device=local.device + ) + padded[: local.shape[0]] = local + gathered = [torch.empty_like(padded) for _ in range(context.world_size)] + dist.all_gather(gathered, padded) + return torch.cat( + [part[: int(length.item())] for part, length in zip(gathered, sizes)], dim=0 + ) + + +def gather_objects(value: Any, context: DistributedContext) -> list[Any]: + if not context.enabled: + return [value] + gathered: list[Any] = [None] * context.world_size + dist.all_gather_object(gathered, value) + return gathered + + +def barrier(context: DistributedContext) -> None: + if context.enabled: + dist.barrier() diff --git a/src/gnn4colliders/distributed/context.py b/src/gnn4colliders/distributed/context.py new file mode 100644 index 0000000000000000000000000000000000000000..c57ad34df20b4d66e2fb38319793a4c79386a310 --- /dev/null +++ b/src/gnn4colliders/distributed/context.py @@ -0,0 +1,103 @@ +"""Process-group context and rank-local device selection.""" + +from __future__ import annotations + +import os +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import timedelta +from typing import Iterator + +import torch +import torch.distributed as dist + + +@dataclass(frozen=True) +class DistributedContext: + enabled: bool = False + rank: int = 0 + local_rank: int = 0 + world_size: int = 1 + device: torch.device = torch.device("cpu") + + @property + def is_main_process(self) -> bool: + return self.rank == 0 + + @classmethod + def single_process(cls, device: torch.device | str = "cpu") -> "DistributedContext": + return cls(device=torch.device(device)) + + +def _launched() -> bool: + return "RANK" in os.environ or "WORLD_SIZE" in os.environ + + +def initialize( + *, + enabled: bool | None = None, + backend: str | None = None, + device: torch.device | str | None = None, + timeout_seconds: int = 1800, +) -> DistributedContext: + """Initialize ``torchrun``'s ``env://`` process group when requested.""" + + launched = _launched() + if enabled is False and launched and int(os.environ.get("WORLD_SIZE", "1")) > 1: + raise ValueError( + "distributed launcher variables are present but distributed.enabled=false" + ) + active = launched if enabled is None else enabled + if not active: + requested = torch.device( + device or ("cuda" if torch.cuda.is_available() else "cpu") + ) + if requested.type == "cuda" and not torch.cuda.is_available(): + requested = torch.device("cpu") + return DistributedContext.single_process(requested) + + try: + rank = int(os.environ["RANK"]) + local_rank = int(os.environ.get("LOCAL_RANK", rank)) + world_size = int(os.environ["WORLD_SIZE"]) + except KeyError as error: + raise ValueError( + "distributed execution requires torchrun RANK and WORLD_SIZE" + ) from error + if world_size < 2: + return DistributedContext.single_process(device or "cpu") + requested = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu")) + if requested.type == "cuda": + if not torch.cuda.is_available(): + raise RuntimeError( + "CUDA was requested for distributed execution but is unavailable" + ) + torch.cuda.set_device(local_rank) + assigned = torch.device("cuda", local_rank) + selected_backend = backend or "nccl" + else: + assigned = torch.device("cpu") + selected_backend = backend or "gloo" + if not dist.is_initialized(): + dist.init_process_group( + backend=selected_backend, + init_method="env://", + timeout=timedelta(seconds=timeout_seconds), + ) + return DistributedContext(True, rank, local_rank, world_size, assigned) + + +def finalize(context: DistributedContext) -> None: + """Destroy a process group created for ``context``.""" + + if context.enabled and dist.is_available() and dist.is_initialized(): + dist.destroy_process_group() + + +@contextmanager +def process_group(**kwargs: object) -> Iterator[DistributedContext]: + context = initialize(**kwargs) + try: + yield context + finally: + finalize(context) diff --git a/src/gnn4colliders/distributed/model.py b/src/gnn4colliders/distributed/model.py new file mode 100644 index 0000000000000000000000000000000000000000..13b53b3c0e257880963587d0ee2e3779499817c7 --- /dev/null +++ b/src/gnn4colliders/distributed/model.py @@ -0,0 +1,20 @@ +"""DDP model wrapping kept outside architecture implementations.""" + +from __future__ import annotations + +from torch import nn +from torch.nn.parallel import DistributedDataParallel + +from .context import DistributedContext + + +def prepare_model(model: nn.Module, context: DistributedContext) -> nn.Module: + model = model.to(context.device) + if not context.enabled: + return model + kwargs = ( + {"device_ids": [context.local_rank], "output_device": context.local_rank} + if context.device.type == "cuda" + else {} + ) + return DistributedDataParallel(model, **kwargs) diff --git a/src/gnn4colliders/distributed/sampler.py b/src/gnn4colliders/distributed/sampler.py new file mode 100644 index 0000000000000000000000000000000000000000..f9f36d63490a7c8e7591bdad7138a67e2dc462bb --- /dev/null +++ b/src/gnn4colliders/distributed/sampler.py @@ -0,0 +1,47 @@ +"""Index sharding for the project's deterministic graph loader.""" + +from __future__ import annotations + +import torch + +from .context import DistributedContext + + +class DistributedIndices: + """Epoch-aware rank partitioning without evaluation padding duplicates.""" + + def __init__( + self, + size: int, + context: DistributedContext, + *, + shuffle: bool, + drop_last: bool = False, + seed: int = 0, + pad: bool = True, + ) -> None: + self.size = size + self.context = context + self.shuffle = shuffle + self.drop_last = drop_last + self.seed = seed + self.pad = pad + self.epoch = 0 + + def set_epoch(self, epoch: int) -> None: + self.epoch = epoch + + def indices(self) -> list[int]: + values = torch.arange(self.size, dtype=torch.long) + if self.shuffle: + generator = torch.Generator().manual_seed(self.seed + self.epoch) + values = values[torch.randperm(self.size, generator=generator)] + if not self.context.enabled: + return values.tolist() + if self.drop_last: + usable = self.size - self.size % self.context.world_size + values = values[:usable] + elif self.pad and self.size % self.context.world_size: + extra = self.context.world_size - self.size % self.context.world_size + values = torch.cat((values, values[:extra])) + return values[self.context.rank :: self.context.world_size].tolist() diff --git a/src/gnn4colliders/export/__init__.py b/src/gnn4colliders/export/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..ed755c7d2371a090bd4cc75d73e3467d4af1d638 --- /dev/null +++ b/src/gnn4colliders/export/__init__.py @@ -0,0 +1,21 @@ +"""Optional ONNX export for already-prepared ROOT-GNN graph tensors.""" + +from .onnx import ( + RootGNNExportAdapter, + RootGNNExportInputs, + export_checkpoint_to_onnx, + export_root_gnn_onnx, + inputs_from_graph_batch, + validate_onnx, + validate_onnx_runtime, +) + +__all__ = [ + "RootGNNExportAdapter", + "RootGNNExportInputs", + "export_checkpoint_to_onnx", + "export_root_gnn_onnx", + "inputs_from_graph_batch", + "validate_onnx", + "validate_onnx_runtime", +] diff --git a/src/gnn4colliders/export/onnx.py b/src/gnn4colliders/export/onnx.py new file mode 100644 index 0000000000000000000000000000000000000000..1d01d54179dc7832446ef1cb4b39f5436c8baea8 --- /dev/null +++ b/src/gnn4colliders/export/onnx.py @@ -0,0 +1,380 @@ +"""Tensor-only ROOT-GNN export and ONNX Runtime validation. + +The native model intentionally remains DGL-based. This module is an +inference/deployment adapter whose public tensor contract is: +``node_features``, ``edge_features``, ``edge_src``, ``edge_dst``, +``node_batch``, and ``global_features``. ``global_features`` is a required +runtime input for a stable ONNX signature; for models without configured +globals it is ignored and may contain one zero column per graph. +""" + +from __future__ import annotations + +import hashlib +import json +import os +import tempfile +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Mapping + +import torch +from torch import nn + +from gnn4colliders.data.metadata import FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork +from gnn4colliders.training import CheckpointManager, load_model_weights + +ONNX_OPSET = 17 +_INPUT_NAMES = ( + "node_features", + "edge_features", + "edge_src", + "edge_dst", + "node_batch", + "global_features", +) + + +@dataclass(frozen=True) +class RootGNNExportInputs: + """Numerical graph representation consumed by the exported model.""" + + node_features: torch.Tensor + edge_features: torch.Tensor + edge_src: torch.Tensor + edge_dst: torch.Tensor + node_batch: torch.Tensor + global_features: torch.Tensor + + @property + def graph_count(self) -> int: + return int(self.global_features.shape[0]) + + def as_tuple(self) -> tuple[torch.Tensor, ...]: + return tuple(getattr(self, name) for name in _INPUT_NAMES) + + +def inputs_from_graph_batch(batch: Any) -> RootGNNExportInputs: + """Convert a ``GraphBatch`` to explicit, metadata-free export inputs.""" + graph = batch.graph if hasattr(batch, "graph") else batch + if not hasattr(graph, "ndata") or not hasattr(graph, "edata"): + raise TypeError("example batch must contain a DGL graph") + node_features = graph.ndata["features"] + edge_features = graph.edata["features"] + if node_features.ndim != 2 or edge_features.ndim != 2: + raise ValueError("graph node and edge features must be rank-2 tensors") + try: + node_counts = graph.batch_num_nodes().to(dtype=torch.long) + edge_counts = graph.batch_num_edges().to(dtype=torch.long) + edge_src, edge_dst = graph.edges(order="eid") + except AttributeError as error: + raise ValueError("export requires a batched DGL graph") from error + if ( + node_counts.numel() == 0 + or int(node_counts.sum()) == 0 + or bool((node_counts == 0).any()) + ): + raise ValueError("empty-node graphs are not valid ROOT-GNN export inputs") + node_batch = torch.repeat_interleave( + torch.arange(len(node_counts), device=node_counts.device), node_counts + ) + edge_batch = torch.repeat_interleave( + torch.arange(len(edge_counts), device=edge_counts.device), edge_counts + ) + # The edge batch is deliberately derived from source node membership in the + # adapter; checking it here catches malformed graph batches early. + if edge_batch.numel() and not torch.equal(edge_batch, node_batch[edge_src]): + raise ValueError("DGL edge ordering does not agree with graph membership") + globals_ = getattr(batch, "global_features", None) + if globals_ is None: + globals_ = node_counts[:, None].to(dtype=node_features.dtype) + elif globals_.ndim == 1: + globals_ = globals_.unsqueeze(0) + if globals_.shape[0] != len(node_counts): + raise ValueError("global_features must have one row per graph") + return RootGNNExportInputs( + node_features=node_features, + edge_features=edge_features, + edge_src=edge_src.to(dtype=torch.long), + edge_dst=edge_dst.to(dtype=torch.long), + node_batch=node_batch.to(dtype=torch.long), + global_features=globals_, + ) + + +def _mean_by_group( + values: torch.Tensor, groups: torch.Tensor, count: int +) -> torch.Tensor: + membership = _group_membership(groups, count, values) + result = membership @ values + sizes = membership.sum(dim=1, keepdim=True) + return result / sizes.clamp_min(1) + + +def _sum_by_group( + values: torch.Tensor, groups: torch.Tensor, count: int +) -> torch.Tensor: + return _group_membership(groups, count, values) @ values + + +def _group_membership( + groups: torch.Tensor, count: int, values: torch.Tensor +) -> torch.Tensor: + """Return a dense group-membership matrix for export-safe reduction. + + ``index_add`` is efficient in native PyTorch but exports to ONNX scatter + operations whose behavior for duplicate indices is not reliable. A + boolean membership matrix has unambiguous reduction semantics for the + repeated node/edge group IDs used by batched graphs. + """ + labels = torch.arange(count, device=groups.device).unsqueeze(1) + return (labels == groups.unsqueeze(0)).to(dtype=values.dtype) + + +class RootGNNExportAdapter(nn.Module): + """Reproduce an ``EdgeNetwork``/``FineTunedEdgeNetwork`` with tensors only.""" + + def __init__(self, model: nn.Module) -> None: + super().__init__() + if isinstance(model, FineTunedEdgeNetwork): + self.backbone = model.backbone + self.classifier = model.classifier + elif isinstance(model, EdgeNetwork): + self.backbone = model + self.classifier = model.classifier + else: + raise TypeError("model must be EdgeNetwork or FineTunedEdgeNetwork") + self.has_global = self.backbone.has_global + + def forward( + self, + node_features: torch.Tensor, + edge_features: torch.Tensor, + edge_src: torch.Tensor, + edge_dst: torch.Tensor, + node_batch: torch.Tensor, + global_features: torch.Tensor, + ) -> torch.Tensor: + graph_count = global_features.shape[0] + node_h = self.backbone.node_encoder(node_features) + edge_h = self.backbone.edge_encoder(edge_features) + if self.has_global: + global_h = self.backbone.global_encoder(global_features) + else: + counts = _sum_by_group( + node_features.new_ones((node_features.shape[0], 1)), + node_batch, + graph_count, + ) + global_h = self.backbone.global_encoder(counts) + edge_batch = node_batch[edge_src] + for _ in range(self.backbone.n_proc_steps): + edge_h = self.backbone.edge_update( + torch.cat( + (edge_h, node_h[edge_src], node_h[edge_dst], global_h[edge_batch]), + 1, + ) + ) + node_messages = _sum_by_group(edge_h, edge_dst, node_h.shape[0]) + node_h = self.backbone.node_update( + torch.cat((node_h, node_messages, global_h[node_batch]), 1) + ) + global_h = self.backbone.global_update( + torch.cat( + ( + global_h, + _mean_by_group(node_h, node_batch, graph_count), + _mean_by_group(edge_h, edge_batch, graph_count), + ), + 1, + ) + ) + return self.classifier(self.backbone.global_decoder(global_h)) + + +def _optional_onnx() -> Any: + try: + import onnx + except ImportError as error: + raise ImportError( + "ONNX export requires the 'onnx' extra. Install with: " + "uv sync --extra root-gnn --extra onnx" + ) from error + return onnx + + +def validate_onnx(path: str | os.PathLike[str]) -> None: + """Run ONNX structural validation and fail clearly on invalid artifacts.""" + onnx = _optional_onnx() + model = onnx.load(str(path)) + onnx.checker.check_model(model) + + +def validate_onnx_runtime( + path: str | os.PathLike[str], + inputs: RootGNNExportInputs, + reference: torch.Tensor, + *, + rtol: float = 1e-4, + atol: float = 1e-5, +) -> torch.Tensor: + """Run CPU ONNX Runtime and compare its raw logits with PyTorch.""" + try: + import numpy as np + import onnxruntime as ort + except ImportError as error: + raise ImportError( + "ONNX Runtime validation requires the 'onnx' extra. Install with: " + "uv sync --extra root-gnn --extra onnx" + ) from error + session = ort.InferenceSession(str(path), providers=["CPUExecutionProvider"]) + values = { + name: tensor.detach().cpu().numpy() + for name, tensor in zip(_INPUT_NAMES, inputs.as_tuple()) + } + result = torch.from_numpy(np.asarray(session.run(["logits"], values)[0])) + torch.testing.assert_close(result, reference.detach().cpu(), rtol=rtol, atol=atol) + return result + + +def _metadata( + model: nn.Module, *, checkpoint: str | None, task_config: Mapping[str, Any] | None +) -> dict[str, Any]: + config = model.checkpoint_config() + return { + "model_family": "root_gnn", + "model_config": config, + "task_config": dict(task_config or {}), + "feature_schema_version": FEATURE_SCHEMA_VERSION, + "graph_schema_version": GRAPH_SCHEMA_VERSION, + "onnx_opset": ONNX_OPSET, + "input_names": list(_INPUT_NAMES), + "output_name": "logits", + "checkpoint": (Path(checkpoint).name if checkpoint else None), + "checkpoint_sha256": ( + hashlib.sha256(Path(checkpoint).read_bytes()).hexdigest() + if checkpoint + else None + ), + } + + +def export_root_gnn_onnx( + model: nn.Module, + example_batch: Any, + output_path: str | os.PathLike[str], + *, + opset: int = ONNX_OPSET, + task_config: Mapping[str, Any] | None = None, + checkpoint: str | None = None, + overwrite: bool = False, +) -> Path: + """Export and validate a ROOT-GNN model from an already prepared batch.""" + if opset != ONNX_OPSET: + raise ValueError(f"only ONNX opset {ONNX_OPSET} is supported") + _optional_onnx() + inputs = ( + example_batch + if isinstance(example_batch, RootGNNExportInputs) + else inputs_from_graph_batch(example_batch) + ) + adapter = RootGNNExportAdapter(model).eval() + output = Path(output_path) + if output.suffix.lower() != ".onnx": + raise ValueError("export output must have a .onnx extension") + if output.exists() and not overwrite: + raise FileExistsError(f"export output already exists: {output}") + output.parent.mkdir(parents=True, exist_ok=True) + with tempfile.TemporaryDirectory( + prefix=f".{output.stem}-", dir=output.parent + ) as temp_dir: + temporary = Path(temp_dir) / output.name + torch.onnx.export( + adapter, + inputs.as_tuple(), + str(temporary), + input_names=list(_INPUT_NAMES), + output_names=["logits"], + opset_version=opset, + dynamic_axes={ + "node_features": {0: "nodes"}, + "edge_features": {0: "edges"}, + "edge_src": {0: "edges"}, + "edge_dst": {0: "edges"}, + "node_batch": {0: "nodes"}, + "global_features": {0: "graphs"}, + "logits": {0: "graphs"}, + }, + do_constant_folding=True, + ) + validate_onnx(temporary) + with torch.inference_mode(): + reference = adapter(*inputs.as_tuple()) + validate_onnx_runtime(temporary, inputs, reference) + temporary.replace(output) + metadata_path = output.with_suffix(output.suffix + ".json") + metadata_path.write_text( + json.dumps( + _metadata(model, checkpoint=checkpoint, task_config=task_config), + indent=2, + sort_keys=True, + ) + + "\n" + ) + return output + + +def _model_from_checkpoint(payload: Mapping[str, Any], example_batch: Any) -> nn.Module: + config = dict(payload.get("model_config") or {}) + if str(config.get("family", "root_gnn")) != "root_gnn": + raise ValueError("checkpoint model family is not root_gnn") + graph = example_batch.graph if hasattr(example_batch, "graph") else None + if graph is None: + raise TypeError("checkpoint reconstruction needs a GraphBatch example") + kwargs = { + key: config[key] + for key in ("hid_size", "n_layers", "n_proc_steps", "dropout") + if key in config + } + model_class = str(config.get("class", "EdgeNetwork")) + if "FineTuned" in model_class or str(config.get("name", "")).startswith("fine"): + backbone = EdgeNetwork( + graph, getattr(example_batch, "global_features", None), out_size=1, **kwargs + ) + model = FineTunedEdgeNetwork( + backbone, + int(config.get("out_size", 1)), + freeze_backbone=bool(config.get("freeze_backbone", False)), + ) + else: + model = EdgeNetwork( + graph, + getattr(example_batch, "global_features", None), + out_size=int(config.get("out_size", 1)), + **kwargs, + ) + load_model_weights(model, payload) + return model + + +def export_checkpoint_to_onnx( + checkpoint_path: str | os.PathLike[str], + output_path: str | os.PathLike[str], + *, + example_batch: Any, + opset: int = ONNX_OPSET, + overwrite: bool = False, +) -> Path: + """Load a Task 11 checkpoint and export its reconstructed ROOT-GNN model.""" + payload = CheckpointManager.load(checkpoint_path, map_location="cpu") + model = _model_from_checkpoint(payload, example_batch) + return export_root_gnn_onnx( + model, + example_batch, + output_path, + opset=opset, + task_config=payload.get("task_config"), + checkpoint=str(checkpoint_path), + overwrite=overwrite, + ) diff --git a/src/gnn4colliders/features/__init__.py b/src/gnn4colliders/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e0e796935fe151dead2757ade59c9592b497732a --- /dev/null +++ b/src/gnn4colliders/features/__init__.py @@ -0,0 +1,5 @@ +"""Collider-domain feature transformations shared across model families.""" + +from .objects import NODE_FEATURE_NAMES, build_node_features, build_object_features + +__all__ = ["NODE_FEATURE_NAMES", "build_node_features", "build_object_features"] diff --git a/src/gnn4colliders/features/objects.py b/src/gnn4colliders/features/objects.py new file mode 100644 index 0000000000000000000000000000000000000000..f9596568d2329a18fdcaee623f192a3142f5b123 --- /dev/null +++ b/src/gnn4colliders/features/objects.py @@ -0,0 +1,120 @@ +"""Architecture-independent collider-object feature construction.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from numbers import Number +from typing import Any + +import numpy as np +import torch + +NODE_FEATURE_NAMES = ( + "pt", + "eta", + "phi", + "energy", + "btag", + "charge", + "node_type", +) + +_CALCULATED_ENERGY = "CALC_E" +_NODE_TYPE = "NODE_TYPE" + + +def _as_feature_tensor(values: list[Any], dtype: torch.dtype) -> torch.Tensor: + return torch.as_tensor(np.asarray(values, dtype=np.float32), dtype=dtype) + + +def build_node_features( + event: Mapping[str, Any], + feature_branches: Sequence[Any], + object_types: Sequence[str], + scales: Sequence[Number] | torch.Tensor, + *, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, list[int]]: + """Build dense node features from one collider event. + + ``feature_branches`` follows the active legacy schema: each item describes + one output column and contains one branch name or constant per object + type. ``object_types`` contains ``"vector"`` or ``"single"`` for each + object type. Rows are concatenated in object-type order. + + The returned tensor has columns in :data:`NODE_FEATURE_NAMES` order for + the active schema, and the second result records rows contributed by each + object type. Inputs are read but never modified. + """ + if not feature_branches: + return torch.empty((0, 0), dtype=dtype), [] + first_column = feature_branches[0] + if len(first_column) != len(object_types): + raise ValueError("one branch specification is required per object type") + + lengths: list[int] = [] + for branch, object_type in zip(first_column, object_types): + if object_type == "single": + lengths.append(1) + elif object_type == "vector": + lengths.append(len(event[branch])) + else: + raise ValueError(f"unknown object type: {object_type!r}") + + columns: list[torch.Tensor] = [] + for column_index, specification in enumerate(feature_branches): + if specification == _CALCULATED_ENERGY: + columns.append(columns[0] * torch.cosh(columns[1])) + continue + if specification == _NODE_TYPE: + values = [ + object_type_index + for object_type_index, length in enumerate(lengths) + for _ in range(length) + ] + columns.append(torch.tensor(values, dtype=dtype)) + continue + if len(specification) != len(object_types): + raise ValueError( + f"feature column {column_index} has the wrong number of branches" + ) + + values: list[Any] = [] + for object_type_index, (length, branch, object_type) in enumerate( + zip(lengths, specification, object_types) + ): + if isinstance(branch, Number): + values.extend([branch] * length) + elif branch == _CALCULATED_ENERGY: + start = sum(lengths[:object_type_index]) + stop = start + length + values.extend( + ( + columns[0][start:stop] * torch.cosh(columns[1][start:stop]) + ).tolist() + ) + elif object_type == "single": + values.append(event[branch]) + else: + values.extend(event[branch]) + columns.append(_as_feature_tensor(values, dtype)) + + features = torch.stack(columns, dim=1) + scale_tensor = torch.as_tensor(scales, dtype=dtype) + if scale_tensor.ndim != 1 or scale_tensor.numel() != features.shape[1]: + raise ValueError("scales must contain one value per feature column") + return features * scale_tensor, lengths + + +def build_object_features( + event: Mapping[str, Any], + feature_branches: Sequence[Any], + object_types: Sequence[str], + scales: Sequence[Number] | torch.Tensor, + *, + dtype: torch.dtype = torch.float32, +) -> tuple[torch.Tensor, list[int]]: + """Alias for :func:`build_node_features` using domain-neutral wording.""" + return build_node_features( + event, feature_branches, object_types, scales, dtype=dtype + ) diff --git a/src/gnn4colliders/graphs/__init__.py b/src/gnn4colliders/graphs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..d9c70a2d3886bfb5077ad9c69ea7334197b5a230 --- /dev/null +++ b/src/gnn4colliders/graphs/__init__.py @@ -0,0 +1,13 @@ +"""Graph representations and graph-specific transformations.""" + +from .dgl import build_dgl_graph +from .edges import build_edge_features +from .topology import clear_topology_cache, fully_connected_edges, topology_cache_info + +__all__ = [ + "build_dgl_graph", + "build_edge_features", + "fully_connected_edges", + "clear_topology_cache", + "topology_cache_info", +] diff --git a/src/gnn4colliders/graphs/dgl.py b/src/gnn4colliders/graphs/dgl.py new file mode 100644 index 0000000000000000000000000000000000000000..c27f418d3b846ac0b5f7ea8e13badf4f21f01f0a --- /dev/null +++ b/src/gnn4colliders/graphs/dgl.py @@ -0,0 +1,44 @@ +"""DGL boundary for ROOT-GNN graph construction.""" + +from __future__ import annotations + +import torch + +from gnn4colliders.features import NODE_FEATURE_NAMES + +from .edges import build_edge_features +from .topology import fully_connected_edges + + +def build_dgl_graph( + node_features: torch.Tensor, + *, + self_loops: bool = False, + eta_index: int = NODE_FEATURE_NAMES.index("eta"), + phi_index: int = NODE_FEATURE_NAMES.index("phi"), +): + """Build a DGL graph with compatible node and edge feature keys. + + DGL is imported only when this representation boundary is used. Inputs + are not moved between devices; callers control placement explicitly. + """ + try: + import dgl + except ImportError as error: # pragma: no cover - depends on optional extra + raise ImportError( + "build_dgl_graph requires the optional 'root-gnn' dependency" + ) from error + + if node_features.ndim != 2: + raise ValueError("node_features must be a two-dimensional tensor") + src, dst = fully_connected_edges( + node_features.shape[0], self_loops=self_loops, device=node_features.device + ) + graph = dgl.graph( + (src, dst), num_nodes=node_features.shape[0], device=node_features.device + ) + graph.ndata["features"] = node_features + graph.edata["features"] = build_edge_features( + node_features, src, dst, eta_index=eta_index, phi_index=phi_index + ) + return graph diff --git a/src/gnn4colliders/graphs/edges.py b/src/gnn4colliders/graphs/edges.py new file mode 100644 index 0000000000000000000000000000000000000000..1b47d20a2e19911ab9d4fa56fa03f098918d8b7e --- /dev/null +++ b/src/gnn4colliders/graphs/edges.py @@ -0,0 +1,42 @@ +"""Edge-feature construction for collider graphs.""" + +from __future__ import annotations + +import math + +import torch + + +def build_edge_features( + node_features: torch.Tensor, + src: torch.Tensor, + dst: torch.Tensor, + *, + eta_index: int, + phi_index: int, +) -> torch.Tensor: + """Calculate edge features in ``[deta, dphi, dR]`` order. + + Differences are signed from ``src`` to ``dst`` and phi is wrapped using + the same strict boundary convention as the legacy implementation. + """ + if node_features.ndim != 2: + raise ValueError("node_features must be a two-dimensional tensor") + if src.ndim != 1 or dst.ndim != 1 or src.shape != dst.shape: + raise ValueError("src and dst must be one-dimensional tensors of equal shape") + if not 0 <= eta_index < node_features.shape[1]: + raise IndexError("eta_index is outside node_features") + if not 0 <= phi_index < node_features.shape[1]: + raise IndexError("phi_index is outside node_features") + + eta_src = node_features[src, eta_index] + eta_dst = node_features[dst, eta_index] + phi_src = node_features[src, phi_index] + phi_dst = node_features[dst, phi_index] + deta = eta_src - eta_dst + dphi = phi_src - phi_dst + pi = math.pi + dphi = torch.where(dphi > pi, dphi - 2 * pi, dphi) + dphi = torch.where(dphi < -pi, dphi + 2 * pi, dphi) + dr = torch.sqrt(deta.square() + dphi.square()) + return torch.stack((deta, dphi, dr), dim=1) diff --git a/src/gnn4colliders/graphs/topology.py b/src/gnn4colliders/graphs/topology.py new file mode 100644 index 0000000000000000000000000000000000000000..b72c159ab44755a732d911f3156ebab3f815664e --- /dev/null +++ b/src/gnn4colliders/graphs/topology.py @@ -0,0 +1,61 @@ +"""Graph topology construction for graph-based representations.""" + +from __future__ import annotations + +from collections import OrderedDict + +import torch + +_TOPOLOGY_CACHE_LIMIT = 32 +_TOPOLOGY_CACHE: OrderedDict[ + tuple[int, bool, str, int | None], tuple[torch.Tensor, torch.Tensor] +] = OrderedDict() + + +def clear_topology_cache() -> None: + """Clear the bounded topology cache, primarily for benchmarks/tests.""" + _TOPOLOGY_CACHE.clear() + + +def topology_cache_info() -> dict[str, int]: + """Return cache size without exposing mutable cache internals.""" + return {"size": len(_TOPOLOGY_CACHE), "max_size": _TOPOLOGY_CACHE_LIMIT} + + +def fully_connected_edges( + num_nodes: int, + *, + self_loops: bool = False, + device: torch.device | str | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Return source-major directed all-pairs edges. + + The legacy ROOT-GNN constructor keeps the self-loop for a one-node graph + even when ``self_loops`` is false. This detail is preserved here. + """ + if not isinstance(num_nodes, int) or isinstance(num_nodes, bool): + raise TypeError("num_nodes must be an integer") + if num_nodes < 0: + raise ValueError("num_nodes must be non-negative") + + target = torch.device(device) if device is not None else torch.device("cpu") + key = (num_nodes, self_loops, target.type, target.index) + cached = _TOPOLOGY_CACHE.get(key) + if cached is not None: + _TOPOLOGY_CACHE.move_to_end(key) + # Preserve the historical fresh-tensor API: callers may safely mutate + # their result without corrupting later graph constructions. + return cached[0].clone(), cached[1].clone() + + indices = torch.arange(num_nodes, dtype=torch.long, device=target) + source = indices.repeat_interleave(num_nodes) + destination = indices.repeat(num_nodes) + if not self_loops and num_nodes > 1: + keep = source != destination + source = source[keep] + destination = destination[keep] + _TOPOLOGY_CACHE[key] = (source, destination) + _TOPOLOGY_CACHE.move_to_end(key) + while len(_TOPOLOGY_CACHE) > _TOPOLOGY_CACHE_LIMIT: + _TOPOLOGY_CACHE.popitem(last=False) + return source, destination diff --git a/src/gnn4colliders/inference/__init__.py b/src/gnn4colliders/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cea69a0cf6ab558dba7e939bcb9bbbb1fade56d5 --- /dev/null +++ b/src/gnn4colliders/inference/__init__.py @@ -0,0 +1,15 @@ +"""Prediction, evaluation, and output writing infrastructure.""" + +from .predictor import Predictor, load_model_and_task_for_inference +from .results import EvaluationResult, PredictionResult +from .root_writer import write_root_scores +from .writers import write_npz + +__all__ = [ + "EvaluationResult", + "PredictionResult", + "Predictor", + "load_model_and_task_for_inference", + "write_npz", + "write_root_scores", +] diff --git a/src/gnn4colliders/inference/predictor.py b/src/gnn4colliders/inference/predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..0c79cb893b9f9c9e7f51bf1d3ebccb3f9dc76ced --- /dev/null +++ b/src/gnn4colliders/inference/predictor.py @@ -0,0 +1,231 @@ +"""Model execution for prediction and full-split evaluation.""" + +from __future__ import annotations + +import inspect +from collections.abc import Callable, Mapping +from types import SimpleNamespace +from typing import Any + +import torch +from torch import nn + +from gnn4colliders.data.graph_dataset import GraphBatch +from gnn4colliders.distributed import ( + DistributedContext, + gather_objects, + gather_tensor, + prepare_model, +) +from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask +from gnn4colliders.training.checkpoint import CheckpointManager, load_model_weights + +from .results import EvaluationResult, PredictionResult + + +def _forward(model: nn.Module, batch: Any) -> torch.Tensor: + """Call the public GraphBatch or graph/global model boundary.""" + if isinstance(batch, GraphBatch): + return model(batch.graph, batch.global_features) + try: + signature = inspect.signature(model.forward) + positional = [ + parameter + for parameter in signature.parameters.values() + if parameter.kind + in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD) + ] + except (TypeError, ValueError): + positional = [] + if len(positional) <= 1: + return model(batch) + if hasattr(batch, "graph"): + return model(batch.graph, getattr(batch, "global_features", None)) + raise TypeError("model requires graph/global_features, but batch has no graph") + + +def _batch_metadata(batch: Any) -> Any: + return getattr(batch, "metadata", None) + + +def _labels(batch: Any) -> torch.Tensor | None: + labels = getattr(batch, "labels", None) + return None if labels is None else torch.as_tensor(labels).detach().cpu() + + +class Predictor: + """Run a task-aware model over an ordered loader.""" + + def __init__( + self, + model: nn.Module, + task: Any, + *, + device: torch.device | str = "cpu", + distributed_context: DistributedContext | None = None, + ) -> None: + self.distributed_context = ( + distributed_context or DistributedContext.single_process(device) + ) + self.device = self.distributed_context.device + self.model = prepare_model(model, self.distributed_context) + self.task = task + + def predict(self, loader: Any) -> PredictionResult: + self.model.eval() + logits_parts: list[torch.Tensor] = [] + label_parts: list[torch.Tensor] = [] + sample_ids: list[str] = [] + metadata_parts: list[Any] = [] + has_labels = True + + with torch.inference_mode(): + for batch in loader: + moved = batch.to(self.device) if hasattr(batch, "to") else batch + logits = _forward(self.model, moved) + if not isinstance(logits, torch.Tensor): + raise TypeError("model must return a torch.Tensor of logits") + logits_parts.append(logits.detach().cpu()) + labels = _labels(moved) + if labels is None: + has_labels = False + else: + label_parts.append(labels) + metadata = _batch_metadata(moved) + if metadata is not None: + ids = getattr(metadata, "sample_id", None) + if ids is not None: + sample_ids.extend(str(value) for value in ids) + metadata_parts.append(metadata) + + if not logits_parts: + raise ValueError("cannot predict on an empty loader") + logits = torch.cat(logits_parts, dim=0) + if self.distributed_context.enabled: + logits = gather_tensor(logits, self.distributed_context) + if label_parts: + labels_local = torch.cat(label_parts, dim=0) + labels_result = gather_tensor(labels_local, self.distributed_context) + else: + labels_result = None + ids = tuple( + value + for group in gather_objects(tuple(sample_ids), self.distributed_context) + for value in group + ) + local_metadata = _combine_metadata(metadata_parts) + if local_metadata is not None and hasattr(local_metadata, "weight"): + metadata = SimpleNamespace( + fold=gather_tensor(local_metadata.fold, self.distributed_context), + weight=gather_tensor( + local_metadata.weight, self.distributed_context + ), + sample_id=ids, + ) + sample_ids = list(ids) + else: + labels_result = torch.cat(label_parts, dim=0) if has_labels else None + task_output = self.task.predictions(logits) + if not self.distributed_context.enabled: + metadata = _combine_metadata(metadata_parts) + if not sample_ids and metadata is not None: + sample_ids = [str(value) for value in getattr(metadata, "sample_id", ())] + if len(sample_ids) != logits.shape[0]: + raise ValueError( + "loader metadata must provide one sample_id for every prediction" + ) + return PredictionResult( + sample_ids=tuple(sample_ids), + logits=logits, + scores=task_output["scores"].detach().cpu(), + predictions=task_output["predictions"].detach().cpu(), + labels=labels_result, + metadata=metadata, + extra=getattr(metadata, "extra", {}) if metadata is not None else {}, + ) + + def evaluate(self, loader: Any) -> EvaluationResult: + result = self.predict(loader) + if result.labels is None: + raise ValueError("evaluation requires labels") + weights = result.weights + batch = SimpleNamespace( + labels=result.labels, + metadata=SimpleNamespace(weight=weights) if weights is not None else None, + ) + if weights is None: + raise ValueError("evaluation requires named metadata.weight") + metrics = { + key: float(value) + for key, value in self.task.metrics(result.logits, batch).items() + } + return EvaluationResult(predictions=result, metrics=metrics) + + +def _combine_metadata(parts: list[Any]) -> Any | None: + if not parts: + return None + first = parts[0] + cls = type(first) + fields = ("fold", "weight", "sample_id", "extra") + values: dict[str, Any] = {} + for name in fields: + if not hasattr(first, name): + continue + current = [getattr(part, name) for part in parts] + if name in {"fold", "weight"}: + values[name] = torch.cat( + [torch.as_tensor(value).cpu() for value in current] + ) + elif name == "sample_id": + values[name] = tuple(value for group in current for value in group) + elif name == "extra": + keys = sorted({key for group in current for key in group}) + values[name] = { + key: tuple(value for group in current for value in group.get(key, ())) + for key in keys + } + try: + return cls(**values) + except TypeError: + return SimpleNamespace(**values) + + +def load_model_and_task_for_inference( + checkpoint: str, + *, + model_factory: Callable[[Mapping[str, Any]], nn.Module], + task_factory: Callable[[Mapping[str, Any]], Any] | None = None, + map_location: torch.device | str = "cpu", +) -> tuple[nn.Module, Any]: + """Build a model from checkpoint metadata and load weights only. + + Factories remain caller-owned because ROOT-GNN construction needs a sample + graph. The helper deliberately never restores optimizer or lifecycle state. + """ + checkpoint_data = CheckpointManager.load(checkpoint, map_location=map_location) + model_config = checkpoint_data.get("model_config") + if not isinstance(model_config, Mapping): + raise ValueError("checkpoint does not contain reconstructable model_config") + model = model_factory(model_config) + load_model_weights(model, checkpoint_data) + task_config = checkpoint_data.get("task_config") or {} + if task_factory is not None: + task = task_factory(task_config) + else: + task_type = str(task_config.get("type", "")).lower() + if "multi" in task_type: + task = MulticlassClassificationTask( + **_task_kwargs(task_config, "absolute_weights") + ) + elif "binary" in task_type: + task = BinaryClassificationTask( + **_task_kwargs(task_config, "absolute_weights", "threshold") + ) + else: + raise ValueError("checkpoint does not contain a supported task_config") + return model, task + + +def _task_kwargs(config: Mapping[str, Any], *names: str) -> dict[str, Any]: + return {name: config[name] for name in names if name in config} diff --git a/src/gnn4colliders/inference/results.py b/src/gnn4colliders/inference/results.py new file mode 100644 index 0000000000000000000000000000000000000000..d18e3df290a262b37819de7ea566811318a677ef --- /dev/null +++ b/src/gnn4colliders/inference/results.py @@ -0,0 +1,52 @@ +"""Typed outputs produced by model inference.""" + +from __future__ import annotations + +from collections.abc import Mapping, Sequence +from dataclasses import dataclass, field +from typing import Any + +import torch + + +@dataclass(frozen=True) +class PredictionResult: + """Predictions in the exact order in which the loader yielded events. + + Tensor fields are CPU tensors. Keeping the tensors detached and on the + CPU makes a result safe to retain after prediction and avoids tying its + lifetime to a CUDA context. + """ + + sample_ids: tuple[str, ...] + logits: torch.Tensor + scores: torch.Tensor + predictions: torch.Tensor + labels: torch.Tensor | None = None + metadata: Any | None = None + extra: Mapping[str, Sequence[Any]] = field(default_factory=dict) + + @property + def fold(self) -> torch.Tensor | None: + return _metadata_value(self.metadata, "fold") + + @property + def weights(self) -> torch.Tensor | None: + return _metadata_value(self.metadata, "weight") + + +@dataclass(frozen=True) +class EvaluationResult: + """A prediction result and metrics calculated over the complete split.""" + + predictions: PredictionResult + metrics: Mapping[str, float] + + +def _metadata_value(metadata: Any, name: str) -> torch.Tensor | None: + if metadata is None: + return None + value = getattr(metadata, name, None) + if value is None and isinstance(metadata, Mapping): + value = metadata.get(name) + return None if value is None else torch.as_tensor(value).cpu() diff --git a/src/gnn4colliders/inference/root_writer.py b/src/gnn4colliders/inference/root_writer.py new file mode 100644 index 0000000000000000000000000000000000000000..9a46ecc379da0d2490062a4a15e72ab3999ef37a --- /dev/null +++ b/src/gnn4colliders/inference/root_writer.py @@ -0,0 +1,79 @@ +"""Optional ROOT score output, kept separate from model execution.""" + +from __future__ import annotations + +import re +from pathlib import Path + +import numpy as np + +from .results import PredictionResult + +_ENTRY_SUFFIX = re.compile(r":(\d+)$") + + +def write_root_scores( + result: PredictionResult, + *, + input_path: str | Path, + output_path: str | Path, + tree_name: str, + score_prefix: str = "score", +) -> Path: + """Clone a ROOT tree's columns and add aligned score branches. + + Sample IDs ending in ``:`` are aligned by source entry. If IDs do + not carry an entry suffix, the result must contain one score per source + entry; this makes accidental misalignment fail loudly. Unselected source + entries receive NaN scores and ``selection_pass=0``. + """ + try: + import uproot + except ImportError as error: # pragma: no cover - optional dependency + raise ImportError("write_root_scores requires uproot") from error + + source = Path(input_path) + target = Path(output_path) + target.parent.mkdir(parents=True, exist_ok=True) + with uproot.open(source) as root_file: + tree = root_file[tree_name] + columns = tree.arrays(library="np") + entries = int(tree.num_entries) + + indices: list[int] = [] + for sample_id in result.sample_ids: + match = _ENTRY_SUFFIX.search(sample_id) + if match is None: + if len(result.sample_ids) != entries: + raise ValueError( + "ROOT alignment requires sample IDs ending in ':' " + "when predictions do not cover every source entry" + ) + indices = list(range(entries)) + break + indices.append(int(match.group(1))) + if len(indices) != len(result.sample_ids) or any( + index < 0 or index >= entries for index in indices + ): + raise ValueError("prediction sample IDs contain invalid ROOT entry indices") + + selection_pass = np.zeros(entries, dtype=np.int32) + selection_pass[indices] = 1 + output = dict(columns) + output["selection_pass"] = selection_pass + scores = np.asarray(result.scores) + if scores.ndim == 1: + branch = np.full(entries, np.nan, dtype=np.float32) + branch[indices] = scores.astype(np.float32, copy=False) + output[score_prefix] = branch + elif scores.ndim == 2: + for class_index in range(scores.shape[1]): + branch = np.full(entries, np.nan, dtype=np.float32) + branch[indices] = scores[:, class_index].astype(np.float32, copy=False) + output[f"{score_prefix}_class_{class_index}"] = branch + else: + raise ValueError("scores must have shape [N] or [N, C]") + + with uproot.recreate(target) as root_file: + root_file[tree_name] = output + return target diff --git a/src/gnn4colliders/inference/writers.py b/src/gnn4colliders/inference/writers.py new file mode 100644 index 0000000000000000000000000000000000000000..aeae06779d08491f4a8b583c6318a3decd55365e --- /dev/null +++ b/src/gnn4colliders/inference/writers.py @@ -0,0 +1,32 @@ +"""Serialization of prediction results.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any + +import numpy as np + +from .results import PredictionResult + + +def write_npz(result: PredictionResult, path: str | Path) -> Path: + """Write named prediction fields to a compressed NumPy archive.""" + target = Path(path) + target.parent.mkdir(parents=True, exist_ok=True) + values: dict[str, Any] = { + "sample_id": np.asarray(result.sample_ids, dtype=str), + "logits": result.logits.numpy(), + "scores": result.scores.numpy(), + "predictions": result.predictions.numpy(), + } + if result.labels is not None: + values["labels"] = result.labels.numpy() + for name, attribute in (("fold", "fold"), ("weight", "weights")): + value = getattr(result, attribute) + if value is not None: + values[name] = value.numpy() + for name, value in result.extra.items(): + values[name] = np.asarray(value) + np.savez_compressed(target, **values) + return target diff --git a/src/gnn4colliders/models/__init__.py b/src/gnn4colliders/models/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..468a273f591cb05c6ee5045f3bdef2458f7dd274 --- /dev/null +++ b/src/gnn4colliders/models/__init__.py @@ -0,0 +1 @@ +"""Architecture-specific machine-learning models.""" diff --git a/src/gnn4colliders/models/root_gnn/__init__.py b/src/gnn4colliders/models/root_gnn/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..37ff1fdfe26b3f68a171f3120355c809b1a4e4b5 --- /dev/null +++ b/src/gnn4colliders/models/root_gnn/__init__.py @@ -0,0 +1,6 @@ +"""ROOT-GNN model components.""" + +from .edge_network import EdgeNetwork +from .transfer import FineTunedEdgeNetwork, load_legacy_edge_network_state_dict + +__all__ = ["EdgeNetwork", "FineTunedEdgeNetwork", "load_legacy_edge_network_state_dict"] diff --git a/src/gnn4colliders/models/root_gnn/blocks.py b/src/gnn4colliders/models/root_gnn/blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..2fe341f9c40bb1ba1edfc2fdf4b2ff61b6aac228 --- /dev/null +++ b/src/gnn4colliders/models/root_gnn/blocks.py @@ -0,0 +1,39 @@ +"""Small neural-network blocks used by the active ROOT-GNN model.""" + +from __future__ import annotations + +from collections.abc import Callable + +from torch import nn + + +def make_mlp( + in_size: int, + hidden_size: int, + out_size: int, + n_layers: int, + *, + dropout: float = 0.0, + activation: Callable[[], nn.Module] = nn.ReLU, +) -> nn.Sequential: + """Build the legacy ``Make_MLP`` block. + + Every linear layer is followed by activation and dropout, including the + final linear layer; LayerNorm is then applied to the output. This order + is part of the ROOT-GNN forward contract. + """ + if n_layers < 1: + raise ValueError("n_layers must be positive") + layers: list[nn.Module] = [] + if n_layers == 1: + sizes = [(in_size, out_size)] + else: + sizes = [(in_size, hidden_size)] + sizes.extend((hidden_size, hidden_size) for _ in range(n_layers - 2)) + sizes.append((hidden_size, out_size)) + for input_size, output_size in sizes: + layers.extend( + (nn.Linear(input_size, output_size), activation(), nn.Dropout(dropout)) + ) + layers.append(nn.LayerNorm(out_size)) + return nn.Sequential(*layers) diff --git a/src/gnn4colliders/models/root_gnn/edge_network.py b/src/gnn4colliders/models/root_gnn/edge_network.py new file mode 100644 index 0000000000000000000000000000000000000000..8aefeefb4f31cd6b54893b59cb3dff50ec37a2d8 --- /dev/null +++ b/src/gnn4colliders/models/root_gnn/edge_network.py @@ -0,0 +1,188 @@ +"""The active edge-message-passing ROOT-GNN architecture.""" + +from __future__ import annotations + +from typing import Any + +import torch +from torch import nn + +from .blocks import make_mlp + + +def _broadcast(values: torch.Tensor, counts: torch.Tensor) -> torch.Tensor: + return torch.repeat_interleave(values, counts, dim=0) + + +def _copy_destination(edges: Any) -> dict[str, torch.Tensor]: + """DGL 2.x has ``copy_u`` but no built-in destination copy primitive.""" + return {"m_v": edges.dst["h"]} + + +class EdgeNetwork(nn.Module): + """ROOT-GNN ``Edge_Network`` with an explicit feature/classifier split.""" + + def __init__( + self, + sample_graph: Any, + sample_global: torch.Tensor | None, + hid_size: int, + out_size: int, + n_layers: int, + n_proc_steps: int, + dropout: float = 0.0, + **_: Any, + ) -> None: + super().__init__() + if n_proc_steps < 0: + raise ValueError("n_proc_steps must be non-negative") + node_features = sample_graph.ndata["features"] + edge_features = sample_graph.edata["features"] + global_width = 0 if sample_global is None else sample_global.shape[1] + self.has_global = global_width != 0 + if not self.has_global: + global_width = 1 + self.hid_size = hid_size + self.n_layers = n_layers + self.n_proc_steps = n_proc_steps + self.node_feature_size = int(node_features.shape[1]) + self.edge_feature_size = int(edge_features.shape[1]) + self.global_feature_size = ( + int(sample_global.shape[1]) if sample_global is not None else 0 + ) + self.dropout = float(dropout) + + self.node_encoder = make_mlp( + node_features.shape[1], hid_size, hid_size, n_layers, dropout=dropout + ) + self.edge_encoder = make_mlp( + edge_features.shape[1], hid_size, hid_size, n_layers, dropout=dropout + ) + self.global_encoder = make_mlp( + global_width, hid_size, hid_size, n_layers, dropout=dropout + ) + self.node_update = make_mlp( + 3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout + ) + self.edge_update = make_mlp( + 4 * hid_size, hid_size, hid_size, n_layers, dropout=dropout + ) + self.global_update = make_mlp( + 3 * hid_size, hid_size, hid_size, n_layers, dropout=dropout + ) + self.global_decoder = make_mlp( + hid_size, hid_size, hid_size, n_layers, dropout=dropout + ) + self.classifier = nn.Linear(hid_size, out_size) + + def checkpoint_config(self) -> dict[str, Any]: + """Return constructor information without serializing a sample graph.""" + return { + "family": "root_gnn", + "class": type(self).__name__, + "node_feature_size": self.node_feature_size, + "edge_feature_size": self.edge_feature_size, + "global_feature_size": self.global_feature_size, + "hid_size": self.hid_size, + "out_size": int(self.classifier.out_features), + "n_layers": self.n_layers, + "n_proc_steps": self.n_proc_steps, + "dropout": self.dropout, + } + + # Compatibility with the legacy public attribute while keeping the new + # name in the model API. + @property + def classify(self) -> nn.Linear: + return self.classifier + + def _features( + self, graph: Any, global_features: torch.Tensor | None + ) -> torch.Tensor: + if not self.has_global: + return graph.batch_num_nodes()[:, None].to( + device=graph.ndata["features"].device, dtype=torch.float32 + ) + if global_features is None: + raise ValueError("global_features are required for this model") + if global_features.ndim == 1: + if graph.batch_num_nodes().numel() != 1: + raise ValueError( + "one-dimensional global_features are valid only for one graph" + ) + global_features = global_features.unsqueeze(0) + return global_features + + def forward_features( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + """Return the decoded graph representation before classification.""" + try: + import dgl + except ImportError as error: # pragma: no cover - optional dependency + raise ImportError("EdgeNetwork requires the 'root-gnn' extra") from error + if hasattr(graph, "graph") and hasattr(graph, "global_features"): + global_features = graph.global_features + graph = graph.graph + with graph.local_scope(): + h = self.node_encoder(graph.ndata["features"]) + e = self.edge_encoder(graph.edata["features"]) + graph.ndata["h"] = h + graph.edata["e"] = e + h_global = self.global_encoder(self._features(graph, global_features)) + counts = graph.batch_num_nodes() + for _ in range(self.n_proc_steps): + graph.apply_edges(dgl.function.copy_u("h", "m_u")) + graph.apply_edges(_copy_destination) + graph.edata["e"] = self.edge_update( + torch.cat( + ( + graph.edata["e"], + graph.edata["m_u"], + graph.edata["m_v"], + _broadcast(h_global, graph.batch_num_edges()), + ), + dim=1, + ) + ) + graph.update_all( + dgl.function.copy_e("e", "m"), dgl.function.sum("m", "h_e") + ) + graph.ndata["h"] = self.node_update( + torch.cat( + ( + graph.ndata["h"], + graph.ndata["h_e"], + _broadcast(h_global, counts), + ), + dim=1, + ) + ) + if "w" in graph.ndata: + mask = torch.any(graph.ndata["features"] != 0, dim=1) + valid_counts = [] + start = 0 + for count in counts.tolist(): + valid_counts.append(mask[start : start + count].sum()) + start += count + denominator = ( + torch.stack(valid_counts).to(h.dtype).clamp_min(1)[:, None] + ) + mean_nodes = dgl.sum_nodes(graph, "h", "w") / denominator + else: + mean_nodes = dgl.mean_nodes(graph, "h") + h_global = self.global_update( + torch.cat((h_global, mean_nodes, dgl.mean_edges(graph, "e")), dim=1) + ) + return self.global_decoder(h_global) + + def forward( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + return self.classifier(self.forward_features(graph, global_features)) + + def representation( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + """Legacy alias for the reusable decoded representation.""" + return self.forward_features(graph, global_features) diff --git a/src/gnn4colliders/models/root_gnn/transfer.py b/src/gnn4colliders/models/root_gnn/transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..ed6c0772d551336f1f775b9ff1be83ddde4dd3be --- /dev/null +++ b/src/gnn4colliders/models/root_gnn/transfer.py @@ -0,0 +1,83 @@ +"""Explicit transfer-learning wrapper for ROOT-GNN.""" + +from __future__ import annotations + +import copy +from collections.abc import Mapping +from typing import Any + +import torch +from torch import nn + +from gnn4colliders.compat.checkpoint import map_legacy_edge_network_state_dict + +from .edge_network import EdgeNetwork + + +def load_legacy_edge_network_state_dict( + model: nn.Module, state: Mapping[str, Any] +) -> None: + """Load a legacy checkpoint/state dict after historical prefix cleanup.""" + state_dict = state.get("model_state_dict", state) + model.load_state_dict(map_legacy_edge_network_state_dict(state_dict), strict=True) + + +class FineTunedEdgeNetwork(nn.Module): + """A pretrained ROOT-GNN backbone with a task-specific classifier.""" + + def __init__( + self, backbone: EdgeNetwork, out_size: int, *, freeze_backbone: bool = False + ) -> None: + super().__init__() + # Own an independent classifier-free copy: constructing a transfer + # model must not mutate the pretrained model supplied by the caller. + self.backbone = copy.deepcopy(backbone) + self.backbone.classifier = nn.Identity() + self.classifier = nn.Linear(backbone.hid_size, out_size) + self.freeze_backbone = freeze_backbone + self.set_backbone_trainable(not freeze_backbone) + + def checkpoint_config(self) -> dict[str, Any]: + config = self.backbone.checkpoint_config() + config.update( + { + "class": type(self).__name__, + "out_size": int(self.classifier.out_features), + "freeze_backbone": self.freeze_backbone, + } + ) + return config + + @classmethod + def from_pretrained( + cls, + pretrained: EdgeNetwork, + out_size: int, + *, + freeze_backbone: bool = False, + state_dict: Mapping[str, Any] | None = None, + ) -> "FineTunedEdgeNetwork": + if state_dict is not None: + load_legacy_edge_network_state_dict(pretrained, state_dict) + return cls(pretrained, out_size, freeze_backbone=freeze_backbone) + + def set_backbone_trainable(self, trainable: bool) -> None: + for parameter in self.backbone.parameters(): + parameter.requires_grad = trainable + for parameter in self.classifier.parameters(): + parameter.requires_grad = True + + def forward_features( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + return self.backbone.forward_features(graph, global_features) + + def forward( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + return self.classifier(self.forward_features(graph, global_features)) + + def representation( + self, graph: Any, global_features: torch.Tensor | None = None + ) -> torch.Tensor: + return self.forward_features(graph, global_features) diff --git a/src/gnn4colliders/tasks/__init__.py b/src/gnn4colliders/tasks/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..812a842da392d1691134f75c6b448895b015c3cb --- /dev/null +++ b/src/gnn4colliders/tasks/__init__.py @@ -0,0 +1,6 @@ +"""Learning-task semantics for interpreting model outputs.""" + +from .binary_classification import BinaryClassificationTask +from .multiclass_classification import MulticlassClassificationTask + +__all__ = ["BinaryClassificationTask", "MulticlassClassificationTask"] diff --git a/src/gnn4colliders/tasks/_common.py b/src/gnn4colliders/tasks/_common.py new file mode 100644 index 0000000000000000000000000000000000000000..957820ca06649eecb83ed15b28bf21d61700c666 --- /dev/null +++ b/src/gnn4colliders/tasks/_common.py @@ -0,0 +1,78 @@ +"""Small shared helpers for classification tasks.""" + +from __future__ import annotations + +from typing import Any + +import torch + + +def batch_labels(batch: Any) -> torch.Tensor: + """Return labels from a GraphBatch-like object.""" + + try: + return batch.labels + except AttributeError as error: + raise TypeError("batch must provide a named 'labels' field") from error + + +def batch_weights( + batch: Any, *, device: torch.device, dtype: torch.dtype +) -> torch.Tensor: + """Return event weights through named metadata, never positional tracking.""" + + try: + weights = batch.metadata.weight + except AttributeError as error: + raise TypeError("batch must provide named metadata.weight") from error + return torch.as_tensor(weights, device=device, dtype=dtype).reshape(-1) + + +def per_label_weighted_mean( + elementwise_loss: torch.Tensor, + labels: torch.Tensor, + weights: torch.Tensor, +) -> torch.Tensor: + """Reproduce the legacy weighted loss reduction. + + Each label receives its own weighted mean, and those means are then + averaged equally. In particular, this is not ordinary weighted BCE/CE. + """ + + result = elementwise_loss.new_zeros(()) + unique_labels = torch.unique(labels) + for label in unique_labels: + mask = labels == label + result = ( + result + + (weights[mask] * elementwise_loss[mask]).sum() / weights[mask].sum() + ) + return result / len(unique_labels) + + +def positive_weight_mask(weights: torch.Tensor) -> torch.Tensor: + """The legacy metric path excludes non-positive original weights.""" + + return weights > 0 + + +def weighted_binary_auc( + labels: torch.Tensor, scores: torch.Tensor, weights: torch.Tensor +) -> float: + """Compute weighted binary ROC AUC, including half-credit ties.""" + + labels = labels.detach().cpu().to(dtype=torch.bool) + scores = scores.detach().cpu().to(dtype=torch.float64) + weights = weights.detach().cpu().to(dtype=torch.float64) + positive = labels + negative = ~labels + positive_weight = weights[positive].sum() + negative_weight = weights[negative].sum() + if positive_weight <= 0 or negative_weight <= 0: + raise ValueError("ROC AUC requires positive and negative samples") + pair_scores = scores[positive, None] - scores[None, negative] + pair_value = (pair_scores > 0).to(torch.float64) + 0.5 * (pair_scores == 0) + pair_weights = weights[positive, None] * weights[None, negative] + return float( + (pair_value * pair_weights).sum() / (positive_weight * negative_weight) + ) diff --git a/src/gnn4colliders/tasks/binary_classification.py b/src/gnn4colliders/tasks/binary_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..00159fb0331e8290110866170a77adbf984d9e6d --- /dev/null +++ b/src/gnn4colliders/tasks/binary_classification.py @@ -0,0 +1,78 @@ +"""Binary classification task semantics.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from ._common import ( + batch_labels, + batch_weights, + per_label_weighted_mean, + positive_weight_mask, + weighted_binary_auc, +) + + +class BinaryClassificationTask: + """Interpret one-logit model outputs as a weighted binary task.""" + + def __init__(self, *, absolute_weights: bool = False, threshold: float = 0.5): + self.absolute_weights = absolute_weights + self.threshold = threshold + + @staticmethod + def _logits(logits: torch.Tensor) -> torch.Tensor: + if logits.ndim == 2 and logits.shape[1] == 1: + return logits[:, 0] + if logits.ndim == 1: + return logits + raise ValueError("binary logits must have shape [B] or [B, 1]") + + @staticmethod + def _targets(batch: Any, *, device: torch.device) -> torch.Tensor: + labels = torch.as_tensor(batch_labels(batch), device=device) + if labels.ndim == 2 and labels.shape[1] == 1: + labels = labels[:, 0] + elif labels.ndim != 1: + raise ValueError("binary labels must have shape [B] or [B, 1]") + return labels.to(dtype=torch.float32) + + def loss(self, logits: torch.Tensor, batch: Any) -> torch.Tensor: + raw_logits = self._logits(logits) + labels = self._targets(batch, device=raw_logits.device) + if labels.shape != raw_logits.shape: + raise ValueError("binary logits and labels must have the same batch size") + weights = batch_weights(batch, device=raw_logits.device, dtype=raw_logits.dtype) + if weights.shape != raw_logits.shape: + raise ValueError("metadata.weight must contain one value per event") + if self.absolute_weights: + weights = weights.abs() + elementwise = torch.nn.functional.binary_cross_entropy_with_logits( + raw_logits, labels.to(dtype=raw_logits.dtype), reduction="none" + ) + return per_label_weighted_mean(elementwise, labels, weights) + + def predictions(self, logits: torch.Tensor) -> dict[str, torch.Tensor]: + scores = torch.sigmoid(self._logits(logits)) + return {"scores": scores, "predictions": scores >= self.threshold} + + predict = predictions + + def metrics(self, logits: torch.Tensor, batch: Any) -> dict[str, float]: + scores = self.predictions(logits)["scores"] + labels = self._targets(batch, device=scores.device) + weights = batch_weights(batch, device=scores.device, dtype=scores.dtype) + metric_weights = weights.abs() if self.absolute_weights else weights + mask = positive_weight_mask(weights) + try: + auc = weighted_binary_auc( + labels[mask] == 1, scores[mask], metric_weights[mask] + ) + except ValueError: + auc = float("nan") + accuracy = ( + (self.predictions(logits)["predictions"] == (labels == 1)).float().mean() + ) + return {"accuracy": float(accuracy), "roc_auc": float(auc), "auc": float(auc)} diff --git a/src/gnn4colliders/tasks/multiclass_classification.py b/src/gnn4colliders/tasks/multiclass_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..5554cac9dbf3bbab68930bf22b78f9ccaba0cd32 --- /dev/null +++ b/src/gnn4colliders/tasks/multiclass_classification.py @@ -0,0 +1,80 @@ +"""Multiclass classification task semantics.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from ._common import ( + batch_labels, + batch_weights, + per_label_weighted_mean, + positive_weight_mask, + weighted_binary_auc, +) + + +class MulticlassClassificationTask: + """Interpret ``[B, C]`` logits with legacy per-class weighting.""" + + def __init__(self, *, absolute_weights: bool = False): + self.absolute_weights = absolute_weights + + @staticmethod + def _targets(batch: Any, *, device: torch.device) -> torch.Tensor: + labels = torch.as_tensor(batch_labels(batch), device=device) + if labels.ndim == 2 and labels.shape[1] == 1: + labels = labels[:, 0] + if labels.ndim != 1: + raise ValueError("multiclass labels must have shape [B] or [B, 1]") + return labels.to(dtype=torch.long) + + def loss(self, logits: torch.Tensor, batch: Any) -> torch.Tensor: + if logits.ndim != 2 or logits.shape[1] < 2: + raise ValueError("multiclass logits must have shape [B, C], with C >= 2") + labels = self._targets(batch, device=logits.device) + if labels.shape[0] != logits.shape[0]: + raise ValueError( + "multiclass logits and labels must have the same batch size" + ) + weights = batch_weights(batch, device=logits.device, dtype=logits.dtype) + if weights.shape != labels.shape: + raise ValueError("metadata.weight must contain one value per event") + if self.absolute_weights: + weights = weights.abs() + elementwise = torch.nn.functional.cross_entropy( + logits, labels, reduction="none" + ) + return per_label_weighted_mean(elementwise, labels, weights) + + @staticmethod + def predictions(logits: torch.Tensor) -> dict[str, torch.Tensor]: + if logits.ndim != 2 or logits.shape[1] < 2: + raise ValueError("multiclass logits must have shape [B, C], with C >= 2") + scores = torch.softmax(logits, dim=1) + return {"scores": scores, "predictions": scores.argmax(dim=1)} + + predict = predictions + + def metrics(self, logits: torch.Tensor, batch: Any) -> dict[str, float]: + output = self.predictions(logits) + labels = self._targets(batch, device=logits.device) + weights = batch_weights(batch, device=logits.device, dtype=logits.dtype) + metric_weights = weights.abs() if self.absolute_weights else weights + mask = positive_weight_mask(weights) + accuracy = (output["predictions"] == labels).float().mean() + try: + class_aucs = [] + for class_index in range(output["scores"].shape[1]): + class_aucs.append( + weighted_binary_auc( + labels[mask] == class_index, + output["scores"][mask, class_index], + metric_weights[mask], + ) + ) + auc = sum(class_aucs) / len(class_aucs) + except ValueError: + auc = float("nan") + return {"accuracy": float(accuracy), "roc_auc": float(auc), "auc": float(auc)} diff --git a/src/gnn4colliders/training/__init__.py b/src/gnn4colliders/training/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..2c3fac43525c6d9767f871db0b4738005d887ebb --- /dev/null +++ b/src/gnn4colliders/training/__init__.py @@ -0,0 +1,37 @@ +"""Architecture-independent training infrastructure.""" + +from .checkpoint import ( + CHECKPOINT_SCHEMA_VERSION, + CheckpointManager, + build_model_from_checkpoint, + load_legacy_checkpoint, + load_model_weights, + load_pretrained_edge_network, + restore_training_state, +) +from .early_stopping import EarlyStopping +from .optim import build_optimizer, trainable_parameters +from .reproducibility import seed_everything +from .schedulers import build_scheduler +from .state import BatchResult, EpochResult, TrainerState, TrainingHistory +from .trainer import Trainer + +__all__ = [ + "BatchResult", + "EarlyStopping", + "EpochResult", + "Trainer", + "TrainerState", + "TrainingHistory", + "build_optimizer", + "build_scheduler", + "seed_everything", + "trainable_parameters", + "CHECKPOINT_SCHEMA_VERSION", + "CheckpointManager", + "build_model_from_checkpoint", + "load_legacy_checkpoint", + "load_model_weights", + "load_pretrained_edge_network", + "restore_training_state", +] diff --git a/src/gnn4colliders/training/checkpoint.py b/src/gnn4colliders/training/checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..3b70b6c3a12a0a9f377722a7336a583487984472 --- /dev/null +++ b/src/gnn4colliders/training/checkpoint.py @@ -0,0 +1,328 @@ +"""Small, explicit checkpoint persistence and compatibility adapters. + +Checkpoint files are trusted PyTorch artifacts. The payload itself contains +only state dictionaries and primitive configuration/metadata, which keeps the +new format independent of live trainer or task objects. +""" + +from __future__ import annotations + +import os +import random +import re +import socket +import subprocess +import tempfile +from collections.abc import Mapping +from datetime import datetime, timezone +from pathlib import Path +from typing import Any, Callable + +import numpy as np +import torch +from torch import nn + +from gnn4colliders.compat.checkpoint import ( + load_legacy_checkpoint, + map_legacy_edge_network_state_dict, + normalize_legacy_state_dict_keys, +) +from gnn4colliders.data.metadata import FEATURE_SCHEMA_VERSION, GRAPH_SCHEMA_VERSION + +from .state import TrainerState + +CHECKPOINT_SCHEMA_VERSION = 1 +_EPOCH_RE = re.compile(r"(?:epoch|model_epoch)[_-](\d+)") + + +def _primitive(value: Any) -> Any: + if value is None or isinstance(value, (str, int, float, bool)): + return value + if isinstance(value, Mapping): + return {str(key): _primitive(item) for key, item in value.items()} + if isinstance(value, (list, tuple)): + return [_primitive(item) for item in value] + if isinstance(value, np.generic): + return value.item() + raise TypeError(f"checkpoint configuration must be primitive, got {type(value)!r}") + + +def _state_dict(model: nn.Module) -> dict[str, Any]: + # Compiled and DDP models expose the useful state on the wrapped module. + target = getattr(model, "_orig_mod", model) + state = target.state_dict() + return { + key.removeprefix("module."): value.detach().cpu() + for key, value in state.items() + } + + +def _rng_state() -> dict[str, Any]: + state: dict[str, Any] = { + "python": random.getstate(), + "numpy": np.random.get_state(), + "torch": torch.get_rng_state(), + } + if torch.cuda.is_available(): + state["torch_cuda"] = torch.cuda.get_rng_state_all() + return state + + +def _restore_rng(state: Mapping[str, Any]) -> None: + if "python" in state: + random.setstate(state["python"]) + if "numpy" in state: + np.random.set_state(state["numpy"]) + if "torch" in state: + torch.set_rng_state(state["torch"]) + if "torch_cuda" in state and torch.cuda.is_available(): + torch.cuda.set_rng_state_all(state["torch_cuda"]) + + +def _git_commit() -> str | None: + try: + result = subprocess.run( + ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True + ) + except (OSError, subprocess.CalledProcessError): + return None + return result.stdout.strip() or None + + +def _model_config(model: nn.Module) -> dict[str, Any]: + config_method = getattr(model, "checkpoint_config", None) + if callable(config_method): + return _primitive(config_method()) + return {"class": f"{type(model).__module__}.{type(model).__qualname__}"} + + +def _task_config(task: Any) -> dict[str, Any]: + values = {"type": type(task).__name__} + for name in ("absolute_weights", "threshold", "num_classes"): + if hasattr(task, name): + values[name] = _primitive(getattr(task, name)) + return values + + +def load_model_weights(model: nn.Module, checkpoint: Mapping[str, Any]) -> None: + """Load only model weights; optimizer and lifecycle state are untouched.""" + state = checkpoint.get("model_state_dict", checkpoint) + if not isinstance(state, Mapping): + raise ValueError("checkpoint model_state_dict must be a mapping") + normalized = normalize_legacy_state_dict_keys(state) + # Task 8 owns the historical classify -> classifier rename. Use its one + # compatibility path when applicable, while preserving generic models. + if type(model).__module__.startswith("gnn4colliders.models.root_gnn"): + normalized = map_legacy_edge_network_state_dict(normalized) + try: + model.load_state_dict(normalized, strict=True) + except RuntimeError as error: + raise ValueError( + f"checkpoint weights are incompatible with {type(model).__name__}" + ) from error + + +class CheckpointManager: + """Save and discover versioned checkpoints in a caller-selected directory.""" + + def __init__(self, directory: str | os.PathLike[str]) -> None: + self.directory = Path(directory) + self.directory.mkdir(parents=True, exist_ok=True) + + def save( + self, + *, + model: nn.Module, + trainer_state: TrainerState, + optimizer: torch.optim.Optimizer | None = None, + scheduler: Any | None = None, + early_stopping: Any | None = None, + model_config: Mapping[str, Any] | None = None, + task_config: Mapping[str, Any] | None = None, + metadata: Mapping[str, Any] | None = None, + name: str | None = None, + monitor_name: str | None = None, + monitor_value: float | None = None, + include_rng: bool = True, + ) -> Path: + epoch = int(trainer_state.epoch) + payload: dict[str, Any] = { + "schema_version": CHECKPOINT_SCHEMA_VERSION, + "epoch": epoch, + "global_step": int(trainer_state.global_step), + "model_state_dict": _state_dict(model), + "optimizer_state_dict": optimizer.state_dict() + if optimizer is not None + else None, + "scheduler_state_dict": scheduler.state_dict() + if scheduler is not None + else None, + "early_stopping_state": early_stopping.state_dict() + if early_stopping is not None + else None, + "trainer_state": { + "epoch": epoch, + "global_step": int(trainer_state.global_step), + }, + "model_config": _primitive(model_config) + if model_config is not None + else _model_config(model), + "task_config": _primitive(task_config) if task_config is not None else None, + "metadata": { + "created_at": datetime.now(timezone.utc).isoformat(), + "hostname": socket.gethostname(), + "git_commit": _git_commit(), + "feature_schema_version": FEATURE_SCHEMA_VERSION, + "graph_schema_version": GRAPH_SCHEMA_VERSION, + **(_primitive(metadata) if metadata is not None else {}), + }, + "monitor_name": monitor_name, + "monitor_value": None if monitor_value is None else float(monitor_value), + } + if include_rng: + payload["rng_state"] = _rng_state() + filename = name or f"epoch_{epoch:04d}.pt" + target = self.directory / filename + with tempfile.NamedTemporaryFile( + dir=self.directory, prefix=f".{filename}.", suffix=".tmp", delete=False + ) as handle: + temporary = Path(handle.name) + try: + torch.save(payload, temporary) + os.replace(temporary, target) + finally: + temporary.unlink(missing_ok=True) + return target + + @staticmethod + def _epoch(path: Path) -> int | None: + match = _EPOCH_RE.search(path.stem) + return int(match.group(1)) if match else None + + def checkpoints(self) -> list[Path]: + return sorted( + ( + path + for path in self.directory.glob("*.pt") + if self._epoch(path) is not None + ), + key=lambda path: self._epoch(path), + ) + + def latest(self) -> Path: + paths = self.checkpoints() + if not paths: + raise FileNotFoundError(f"no epoch checkpoints found in {self.directory}") + return paths[-1] + + def best(self, *, mode: str = "min", monitor: str | None = None) -> Path: + if mode not in {"min", "max"}: + raise ValueError("mode must be 'min' or 'max'") + candidates: list[tuple[float, int, Path]] = [] + for path in self.checkpoints(): + payload = self.load(path) + if monitor is not None and payload.get("monitor_name") != monitor: + continue + value = payload.get("monitor_value") + if value is not None: + candidates.append((float(value), self._epoch(path) or -1, path)) + if not candidates: + raise FileNotFoundError( + "no checkpoints contain the requested monitor value" + ) + return (min if mode == "min" else max)( + candidates, key=lambda item: (item[0], item[1]) + )[2] + + @staticmethod + def load( + path: str | os.PathLike[str], *, map_location: Any = "cpu" + ) -> dict[str, Any]: + payload = torch.load(path, map_location=map_location, weights_only=False) + if not isinstance(payload, Mapping): + raise ValueError("checkpoint payload must be a mapping") + if "schema_version" not in payload: + return load_legacy_checkpoint(payload) + version = payload["schema_version"] + if version != CHECKPOINT_SCHEMA_VERSION: + raise ValueError(f"unsupported checkpoint schema version: {version!r}") + return dict(payload) + + +def restore_training_state( + checkpoint: Mapping[str, Any], + *, + model: nn.Module, + trainer: Any | None = None, + optimizer: torch.optim.Optimizer | None = None, + scheduler: Any | None = None, + early_stopping: Any | None = None, + restore_rng: bool = True, +) -> TrainerState: + """Restore explicitly selected resume components and return trainer state.""" + load_model_weights(model, checkpoint) + if optimizer is not None and checkpoint.get("optimizer_state_dict") is not None: + optimizer.load_state_dict(checkpoint["optimizer_state_dict"]) + if scheduler is not None and checkpoint.get("scheduler_state_dict") is not None: + scheduler.load_state_dict(checkpoint["scheduler_state_dict"]) + if ( + early_stopping is not None + and checkpoint.get("early_stopping_state") is not None + ): + early_stopping.load_state_dict(checkpoint["early_stopping_state"]) + state_data = checkpoint.get("trainer_state", checkpoint) + state = TrainerState( + epoch=int(state_data.get("epoch", -1)), + global_step=int(state_data.get("global_step", 0)), + ) + if trainer is not None: + trainer.state = state + if restore_rng and checkpoint.get("rng_state") is not None: + _restore_rng(checkpoint["rng_state"]) + return state + + +def build_model_from_checkpoint( + checkpoint: Mapping[str, Any], factory: Callable[[Mapping[str, Any]], nn.Module] +) -> nn.Module: + """Build a model through an explicit caller-owned factory.""" + config = checkpoint.get("model_config") + if not isinstance(config, Mapping): + raise ValueError("checkpoint does not contain reconstructable model_config") + model = factory(config) + load_model_weights(model, checkpoint) + return model + + +def load_pretrained_edge_network( + checkpoint: Mapping[str, Any] | str | os.PathLike[str], + *, + sample_graph: Any, + sample_global: torch.Tensor | None = None, + map_location: Any = "cpu", + **overrides: Any, +) -> nn.Module: + """Construct and load a new ROOT-GNN backbone for transfer learning.""" + from gnn4colliders.models.root_gnn import EdgeNetwork + + payload = ( + CheckpointManager.load(checkpoint, map_location=map_location) + if not isinstance(checkpoint, Mapping) + else checkpoint + ) + config = dict(payload.get("model_config") or {}) + config.update(overrides) + config.pop("class", None) + required = {"hid_size", "out_size", "n_layers", "n_proc_steps"} + missing = required - config.keys() + if missing: + raise ValueError( + f"checkpoint lacks EdgeNetwork configuration: {sorted(missing)}" + ) + model = EdgeNetwork( + sample_graph, + sample_global, + **{key: config[key] for key in required | {"dropout"} if key in config}, + ) + load_model_weights(model, payload) + return model diff --git a/src/gnn4colliders/training/early_stopping.py b/src/gnn4colliders/training/early_stopping.py new file mode 100644 index 0000000000000000000000000000000000000000..f6703beef38ff0d0b2c05fe0b2860c4850b805bf --- /dev/null +++ b/src/gnn4colliders/training/early_stopping.py @@ -0,0 +1,72 @@ +"""In-memory early-stopping policy.""" + +from __future__ import annotations + +import math + + +class EarlyStopping: + """Stop after ``patience`` consecutive non-improvements. + + Equality is a non-improvement, matching the legacy ``EarlyStop`` policy. + """ + + def __init__( + self, + *, + monitor: str = "loss", + mode: str = "min", + patience: int = 15, + min_delta: float = 1e-8, + ) -> None: + if mode not in {"min", "max"}: + raise ValueError("mode must be 'min' or 'max'") + if patience < 1: + raise ValueError("patience must be positive") + if min_delta < 0: + raise ValueError("min_delta must be non-negative") + self.monitor, self.mode, self.patience, self.min_delta = ( + monitor, + mode, + patience, + min_delta, + ) + self.best = math.inf if mode == "min" else -math.inf + self.num_bad_epochs = 0 + self.should_stop = False + + def update(self, value: float) -> bool: + value = float(value) + improved = ( + value < self.best - self.min_delta + if self.mode == "min" + else value > self.best + self.min_delta + ) + if improved: + self.best = value + self.num_bad_epochs = 0 + else: + self.num_bad_epochs += 1 + self.should_stop = self.num_bad_epochs >= self.patience + return self.should_stop + + def state_dict(self) -> dict[str, object]: + """Return the primitive state needed to resume this policy.""" + return { + "monitor": self.monitor, + "mode": self.mode, + "patience": self.patience, + "min_delta": self.min_delta, + "best": self.best, + "num_bad_epochs": self.num_bad_epochs, + "should_stop": self.should_stop, + } + + def load_state_dict(self, state: dict[str, object]) -> None: + """Restore policy state, rejecting incompatible policy settings.""" + for name in ("monitor", "mode", "patience", "min_delta"): + if name in state and getattr(self, name) != state[name]: + raise ValueError(f"early-stopping {name} does not match checkpoint") + self.best = float(state["best"]) + self.num_bad_epochs = int(state["num_bad_epochs"]) + self.should_stop = bool(state["should_stop"]) diff --git a/src/gnn4colliders/training/optim.py b/src/gnn4colliders/training/optim.py new file mode 100644 index 0000000000000000000000000000000000000000..7f2b60e9c26851488f6f8a83309202b74d117410 --- /dev/null +++ b/src/gnn4colliders/training/optim.py @@ -0,0 +1,38 @@ +"""Explicit construction helpers for the active optimizer workflow.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import torch +from torch import nn + + +def trainable_parameters(model: nn.Module) -> Iterable[nn.Parameter]: + """Yield only parameters enabled for optimization.""" + + return (parameter for parameter in model.parameters() if parameter.requires_grad) + + +def build_optimizer( + parameters: Iterable[nn.Parameter] | nn.Module, + *, + name: str = "adam", + learning_rate: float = 1e-4, + weight_decay: float = 0.0, + **kwargs: object, +) -> torch.optim.Optimizer: + """Build an optimizer used by standard ROOT-GNN configurations.""" + + if isinstance(parameters, nn.Module): + parameters = trainable_parameters(parameters) + normalized = name.lower().replace("_", "") + if normalized == "adam": + return torch.optim.Adam( + parameters, lr=learning_rate, weight_decay=weight_decay, **kwargs + ) + if normalized == "adamw": + return torch.optim.AdamW( + parameters, lr=learning_rate, weight_decay=weight_decay, **kwargs + ) + raise ValueError(f"unsupported optimizer {name!r}; use 'adam' or 'adamw'") diff --git a/src/gnn4colliders/training/reproducibility.py b/src/gnn4colliders/training/reproducibility.py new file mode 100644 index 0000000000000000000000000000000000000000..8234af4fdcae4bdce7696360b517b56aa51d665e --- /dev/null +++ b/src/gnn4colliders/training/reproducibility.py @@ -0,0 +1,20 @@ +"""Explicit process-level reproducibility configuration.""" + +from __future__ import annotations + +import random + +import numpy as np +import torch + + +def seed_everything(seed: int, *, deterministic: bool = False) -> None: + """Seed Python, NumPy, and Torch; optionally request strict Torch behavior.""" + + random.seed(seed) + np.random.seed(seed) + torch.manual_seed(seed) + if torch.cuda.is_available(): + torch.cuda.manual_seed_all(seed) + if deterministic: + torch.use_deterministic_algorithms(True) diff --git a/src/gnn4colliders/training/schedulers.py b/src/gnn4colliders/training/schedulers.py new file mode 100644 index 0000000000000000000000000000000000000000..228b309d904525e706e2f7619ab68bef3818f5ca --- /dev/null +++ b/src/gnn4colliders/training/schedulers.py @@ -0,0 +1,22 @@ +"""Scheduler construction with explicit epoch/metric stepping.""" + +from __future__ import annotations + +import torch + + +def build_scheduler( + optimizer: torch.optim.Optimizer, + *, + name: str = "exponential", + gamma: float = 1.0, + **kwargs: object, +) -> torch.optim.lr_scheduler.LRScheduler | torch.optim.lr_scheduler.ReduceLROnPlateau: + """Build the active epoch scheduler or a validation-driven scheduler.""" + + normalized = name.lower().replace("_", "") + if normalized in {"exponential", "exponentiallr"}: + return torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=gamma) + if normalized in {"reducelronplateau", "plateau"}: + return torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, **kwargs) + raise ValueError(f"unsupported scheduler {name!r}") diff --git a/src/gnn4colliders/training/state.py b/src/gnn4colliders/training/state.py new file mode 100644 index 0000000000000000000000000000000000000000..e0a507bc2359bd15e149747522d0d1f12558636a --- /dev/null +++ b/src/gnn4colliders/training/state.py @@ -0,0 +1,31 @@ +"""Small in-memory types used by the training lifecycle.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Mapping + + +@dataclass(frozen=True) +class BatchResult: + loss: float + logits: object + + +@dataclass(frozen=True) +class EpochResult: + loss: float + metrics: Mapping[str, float] = field(default_factory=dict) + num_samples: int = 0 + + +@dataclass +class TrainingHistory: + train: list[EpochResult] = field(default_factory=list) + validation: list[EpochResult] = field(default_factory=list) + + +@dataclass +class TrainerState: + epoch: int = -1 + global_step: int = 0 diff --git a/src/gnn4colliders/training/trainer.py b/src/gnn4colliders/training/trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..352eb060a0559fa130b925f5b920f16e1b9fdf5c --- /dev/null +++ b/src/gnn4colliders/training/trainer.py @@ -0,0 +1,272 @@ +"""Single-process, architecture-independent model training lifecycle.""" + +from __future__ import annotations + +import inspect +import logging +from types import SimpleNamespace +from typing import Any + +import torch +from torch import nn + +from gnn4colliders.data.graph_dataset import GraphBatch +from gnn4colliders.distributed import ( + DistributedContext, + broadcast_bool, + gather_objects, + gather_tensor, + prepare_model, +) + +from .early_stopping import EarlyStopping +from .state import BatchResult, EpochResult, TrainerState, TrainingHistory + +logger = logging.getLogger(__name__) + + +def _model_forward(model: nn.Module, batch: GraphBatch) -> torch.Tensor: + """Call both supported model boundaries without masking model errors.""" + + try: + signature = inspect.signature(model.forward) + positional = [ + parameter + for parameter in signature.parameters.values() + if parameter.kind + in (parameter.POSITIONAL_ONLY, parameter.POSITIONAL_OR_KEYWORD) + ] + except (TypeError, ValueError): + positional = [] + if len(positional) <= 1: + return model(batch) + return model(batch.graph, batch.global_features) + + +def _combined_batch(batches: list[GraphBatch], logits: torch.Tensor) -> Any: + labels = torch.cat([batch.labels.reshape(-1) for batch in batches]) + weights = torch.cat([batch.metadata.weight.reshape(-1) for batch in batches]) + sample_ids = tuple( + sample_id for batch in batches for sample_id in batch.metadata.sample_id + ) + return SimpleNamespace( + labels=labels, + metadata=SimpleNamespace(weight=weights, sample_id=sample_ids), + logits=logits, + ) + + +class Trainer: + """Orchestrate model, task, optimizer, and optional scheduler.""" + + def __init__( + self, + model: nn.Module, + task: Any, + optimizer: torch.optim.Optimizer, + scheduler: Any | None = None, + *, + device: torch.device | str = "cpu", + early_stopping: EarlyStopping | None = None, + scheduler_step: str = "epoch", + distributed_context: DistributedContext | None = None, + ) -> None: + if scheduler_step not in {"epoch", "validation"}: + raise ValueError("scheduler_step must be 'epoch' or 'validation'") + self.distributed_context = ( + distributed_context or DistributedContext.single_process(device) + ) + self.device = self.distributed_context.device + self.model = prepare_model(model, self.distributed_context) + self.task = task + self.optimizer = optimizer + self.scheduler = scheduler + self.early_stopping = early_stopping + self.scheduler_step = scheduler_step + self.state = TrainerState() + + def _train_moved_batch(self, batch: GraphBatch) -> BatchResult: + self.optimizer.zero_grad(set_to_none=True) + logits = _model_forward(self.model, batch) + loss = self.task.loss(logits, batch) + if not torch.isfinite(loss): + raise FloatingPointError("non-finite training loss") + loss.backward() + self.optimizer.step() + self.state.global_step += 1 + return BatchResult(loss=float(loss.detach().cpu()), logits=logits.detach()) + + def train_batch(self, batch: GraphBatch) -> BatchResult: + """Run one optimizer step, moving the supplied batch to the device.""" + + return self._train_moved_batch(batch.to(self.device)) + + def train_epoch(self, loader: Any) -> EpochResult: + self.model.train() + losses: list[float] = [] + batches: list[GraphBatch] = [] + outputs: list[torch.Tensor] = [] + num_samples = 0 + for batch in loader: + moved = batch.to(self.device) + result = self._train_moved_batch(moved) + losses.append(result.loss) + batches.append(moved) + outputs.append(result.logits) + num_samples += int(moved.labels.shape[0]) + if not losses: + raise ValueError("cannot train on an empty loader") + logits = torch.cat(outputs, dim=0) + combined = _combined_batch(batches, logits) + if self.distributed_context.enabled: + combined = self._gather_combined(combined) + # Compute the reported loss over the global event set. Gradients + # still came from local losses and were synchronized by DDP. + global_loss = float(self.task.loss(combined.logits, combined).cpu()) + else: + global_loss = sum(losses) / len(losses) + metrics = { + key: float(value) + for key, value in self.task.metrics(logits, combined).items() + } + # Legacy training reports the arithmetic mean of batch losses, while + # metrics are evaluated over every event in the complete epoch. + return EpochResult( + loss=global_loss, metrics=metrics, num_samples=int(combined.labels.shape[0]) + ) + + def evaluate( + self, loader: Any, *, return_outputs: bool = False + ) -> EpochResult | Any: + self.model.eval() + batches: list[GraphBatch] = [] + outputs: list[torch.Tensor] = [] + with torch.inference_mode(): + for batch in loader: + moved = batch.to(self.device) + outputs.append(_model_forward(self.model, moved).detach()) + batches.append(moved) + if not outputs: + raise ValueError("cannot evaluate an empty loader") + logits = torch.cat(outputs, dim=0) + combined = _combined_batch(batches, logits) + if self.distributed_context.enabled: + combined = self._gather_combined(combined) + loss = self.task.loss(logits, combined) + metrics = { + key: float(value) + for key, value in self.task.metrics(logits, combined).items() + } + result = EpochResult( + loss=float(loss.cpu()), + metrics=metrics, + num_samples=int(logits.shape[0]), + ) + if return_outputs: + return ( + result, + logits.cpu(), + combined.labels.cpu(), + combined.metadata.weight.cpu(), + combined.metadata.sample_id, + ) + return result + + def _gather_combined(self, combined: Any) -> Any: + logits = gather_tensor(combined.logits, self.distributed_context) + labels = gather_tensor(combined.labels, self.distributed_context) + weights = gather_tensor(combined.metadata.weight, self.distributed_context) + ids = tuple( + item + for group in gather_objects( + combined.metadata.sample_id, self.distributed_context + ) + for item in group + ) + return SimpleNamespace( + labels=labels, + metadata=SimpleNamespace(weight=weights, sample_id=ids), + logits=logits, + ) + + def fit( + self, + train_loader: Any, + validation_loader: Any | None = None, + *, + epochs: int, + ) -> TrainingHistory: + """Train using validation data for every-epoch model selection. + + The held-out test split is intentionally absent from this method and + should be passed to :meth:`evaluate` only after model selection is + complete. This avoids the legacy train/test/val naming inversion. + """ + + if epochs < 1: + raise ValueError("epochs must be positive") + if self.early_stopping is not None and validation_loader is None: + raise ValueError("early stopping requires a validation loader") + history = TrainingHistory() + # ``state.epoch`` is the last completed epoch. This makes a saved + # state unambiguous: a resumed fit starts at ``epoch + 1``. + start_epoch = self.state.epoch + 1 + for epoch in range(start_epoch, start_epoch + epochs): + if hasattr(train_loader, "set_epoch"): + train_loader.set_epoch(epoch) + if validation_loader is not None and hasattr( + validation_loader, "set_epoch" + ): + validation_loader.set_epoch(epoch) + train_result = self.train_epoch(train_loader) + history.train.append(train_result) + validation_result = None + if validation_loader is not None: + validation_result = self.evaluate(validation_loader) + history.validation.append(validation_result) + if self.distributed_context.is_main_process: + logger.info("epoch=%d train_loss=%s", epoch, train_result.loss) + if ( + validation_result is not None + and self.distributed_context.is_main_process + ): + logger.info( + "epoch=%d validation_loss=%s metrics=%s", + epoch, + validation_result.loss, + validation_result.metrics, + ) + # Record the completed epoch before policy logic can terminate + # the loop, so a checkpoint after early stopping is resumable. + self.state.epoch = epoch + if self.early_stopping is not None: + if self.early_stopping.monitor == "loss": + monitored = validation_result.loss + else: + if validation_result is None: + raise ValueError("early stopping monitor requires validation") + try: + monitored = validation_result.metrics[ + self.early_stopping.monitor + ] + except KeyError as error: + raise KeyError( + f"metric {self.early_stopping.monitor!r} was not produced" + ) from error + stop = ( + self.early_stopping.update(monitored) + if self.distributed_context.is_main_process + else False + ) + if broadcast_bool(stop, self.distributed_context): + break + if self.scheduler is not None: + if self.scheduler_step == "validation": + if validation_result is None: + raise ValueError( + "validation scheduler requires a validation loader" + ) + self.scheduler.step(validation_result.loss) + else: + self.scheduler.step() + return history diff --git a/tasks/task1.md b/tasks/task1.md new file mode 100644 index 0000000000000000000000000000000000000000..2d10ecff7a143a53437026d62f1f52a73dabea8b --- /dev/null +++ b/tasks/task1.md @@ -0,0 +1,389 @@ +# Task: Initialize the GNN4Colliders project structure + +We are rewriting an existing collider-ML repository from scratch. + +The repository is named **GNN4Colliders**. + +The installable Python package should be: + +```text +gnn4colliders +``` + +The initial model family being rewritten is called: + +```text +root_gnn +``` + +In the future, the project may contain other architectures such as: + +```text +root_transformer +``` + +The goal is therefore to build a reusable collider-ML package rather than make the entire package GNN-specific. + +## Important context + +Read these files before making changes: + +```text +docs/architecture.md +``` + +and any existing repository-level documentation. + +The legacy implementation should be treated as a behavioral reference only. + +If a `legacy/` directory exists: + +* do not modify it +* do not reorganize it +* do not copy its architecture blindly +* do not begin migrating implementation code in this task + +This task is only about establishing the new project structure and documentation. + +--- + +# Desired architecture + +Create this initial structure: + +```text +GNN4Colliders/ +├── README.md +├── AGENTS.md +├── pyproject.toml +│ +├── configs/ +│ ├── data/ +│ ├── model/ +│ │ └── root_gnn/ +│ ├── task/ +│ ├── trainer/ +│ └── environment/ +│ +├── src/ +│ └── gnn4colliders/ +│ ├── __init__.py +│ │ +│ ├── data/ +│ │ └── __init__.py +│ │ +│ ├── features/ +│ │ └── __init__.py +│ │ +│ ├── graphs/ +│ │ └── __init__.py +│ │ +│ ├── models/ +│ │ ├── __init__.py +│ │ └── root_gnn/ +│ │ └── __init__.py +│ │ +│ ├── training/ +│ │ └── __init__.py +│ │ +│ ├── inference/ +│ │ └── __init__.py +│ │ +│ └── cli/ +│ └── __init__.py +│ +├── tests/ +│ ├── unit/ +│ ├── integration/ +│ ├── parity/ +│ └── fixtures/ +│ +├── notebooks/ +│ +├── scripts/ +│ ├── slurm/ +│ └── dev/ +│ +├── docs/ +│ ├── architecture.md +│ └── migration.md +│ +└── legacy/ +``` + +Do not create unnecessary placeholder Python modules yet. + +Empty directories may contain `.gitkeep` files where required. + +--- + +# Architectural intent + +The package should follow these responsibilities. + +## `gnn4colliders.data` + +Generic data access and dataset infrastructure. + +Future responsibilities may include: + +* ROOT/Awkward input +* dataset abstractions +* cache handling +* batching +* folds/splits + +This package should not contain GNN-specific logic. + +## `gnn4colliders.features` + +Physics-domain transformations that can be shared across architectures. + +Future responsibilities may include: + +* collider object features +* selections +* feature scaling +* derived physics quantities + +These should be usable by both GNNs and future transformer-like architectures. + +## `gnn4colliders.graphs` + +Graph-specific representation logic. + +Future responsibilities may include: + +* graph topology +* graph construction +* edge feature construction +* graph padding + +Keep graph representation separate from generic ROOT input where possible. + +## `gnn4colliders.models` + +Architecture-specific ML models. + +Initially: + +```text +gnn4colliders.models.root_gnn +``` + +Future architectures may include: + +```text +gnn4colliders.models.root_transformer +``` + +Shared infrastructure should not be placed inside `root_gnn` unless it is genuinely GNN-specific. + +## `gnn4colliders.training` + +Architecture-independent training infrastructure where practical. + +Future responsibilities may include: + +* training lifecycle +* losses +* metrics +* checkpointing +* distributed training +* reproducibility utilities + +Do not implement these yet. + +## `gnn4colliders.inference` + +Architecture-independent inference/application infrastructure where practical. + +Future responsibilities may include: + +* prediction +* evaluation +* output writers +* model export + +## `gnn4colliders.cli` + +Thin command-line entry points. + +Eventually the project should support commands conceptually similar to: + +```bash +gnn4colliders prepare +gnn4colliders train +gnn4colliders evaluate +gnn4colliders predict +gnn4colliders export +``` + +Do not implement these workflows yet unless minimal CLI scaffolding is necessary for packaging. + +--- + +# Configuration philosophy + +Create the config directory structure, but do not reproduce the legacy dynamic-import configuration system. + +Configuration should eventually describe experiments semantically. + +For example: + +```yaml +model: + type: root_gnn +``` + +rather than exposing Python internals such as: + +```yaml +module: some.python.module +class: SomeClass +``` + +We expect to use composable YAML configuration, likely with Hydra, but this task should not build the configuration system beyond any minimal dependency or documentation decision that is clearly justified. + +--- + +# `pyproject.toml` + +Create a modern minimal `pyproject.toml`. + +Requirements: + +* project name: `gnn4colliders` +* use `src/` layout +* package discovery should find `src/gnn4colliders` +* specify an appropriate modern Python minimum version +* include only dependencies that are clearly required for the initial project skeleton +* development tooling may include: + + * pytest + * ruff + +Do not prematurely add the full legacy dependency environment. + +Do not pin CUDA, PyTorch, DGL, ROOT, or other scientific dependencies until their compatibility strategy is addressed separately. + +If adding a CLI entry point now would require inventing implementation, leave it out and document the intended future CLI instead. + +--- + +# README.md + +Create a useful initial README that explains: + +1. What GNN4Colliders is. +2. That it is a collider-ML toolkit intended to support multiple model families. +3. That `root_gnn` is the first model family being rewritten. +4. That additional architectures such as `root_transformer` may be added later. +5. The high-level package layout. +6. The distinction between: + + * shared collider/data infrastructure + * representation-specific code + * model-specific code + * experiment configuration +7. The intended future CLI/config workflow. +8. Current development status: + + * architecture and migration scaffold + * implementation not yet complete +9. Basic developer setup using an editable install. +10. A short development philosophy emphasizing: + + * testability + * reproducibility + * explicit interfaces + * incremental migration from legacy behavior + +Do not claim features already work when they do not. + +--- + +# AGENTS.md + +Create the repository-level `AGENTS.md` using the project instructions supplied in this task. + +It should establish durable rules for future coding-agent work, including: + +* repository purpose +* architecture boundaries +* legacy code rules +* testing expectations +* reproducibility expectations +* configuration conventions +* code-quality conventions +* incremental migration workflow +* requirement to inspect relevant existing code before changing behavior +* requirement to avoid broad unrelated refactors + +The file should remain concise enough to serve as a practical agent instruction file. + +--- + +# migration.md + +If `docs/migration.md` does not already exist, create a minimal one. + +It should list the intended migration stages without implementing them: + +1. project/package skeleton +2. characterization/parity tests +3. configuration schema +4. ROOT/Awkward I/O +5. physics feature extraction and selections +6. graph construction +7. dataset caching/loading/batching +8. active ROOT-GNN model +9. losses and metrics +10. training lifecycle +11. checkpoints +12. inference +13. export +14. distributed/HPC workflows +15. legacy removal after parity + +If `docs/migration.md` already exists, preserve its content unless a small structural update is clearly necessary. + +--- + +# Constraints + +Do not: + +* rewrite legacy model code +* migrate training logic +* implement datasets +* implement graph construction +* introduce Lightning +* introduce Kedro +* create speculative abstractions +* add unnecessary dependencies +* modify legacy code +* remove existing documentation +* perform broad cleanup unrelated to this task + +Prefer the smallest structure that establishes clear long-term boundaries. + +--- + +# Validation + +After making changes: + +1. Show the resulting directory tree. +2. Verify the package can be discovered/imported if practical. +3. Run any formatter/linter/tests that are available and relevant. +4. Check that `pyproject.toml` is valid. +5. Review README.md and AGENTS.md for claims about functionality that does not exist yet. + +Then report: + +* files created +* files modified +* validation performed +* any decisions you intentionally deferred +* any conflicts you found with the existing architecture documentation diff --git a/tasks/task10.md b/tasks/task10.md new file mode 100644 index 0000000000000000000000000000000000000000..c248168f0dd8b0807c369022419d67d376399bff --- /dev/null +++ b/tasks/task10.md @@ -0,0 +1,1360 @@ +# Task 10: Implement the Training and Evaluation Lifecycle + +Implement the core training/evaluation lifecycle around the model, task, and `GraphBatch` abstractions completed in previous tasks. + +This task builds on: + +```text +Task 4: shared physics feature construction +Task 5: ROOT-GNN graph construction +Task 6: ROOT/Awkward ingestion and EventSample +Task 7: metadata-aware caching, splits, batching, and GraphBatch +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: classification tasks, weighted losses, metrics, and output handling +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/models/ +src/gnn4colliders/tasks/ +src/gnn4colliders/training/ +tests/unit/ +tests/integration/ +tests/parity/ +legacy/root_gnn_dgl/scripts/training_script.py +legacy/root_gnn_dgl/root_gnn_base/custom_scheduler.py +``` + +Focus on the active training behavior required by standard ROOT-GNN workflows. + +Do not implement checkpoint persistence/resume, inference output writing, Hydra/CLI wiring, or DDP in this task. + +--- + +# Goal + +Implement a clean trainer that orchestrates: + +```text +GraphBatch + -> +Model + -> +logits + -> +Task.loss(...) + -> +backward + -> +optimizer step +``` + +and, for evaluation: + +```text +GraphBatch + -> +Model + -> +logits + -> +accumulate outputs + -> +Task.metrics(...) +``` + +The trainer should know how to execute training. + +It should not know collider-specific preprocessing rules, graph-construction rules, task-specific loss mathematics, or legacy tracking-column semantics. + +--- + +# 1. Training package structure + +Prefer a structure such as: + +```text +src/gnn4colliders/training/ + __init__.py + trainer.py + optim.py + schedulers.py + early_stopping.py + state.py +``` + +Use fewer modules if clearer. + +Do not create a large framework. + +Possible responsibilities: + +```text +trainer.py + train/eval loops + +optim.py + optimizer construction + +schedulers.py + active scheduler construction + +early_stopping.py + stopping policy + +state.py + lightweight epoch/history/result types +``` + +Do not implement checkpoint serialization in these modules yet. + +--- + +# 2. Trainer responsibility + +The trainer should orchestrate: + +* device placement +* training mode +* evaluation mode +* zeroing gradients +* forward pass +* task loss +* backward pass +* optimizer step +* scheduler step where appropriate +* epoch-level output accumulation +* epoch-level metrics +* early-stopping decisions +* training history + +The trainer should not own: + +* ROOT file reading +* feature construction +* graph construction +* fold semantics +* loss mathematics +* metric mathematics +* checkpoint serialization +* CLI parsing +* YAML configuration + +--- + +# 3. Trainer API + +Prefer a small explicit API. + +Conceptually: + +```python +trainer = Trainer( + model=model, + task=task, + optimizer=optimizer, + scheduler=scheduler, + device=device, +) + +history = trainer.fit( + train_loader=train_loader, + validation_loader=validation_loader, + epochs=100, +) +``` + +The exact API may differ. + +Keep it easy to instantiate directly from Python. + +Do not wire it to Hydra yet. + +--- + +# 4. Train one batch + +Implement the batch-level training operation clearly. + +Conceptually: + +```python +optimizer.zero_grad() + +logits = model(batch) + +loss = task.loss_from_batch( + logits, + batch, +) + +loss.backward() + +optimizer.step() +``` + +Do not hide unusual side effects. + +Return a lightweight result if useful. + +For example: + +```python +BatchResult( + loss=..., + logits=..., +) +``` + +Avoid retaining unnecessary autograd graphs. + +--- + +# 5. Train one epoch + +Implement a clear method/function for one training epoch. + +Conceptually: + +```python +train_result = trainer.train_epoch(train_loader) +``` + +It should: + +1. call `model.train()` +2. iterate over batches +3. move batches to the configured device +4. compute logits +5. compute task loss +6. backpropagate +7. update optimizer +8. accumulate detached outputs needed for epoch metrics +9. calculate epoch-level metrics after iteration + +Do not compute ROC AUC independently per mini-batch. + +--- + +# 6. Evaluate one epoch + +Implement evaluation separately from training. + +Conceptually: + +```python +eval_result = trainer.evaluate(validation_loader) +``` + +It should: + +* call `model.eval()` +* use `torch.no_grad()` or `torch.inference_mode()` +* never call backward +* never modify optimizer state +* accumulate outputs across the full split +* compute task metrics at split level + +Avoid duplicating substantial train/eval logic where a small shared helper is clearer. + +--- + +# 7. Epoch-level metric accumulation + +Metrics such as ROC AUC should operate over the full split, not independently on each batch. + +Accumulate only the data needed for task-level metrics. + +Conceptually: + +```text +batch logits +batch targets +batch weights +sample IDs + ↓ +detach + ↓ +CPU or compact accumulator + ↓ +end of epoch + ↓ +task.metrics(...) +``` + +Do not retain computation graphs. + +--- + +# 8. Epoch result type + +Introduce a small explicit result type. + +Conceptually: + +```python +@dataclass(frozen=True) +class EpochResult: + loss: float + metrics: Mapping[str, float] + num_samples: int +``` + +Additional fields may be added if useful. + +Do not expose internal accumulator details. + +--- + +# 9. Training history + +Provide a lightweight history representation. + +Conceptually: + +```python +@dataclass +class TrainingHistory: + train: list[EpochResult] + validation: list[EpochResult] +``` + +or equivalent. + +Avoid introducing experiment-tracking infrastructure yet. + +History should be easy to serialize later. + +--- + +# 10. Batch device movement + +Define one explicit way to move `GraphBatch` to a device. + +If Task 7 does not already provide: + +```python +batch.to(device) +``` + +consider adding a small method/helper. + +It should correctly move: + +* DGL graph +* labels +* numeric batched metadata needed by training +* global features + +Strings such as `sample_id` should remain on CPU as ordinary Python data. + +Do not silently discard metadata. + +--- + +# 11. Model device + +The trainer should move the model to the requested device explicitly. + +For example: + +```python +model.to(device) +``` + +Do not hardcode CUDA. + +Tests must run on CPU. + +--- + +# 12. Device configuration + +Support explicit devices such as: + +```text +cpu +cuda +cuda:0 +``` + +Do not silently choose a GPU based on availability unless that is explicitly documented and desired. + +Prefer caller control. + +A convenience `"auto"` option may be added only if it remains obvious and testable. + +--- + +# 13. Optimizer construction + +Implement the optimizer(s) used by the active standard workflow. + +Inspect the legacy configuration and training script to determine the active optimizer. + +Support the smallest useful API. + +Conceptually: + +```python +optimizer = build_optimizer( + model.parameters(), + name="adam", + learning_rate=..., + weight_decay=..., +) +``` + +Do not implement every optimizer supported by PyTorch. + +Only support what the active standard configs require plus perhaps one obvious default. + +--- + +# 14. Fine-tuning optimizer behavior + +Ensure transfer/fine-tuning works correctly with Task 8's freeze/unfreeze behavior. + +If the backbone is frozen: + +```text +optimizer + -> +only trainable parameters +``` + +The optimizer should not accidentally optimize parameters with `requires_grad=False`. + +Prefer: + +```python +parameters = ( + parameter + for parameter in model.parameters() + if parameter.requires_grad +) +``` + +or equivalent. + +Add tests. + +--- + +# 15. Learning-rate scheduler + +Inspect the active legacy scheduler behavior. + +Only implement scheduler behavior required by current standard configs. + +If the legacy active workflow uses a custom scheduler, characterize its actual external behavior before migrating it. + +Prefer standard PyTorch schedulers where they reproduce the required behavior exactly. + +Do not preserve a custom implementation solely because it exists. + +--- + +# 16. Scheduler stepping semantics + +Make scheduler stepping explicit. + +Distinguish between: + +```text +step per batch +step per epoch +step on validation metric +``` + +Do not guess. + +Verify against the active legacy workflow. + +Add tests for the chosen stepping behavior. + +--- + +# 17. No checkpoint persistence yet + +The trainer may expose in-memory state such as: + +```text +current epoch +optimizer +scheduler +early stopping state +``` + +but do not write checkpoint files yet. + +Do not implement: + +```text +model_epoch_N.pt +resume from disk +best checkpoint +last checkpoint +``` + +Those belong to Task 11. + +Design the trainer state so Task 11 can serialize it cleanly later. + +--- + +# 18. Trainer state + +A small state object is acceptable. + +Conceptually: + +```python +@dataclass +class TrainerState: + epoch: int + global_step: int +``` + +It may also track: + +```text +best metric +epochs without improvement +``` + +if early stopping requires it. + +Do not include serialized model tensors inside this object. + +--- + +# 19. Early stopping + +Implement early stopping as a separate policy. + +Conceptually: + +```python +early_stopping = EarlyStopping( + monitor="loss", + mode="min", + patience=10, + min_delta=0.0, +) +``` + +or equivalent. + +The trainer should ask the policy whether to stop. + +Do not mix early-stopping logic throughout the epoch loop. + +--- + +# 20. Early-stopping behavior + +Characterize legacy semantics: + +* metric monitored +* comparison direction +* patience interpretation +* min delta if any +* what happens on equality +* whether the counter resets on improvement + +Preserve externally relevant behavior where appropriate. + +Document intentional improvements. + +--- + +# 21. Best state vs best checkpoint + +The early-stopping policy may track the best metric in memory. + +Do not save the best model to disk in this task. + +Task 11 will connect early-stopping state to checkpoint persistence. + +--- + +# 22. Loss aggregation + +Be explicit about epoch loss reporting. + +Do not simply average batch losses unless that reproduces intended semantics. + +Determine whether reported epoch loss should be: + +```text +mean across batches +mean across samples +weighted aggregate +``` + +based on legacy behavior and Task 9 semantics. + +Implement the active behavior explicitly. + +Add tests with different batch sizes to ensure aggregation semantics are correct. + +--- + +# 23. Metadata usage + +The trainer should only use metadata needed for training/task execution. + +For example: + +```text +weight +``` + +may be used by the task. + +The trainer should not contain code like: + +```python +tracking[:, 1] +``` + +or direct knowledge that: + +```text +fold +sample_id +source_file +``` + +exist unless needed for logging/debugging. + +Keep task/domain semantics outside the trainer. + +--- + +# 24. Sample IDs + +Preserve `sample_id` through batch execution. + +The trainer does not need to use it for loss calculation. + +However, evaluation accumulators may optionally retain sample IDs to support later debugging/inference workflows. + +Do not make this required if it adds unnecessary memory use. + +--- + +# 25. Reproducibility + +Add an explicit seeding utility if one does not already exist. + +Conceptually: + +```python +seed_everything(seed) +``` + +Account for: + +* Python random +* NumPy +* PyTorch CPU +* PyTorch CUDA where available + +Do not mutate seeds from inside model constructors. + +Do not reseed on every epoch unless explicitly required. + +--- + +# 26. DataLoader reproducibility + +Coordinate with Task 7's explicit DataLoader generator/seeding. + +Do not silently override DataLoader seed behavior in the trainer. + +The trainer-level seed should be documented as controlling model/optimization randomness, while loader ordering should remain governed by the loader's own explicit seed/generator. + +Avoid two hidden competing seed systems. + +--- + +# 27. Determinism mode + +Do not promise full bitwise GPU determinism automatically. + +If useful, support an explicit strict-determinism option. + +For example: + +```python +configure_reproducibility( + seed=42, + deterministic=False, +) +``` + +If strict PyTorch deterministic algorithms are enabled, make it explicit. + +Document performance/compatibility tradeoffs. + +--- + +# 28. Gradient handling + +Implement ordinary gradient behavior first. + +If the active legacy path requires: + +```text +gradient clipping +``` + +characterize and implement it explicitly. + +Otherwise defer it. + +Do not add clipping as an arbitrary default. + +--- + +# 29. Gradient accumulation + +Do not implement gradient accumulation unless active configurations require it. + +Keep one optimizer step per training batch by default. + +Document this as deferred if relevant. + +--- + +# 30. Mixed precision + +Do not implement AMP/mixed precision yet. + +Avoid introducing: + +```text +autocast +GradScaler +bf16/fp16 policies +``` + +in this task. + +Mixed precision belongs to later performance/HPC work after correctness is established. + +--- + +# 31. torch.compile + +Do not use: + +```python +torch.compile(...) +``` + +in the baseline trainer. + +Task 15 can benchmark it later. + +The baseline trainer should remain simple and parity-focused. + +--- + +# 32. Evaluation cadence + +Support the simple default: + +```text +train one epoch + -> +evaluate validation split +``` + +Do not implement arbitrary callback scheduling or evaluation every N steps unless needed by active workflows. + +--- + +# 33. Validation loader optionality + +Allow training without a validation loader if useful. + +For example: + +```python +trainer.fit( + train_loader, + validation_loader=None, +) +``` + +If early stopping requires validation, raise a clear error when configured without the required monitored data. + +--- + +# 34. Test-loader semantics + +The legacy implementation may use "test" loaders in places where validation would normally occur. + +Do not reproduce confused naming. + +Use clear new concepts: + +```text +train +validation +test +``` + +Preserve dataset membership semantics where required, but modernize the API naming. + +Document intentional terminology cleanup. + +--- + +# 35. Evaluation result collection + +Expose enough evaluation data for future Task 12 inference/evaluation output. + +Conceptually, evaluation may optionally return: + +```python +EvaluationResult( + epoch=..., + metrics=..., + logits=..., + targets=..., + weights=..., + sample_ids=..., +) +``` + +However, do not force the trainer to retain large arrays during normal training. + +Prefer an option such as: + +```python +return_outputs=False +``` + +or a separate evaluator method. + +Keep memory use predictable. + +--- + +# 36. Logging + +Use the standard Python `logging` module for trainer-level logging. + +Do not add: + +```text +Weights & Biases +MLflow +TensorBoard +``` + +as hard dependencies in this task. + +Log useful events such as: + +```text +epoch +train loss +validation loss +metrics +learning rate +early-stop status +``` + +Keep logging separate from metric computation. + +--- + +# 37. Progress bars + +Do not make `tqdm` a required dependency unless it is already present and useful. + +If added, make progress display optional. + +Correctness should not depend on interactive terminal behavior. + +--- + +# 38. Callback systems + +Do not build a generic callback/plugin framework. + +Task 10 needs a trainer, not a mini-Lightning clone. + +Use direct explicit composition: + +```text +Trainer +Task +Optimizer +Scheduler +EarlyStopping +``` + +If later requirements justify callbacks, they can be added later. + +--- + +# 39. Fit lifecycle + +A typical lifecycle should conceptually be: + +```text +configure model/device +configure optimizer +configure scheduler +configure early stopping + +for epoch: + train_result = train_epoch(...) + + if validation_loader: + validation_result = evaluate(...) + + scheduler.step(...) + + update early stopping + + append history + + if stop: + break +``` + +Verify scheduler/early-stopping ordering against legacy behavior where it matters. + +--- + +# 40. Fine-tuning integration + +Add tests proving the trainer works with: + +```text +FineTunedEdgeNetwork +``` + +in both: + +```text +frozen backbone +unfrozen backbone +``` + +modes. + +For the frozen case, verify only classifier parameters change after one optimizer step. + +For the unfrozen case, verify at least one backbone parameter receives a gradient/updates. + +Use tiny deterministic inputs. + +--- + +# 41. Binary training integration + +Add a small test: + +```text +GraphBatch + -> +FineTunedEdgeNetwork + -> +BinaryClassificationTask + -> +Trainer.train_epoch +``` + +Verify: + +* finite loss +* backward succeeds +* optimizer updates trainable parameters +* metadata weight is honored by task +* no positional tracking access exists + +--- + +# 42. Multiclass training integration + +Add a small test: + +```text +GraphBatch + -> +EdgeNetwork + -> +MulticlassClassificationTask + -> +Trainer.train_epoch +``` + +Verify: + +* finite loss +* backward succeeds +* expected logits/task semantics +* epoch result is produced + +Do not attempt convergence testing. + +--- + +# 43. No full training convergence tests + +Do not require a model to learn a real collider task. + +Integration tests should prove execution semantics, not scientific convergence. + +A tiny synthetic overfit test may be added only if it is fast and genuinely useful. + +Do not make the suite slow. + +--- + +# 44. Legacy parity + +Characterize and compare the active legacy training lifecycle where practical. + +Focus on deterministic behavior such as: + +```text +one batch loss +one optimizer update +scheduler step +epoch loss aggregation +early-stopping update +``` + +Do not expect long multi-epoch GPU training runs to be bit-identical. + +Use fixed seeds and tiny fixtures. + +--- + +# 45. One-step optimizer parity + +Where practical: + +1. initialize legacy/new model with equivalent weights +2. use identical batch +3. use equivalent optimizer parameters +4. execute one forward/backward/update +5. compare updated active model parameters + +Use appropriate floating-point tolerances. + +If exact optimizer parity is impractical due to architecture cleanup, characterize and document why. + +--- + +# 46. Scheduler parity + +If an active scheduler exists, add a deterministic test comparing learning-rate progression across a handful of steps/epochs. + +Prefer explicit expected LR values. + +Do not test through a long training run. + +--- + +# 47. Early-stopping unit tests + +Test sequences such as: + +```text +improving +plateau +degrading +equality +reset after improvement +patience exhausted +``` + +Ensure behavior is deterministic and independent of model execution. + +--- + +# 48. Training-history tests + +Verify: + +* one result per executed epoch +* validation history aligns with train history when enabled +* stopped training records only executed epochs +* numeric outputs are ordinary Python numbers where appropriate + +--- + +# 49. No checkpoint implementation + +Do not save files during `Trainer.fit()` in this task. + +No: + +```text +model_epoch_0.pt +best.pt +last.pt +``` + +Task 11 will add a checkpoint manager around trainer/model/optimizer/scheduler state. + +Keep Task 10 trainer easy to integrate with that later. + +--- + +# 50. No Hydra/CLI wiring + +Do not expose training through: + +```bash +gnn4colliders train +``` + +yet. + +Do not instantiate Trainer from YAML in this task. + +Use direct Python APIs in tests. + +Hydra/CLI wiring belongs to Task 13. + +--- + +# 51. No DDP + +Do not implement: + +```text +DistributedDataParallel +NCCL +DistributedSampler +rank-aware metrics +torchrun +``` + +yet. + +Keep the trainer single-process. + +Design it cleanly enough that distributed execution can wrap it later. + +--- + +# 52. Public exports + +Expose intended training APIs. + +For example: + +```python +from gnn4colliders.training import ( + Trainer, + EarlyStopping, + build_optimizer, +) +``` + +Do not expose internal accumulator helpers unnecessarily. + +--- + +# 53. Documentation + +Update: + +```text +docs/architecture.md +``` + +to clarify: + +```text +DataLoader + -> +GraphBatch + -> +Model + -> +Task + -> +Trainer +``` + +Document responsibility boundaries: + +```text +Model: + logits + +Task: + loss, predictions, metrics + +Trainer: + execution lifecycle + +Checkpoint manager: + future persistence/resume +``` + +Document named metadata usage. + +Update: + +```text +docs/migration.md +``` + +after Task 10 passes validation. + +--- + +# Validation + +Run focused training tests: + +```bash +uv run pytest tests/unit/training -v +``` + +Run task/model tests: + +```bash +uv run pytest tests/unit/tasks -v +uv run pytest tests/unit/models/root_gnn -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy code remains unchanged +* no positional tracking access has returned +* trainer does not contain task-specific loss formulas +* no checkpoint files are written +* no Hydra/CLI integration was added +* no DDP implementation was added +* no AMP/torch.compile optimization was added +* no unrelated changes are included + +--- + +# Completion criteria + +Task 10 is complete when: + +1. A clean `Trainer` exists. +2. One-batch training works. +3. One-epoch training works. +4. Evaluation works without gradients. +5. Epoch-level metrics operate on accumulated split outputs. +6. Epoch loss aggregation is explicit and tested. +7. Model/task/trainer responsibilities remain separated. +8. CPU device execution works. +9. Explicit device movement works. +10. Optimizer construction supports the active workflows. +11. Frozen-backbone fine-tuning optimizes only trainable parameters. +12. Unfrozen fine-tuning trains the backbone where expected. +13. Active scheduler behavior is implemented if required. +14. Early stopping is a separate testable policy. +15. Reproducibility configuration is explicit. +16. Training history is returned cleanly. +17. Binary training integration passes. +18. Multiclass training integration passes. +19. Relevant deterministic legacy parity checks pass. +20. Unit tests pass. +21. Integration tests pass. +22. Full tests pass. +23. No checkpoint persistence/resume has been implemented. +24. No DDP/Hydra/CLI implementation has been added. +25. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. Trainer API +4. train-epoch lifecycle +5. evaluation lifecycle +6. EpochResult/history design +7. device-handling strategy +8. optimizer support +9. fine-tuning optimizer behavior +10. scheduler behavior +11. scheduler stepping semantics +12. early-stopping behavior +13. epoch-loss aggregation semantics +14. metric accumulation strategy +15. reproducibility/seeding strategy +16. logging behavior +17. binary training integration results +18. multiclass training integration results +19. one-step legacy parity results +20. scheduler parity results +21. intentional deviations from legacy internals +22. performance features intentionally deferred +23. unresolved ambiguities +24. validation commands and results + +After validation succeeds, create one Git commit containing only Task 10 changes. + +Use: + +```text +feat: implement training and evaluation lifecycle +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task11.md b/tasks/task11.md new file mode 100644 index 0000000000000000000000000000000000000000..34f3ff48decd57f4dc98e8549fd9d057c5de7ac1 --- /dev/null +++ b/tasks/task11.md @@ -0,0 +1,1486 @@ +# Task 11: Implement Checkpointing, Resume, and Pretrained Compatibility + +Implement durable checkpoint save/load behavior for GNN4Colliders training. + +This task builds on: + +```text +Task 7: metadata-aware dataset orchestration +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: task/loss/metric layer +Task 10: training and evaluation lifecycle +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/models/ +src/gnn4colliders/training/ +src/gnn4colliders/tasks/ +tests/unit/ +tests/integration/ +tests/parity/ + +legacy/root_gnn_dgl/scripts/training_script.py +legacy/root_gnn_dgl/root_gnn_base/utils.py +legacy/root_gnn_dgl/models/GCN.py +``` + +Focus on the active legacy checkpoint behavior and the pretrained ROOT-GNN / fine-tuning workflow. + +Do not implement inference output writing, ONNX export, Hydra/CLI wiring, or distributed checkpointing in this task. + +--- + +# Goal + +Implement a clean checkpoint layer that supports three primary workflows: + +```text +1. Resume training +2. Load a model for evaluation +3. Load a pretrained ROOT-GNN backbone for fine-tuning +``` + +The intended architecture is: + +```text +Trainer +Model +Optimizer +Scheduler +EarlyStopping +TrainerState + ↓ +CheckpointManager + ↓ +checkpoint file +``` + +Checkpoint persistence should be separate from the core `Trainer`. + +The trainer owns training state. + +The checkpoint layer owns serialization and restoration. + +--- + +# 1. Package structure + +Prefer: + +```text +src/gnn4colliders/training/ + checkpoint.py +``` + +If compatibility code becomes substantial, a small additional module is acceptable: + +```text +src/gnn4colliders/training/ + checkpoint.py + legacy_checkpoint.py +``` + +Do not create a large checkpoint framework. + +--- + +# 2. Explicit checkpoint schema version + +Introduce a checkpoint schema version. + +For example: + +```python +CHECKPOINT_SCHEMA_VERSION = 1 +``` + +Every new-format checkpoint must include it. + +On load: + +* recognize supported schema versions +* fail clearly for unsupported future versions +* do not silently guess + +Do not use package version alone as the checkpoint format version. + +--- + +# 3. New checkpoint payload + +Define an explicit checkpoint payload. + +Conceptually: + +```python +{ + "schema_version": 1, + "epoch": 12, + "global_step": 5400, + + "model_state_dict": ..., + "optimizer_state_dict": ..., + "scheduler_state_dict": ..., + "early_stopping_state": ..., + "trainer_state": ..., + + "model_config": ..., + "task_config": ..., + + "metadata": { + ... + }, +} +``` + +The exact shape may differ. + +Keep field names clear and stable. + +Do not serialize arbitrary live Python objects when structured primitive data/state dicts are sufficient. + +--- + +# 4. Checkpoint should be self-describing + +Store enough information to understand and reconstruct the model/task. + +At minimum consider: + +```text +model family +model architecture/config +output size +task type +checkpoint schema version +feature schema version +graph schema version +``` + +Where available, also consider: + +```text +GNN4Colliders version +git commit +creation timestamp +training seed +``` + +Do not make Git availability a hard requirement. + +The checkpoint should not require someone to remember which YAML file created it just to know the model shape. + +--- + +# 5. Model configuration + +Store the model configuration needed to reconstruct: + +```text +EdgeNetwork +FineTunedEdgeNetwork +``` + +This may include things such as: + +```text +input dimensions +hidden dimensions +number of processing steps +dropout +output size +global feature dimensions +classifier configuration +``` + +Use the actual Task 8 constructor/config interface. + +Do not duplicate model defaults in checkpoint code. + +--- + +# 6. Task configuration + +Store enough task information to interpret outputs later. + +Examples: + +```text +binary_classification +multiclass_classification +num_classes +absolute-weight behavior +prediction threshold +``` + +Only persist behavior that is relevant to reproducing the task. + +Do not store transient metric outputs as task configuration. + +--- + +# 7. Trainer state + +Persist the in-memory trainer state created in Task 10. + +At minimum: + +```text +epoch +global_step +``` + +If Task 10 tracks additional lifecycle state, preserve it where needed. + +The restored trainer should know where to continue. + +Define clearly whether: + +```text +epoch +``` + +means: + +```text +last completed epoch +``` + +or: + +```text +next epoch to execute +``` + +Choose one convention and test it. + +--- + +# 8. Optimizer state + +For resume training, save: + +```python +optimizer.state_dict() +``` + +and restore it. + +Test that optimizer state survives round-trip. + +For optimizers such as Adam, this includes internal moment estimates. + +Do not claim training resume parity if only model weights are restored. + +--- + +# 9. Scheduler state + +If a scheduler is configured, save and restore: + +```python +scheduler.state_dict() +``` + +If no scheduler exists, store a clear `None` or omit according to the chosen schema. + +Keep behavior consistent. + +Test LR progression after resume. + +--- + +# 10. Early-stopping state + +Persist the Task 10 early-stopping policy state where needed. + +Examples: + +```text +best metric +bad epoch count +stopped flag +``` + +Do not serialize arbitrary class instances if a small state dictionary is sufficient. + +Add: + +```python +state_dict() +load_state_dict(...) +``` + +to the policy if that is the cleanest design. + +--- + +# 11. RNG state for reproducible resume + +Evaluate whether exact-ish training continuation requires RNG state. + +Where practical, support persisting: + +```text +Python random state +NumPy random state +Torch CPU RNG state +Torch CUDA RNG state +``` + +This is particularly useful if training resumes mid-experiment. + +If implemented, keep it isolated and optional. + +Document that exact GPU reproducibility may still depend on CUDA/kernel behavior. + +Do not overpromise bitwise deterministic resume. + +--- + +# 12. CheckpointManager API + +Prefer a small explicit API. + +Conceptually: + +```python +manager = CheckpointManager(directory) + +manager.save( + name=..., + model=model, + optimizer=optimizer, + scheduler=scheduler, + trainer_state=trainer.state, + early_stopping=early_stopping, + model_config=model_config, + task_config=task_config, +) +``` + +and: + +```python +checkpoint = manager.load(path) +``` + +or a cleaner equivalent. + +Keep serialization details out of the trainer. + +--- + +# 13. Restore API + +Provide an explicit restore path. + +Conceptually: + +```python +restore_training_state( + checkpoint, + model=model, + optimizer=optimizer, + scheduler=scheduler, + early_stopping=early_stopping, + trainer=trainer, +) +``` + +or equivalent. + +Avoid a single giant magic function that silently mutates everything without making restored components clear. + +--- + +# 14. Weight-only loading + +Support loading only model weights for evaluation. + +Conceptually: + +```python +load_model_weights( + model, + checkpoint, +) +``` + +This path should not require: + +```text +optimizer +scheduler +early stopping +``` + +Use this later for Task 12 evaluation/inference. + +--- + +# 15. Resume-training workflow + +Support: + +```text +new process + -> +construct model/task/trainer/optimizer + -> +load checkpoint + -> +restore state + -> +continue at correct epoch/global step +``` + +Add an integration test that trains for a small number of steps, saves, restores into fresh objects, and continues. + +--- + +# 16. Resume parity test + +Use a tiny deterministic setup. + +Compare: + +```text +Run A: +train continuously for N epochs +``` + +against: + +```text +Run B: +train for K epochs +save +restore +train remaining N-K epochs +``` + +Compare final model parameters where determinism permits. + +Use CPU and explicit seeds to maximize reproducibility. + +If exact equality is not realistic, use a tight justified tolerance. + +--- + +# 17. Latest checkpoint selection + +Implement explicit latest-checkpoint discovery. + +Do not depend on lexicographic filename ordering that breaks on: + +```text +epoch_9 +epoch_10 +``` + +Parse numeric epoch information or use metadata. + +Prefer a stable convention. + +For example: + +```text +checkpoints/ + epoch_0000.pt + epoch_0001.pt + epoch_0012.pt +``` + +Zero-padded filenames are encouraged. + +--- + +# 18. Best checkpoint selection + +Support identifying a "best" checkpoint based on a monitored metric. + +Do not infer best by scanning training logs. + +Prefer checkpoint metadata or a small manifest/index. + +For example, each checkpoint may record: + +```text +monitor_name +monitor_value +``` + +and the manager can select according to: + +```text +mode = min | max +``` + +Keep the implementation small. + +--- + +# 19. Last/best aliases + +If useful, support convenience paths such as: + +```text +last.pt +best.pt +``` + +but avoid duplicating large checkpoint files unnecessarily if symlinks or small manifests are more appropriate. + +On HPC filesystems, portability matters. + +Use the simplest robust approach. + +--- + +# 20. Atomic writes + +Avoid leaving corrupted checkpoints if a job terminates during save. + +Prefer: + +```text +write temporary file + -> +flush/close + -> +atomic rename +``` + +where practical. + +Do not overwrite a good checkpoint in-place if a safer temporary-write pattern is straightforward. + +Add a small unit test if possible. + +--- + +# 21. Directory layout + +Use a clear checkpoint directory layout. + +Conceptually: + +```text +outputs/ + checkpoints/ + epoch_0000.pt + epoch_0001.pt + ... +``` + +Do not hardcode `/pscratch` or `/global/cfs`. + +The caller chooses the directory. + +Do not implement the full experiment-output directory system yet. + +--- + +# 22. Legacy checkpoint compatibility + +Implement compatibility with active legacy ROOT-GNN checkpoints where needed. + +Legacy checkpoint files contain fields such as: + +```text +epoch +model_state_dict +optimizer_state_dict +early_stop +``` + +and may include state-dict key prefixes such as: + +```text +module. +_orig_mod. +``` + +Verify actual legacy behavior before implementing compatibility. + +--- + +# 23. Isolate legacy compatibility + +Do not contaminate the clean new checkpoint schema with legacy assumptions. + +Prefer: + +```python +load_legacy_checkpoint(...) +``` + +or: + +```python +convert_legacy_checkpoint(...) +``` + +that produces a normalized internal representation. + +Conceptually: + +```text +legacy file + ↓ +legacy adapter + ↓ +normalized checkpoint representation + ↓ +new model/load APIs +``` + +--- + +# 24. Legacy state-dict key normalization + +Support required prefix cleanup such as: + +```text +module. +_orig_mod. +``` + +only where verified. + +Do not blindly strip arbitrary prefixes. + +Add unit tests for each supported legacy case. + +--- + +# 25. Legacy model-name mapping + +If Task 8 changed module/class names, reuse Task 8's explicit legacy-to-new state-dict mapping. + +Do not duplicate a second independent mapping inside checkpoint code. + +There should be one compatibility path. + +Fail clearly if active parameters remain unmapped. + +--- + +# 26. Legacy optimizer restoration + +Only support legacy optimizer-state restoration if it is needed and can be done safely. + +The highest-priority legacy compatibility workflows are: + +```text +load pretrained model +evaluate old model +fine-tune old pretrained backbone +``` + +Full continuation of historical optimizer state is useful but less important if architectural/module mappings make it unreliable. + +If unsupported, document it explicitly. + +Do not pretend full resume compatibility exists if it does not. + +--- + +# 27. Pretrained ROOT-GNN loading + +Provide a clean workflow for Task 8 transfer learning. + +Conceptually: + +```python +pretrained = load_pretrained_edge_network( + checkpoint_path, +) +``` + +followed by: + +```python +model = FineTunedEdgeNetwork.from_pretrained( + pretrained, + out_size=1, + freeze_backbone=True, +) +``` + +or equivalent. + +Do not require constructing the legacy class to use new code unless unavoidable for compatibility conversion. + +--- + +# 28. Pretrained backbone-only loading + +Support loading the reusable representation/backbone separately from the old classifier. + +The workflow should make it easy to: + +```text +load pretrained multiclass model + -> +reuse backbone + -> +discard 12-class classifier + -> +attach binary classifier +``` + +This is a core project workflow. + +Add a direct test for it. + +--- + +# 29. Classifier mismatch handling + +When loading a pretrained checkpoint into a fine-tuning model: + +* backbone weights should load strictly +* old classifier mismatch should be intentional +* unexpected backbone mismatches should fail + +Do not use broad: + +```python +strict=False +``` + +without inspecting missing/unexpected keys. + +If `strict=False` is necessary at the transfer boundary, validate exactly which keys are allowed to differ. + +--- + +# 30. New-format model reconstruction + +Provide a way to reconstruct a new model from checkpoint configuration. + +Conceptually: + +```python +model = build_model_from_checkpoint(checkpoint) +``` + +The exact API may differ. + +Avoid circular dependencies between training and configuration layers. + +A small model factory based on stored model config is acceptable. + +Do not implement Hydra config composition here. + +--- + +# 31. Task reconstruction + +Similarly, provide enough stored data so Task 12 can determine whether the checkpoint corresponds to: + +```text +binary classification +multiclass classification +``` + +A full task factory may be added if simple. + +Do not over-engineer registries. + +--- + +# 32. Map location + +Support loading checkpoints on a different device. + +For example: + +```python +torch.load( + path, + map_location="cpu", +) +``` + +The default load path should work on CPU even for checkpoints produced on GPU where feasible. + +Do not assume CUDA availability. + +--- + +# 33. Safe loading behavior + +Use current supported PyTorch loading semantics carefully. + +Do not deserialize arbitrary untrusted Python objects unnecessarily. + +Prefer state dictionaries and primitive structured data. + +Avoid placing custom live objects in checkpoint payloads. + +Document that checkpoint files should be treated as trusted artifacts if PyTorch serialization semantics require it. + +--- + +# 34. Checkpoint metadata + +Include useful human-readable metadata. + +For example: + +```text +created_at +epoch +global_step +model_family +task_type +feature_schema_version +graph_schema_version +``` + +Optional: + +```text +git_commit +hostname +``` + +Do not store machine-specific absolute paths unless they are genuinely useful provenance and clearly treated as informational. + +--- + +# 35. Dataset/cache provenance reference + +Do not copy entire datasets or cache metadata into checkpoints. + +It is acceptable to store a compact reference/fingerprint such as: + +```text +preprocessing_fingerprint +feature_schema_version +graph_schema_version +``` + +This lets later inference detect obvious incompatibility. + +--- + +# 36. Compatibility checks on load + +When loading a checkpoint for a model, verify relevant compatibility where practical. + +Examples: + +```text +model family +model config +feature schema +graph schema +output size +``` + +For fine-tuning, classifier output-size mismatch is intentionally allowed. + +Backbone incompatibilities should fail clearly. + +--- + +# 37. Checkpoint save cadence + +Do not embed a fixed save cadence into the manager. + +The future CLI/configuration layer may decide: + +```text +every epoch +every N epochs +best only +last only +``` + +For Task 11, implement the mechanism. + +A small helper for "save each completed epoch" is acceptable if needed for integration with Task 10. + +--- + +# 38. Trainer integration + +Add minimal integration between `Trainer.fit()` and checkpointing if necessary. + +Prefer optional explicit composition. + +For example: + +```python +trainer.fit( + ..., + checkpoint_manager=manager, +) +``` + +or perform checkpoint saves outside the trainer loop using returned epoch state. + +Choose the approach that keeps responsibilities clear. + +Do not build a generic callback system. + +--- + +# 39. Resume start epoch + +Define and test resume semantics precisely. + +For example, if checkpoint says: + +```text +epoch = 4 +``` + +and that means epoch 4 completed, resumed training should begin at: + +```text +epoch = 5 +``` + +Make this unambiguous. + +Avoid accidentally repeating or skipping an epoch. + +--- + +# 40. Global step restoration + +Restore `global_step` exactly. + +Add a test verifying continuation increments from restored value rather than resetting to zero. + +--- + +# 41. Scheduler resume correctness + +After restore, scheduler behavior must continue from the previous state. + +Test learning-rate sequences. + +For example: + +```text +continuous run LR sequence +== +save/resume LR sequence +``` + +for a tiny deterministic scenario. + +--- + +# 42. Fine-tuning does not restore source optimizer by default + +When loading a pretrained model for a new fine-tuning task, do not automatically restore the source pretraining optimizer/scheduler. + +The default transfer workflow should be: + +```text +load pretrained model weights/backbone + -> +construct new classifier + -> +construct new optimizer + -> +fine-tune +``` + +This is different from resume training. + +Make the APIs distinguish these cases clearly. + +--- + +# 43. Distinguish resume vs pretrained load + +Avoid one ambiguous `load_checkpoint()` behavior that sometimes restores everything and sometimes does not. + +Make the distinction explicit: + +```text +resume_training(...) +load_weights(...) +load_pretrained_backbone(...) +``` + +or equivalent. + +This is an important API boundary. + +--- + +# 44. Unit tests + +Add tests under: + +```text +tests/unit/training/ +``` + +Suggested coverage: + +```text +test_checkpoint_roundtrip.py +test_checkpoint_selection.py +test_legacy_checkpoint.py +test_resume_state.py +``` + +Use a smaller organization if clearer. + +--- + +# 45. New checkpoint round-trip test + +Verify: + +```text +model state +optimizer state +scheduler state +early stopping state +trainer state +model/task metadata +``` + +survive save/load. + +Use tiny models/fixtures. + +--- + +# 46. Weight-only round-trip + +Verify a fresh model loaded from a checkpoint produces the same output on the same deterministic `GraphBatch`. + +Use tight numerical comparison. + +--- + +# 47. Pretrained fine-tuning test + +Test: + +```text +multiclass EdgeNetwork checkpoint + -> +load pretrained backbone + -> +FineTunedEdgeNetwork(out_size=1) +``` + +Verify: + +* backbone weights match +* old classifier is not reused +* new classifier has correct output dimension +* forward pass succeeds + +--- + +# 48. Freeze behavior after pretrained loading + +Verify: + +```text +freeze_backbone=True +``` + +still freezes transferred parameters after checkpoint loading. + +Verify the classifier remains trainable. + +--- + +# 49. Best/latest selection tests + +Create multiple tiny checkpoint files with different: + +```text +epochs +monitor values +``` + +and verify: + +```text +latest +best +``` + +selection independently. + +Do not rely on filesystem modification time. + +--- + +# 50. Corrupt/incomplete checkpoint behavior + +Add at least basic validation for malformed new checkpoints. + +Examples: + +```text +missing schema_version +missing model_state_dict +unsupported schema_version +``` + +Raise clear errors. + +Do not attempt automatic repair. + +--- + +# 51. Atomic-write failure behavior + +Where practical, ensure temporary checkpoint files do not become valid final checkpoints if save fails. + +Clean up temporary artifacts if possible. + +Do not build elaborate crash recovery. + +--- + +# 52. Integration test: train-save-resume + +Add a small CPU integration test: + +```text +tiny GraphBatch/DataLoader + -> +Trainer + -> +train one or two epochs + -> +save checkpoint + -> +fresh model/trainer/optimizer + -> +restore + -> +continue training +``` + +Verify: + +* epoch resumes correctly +* global step resumes correctly +* optimizer state is restored +* scheduler state is restored +* final parameters match a continuous reference run where practical + +--- + +# 53. Integration test: pretrained fine-tuning + +Add a second integration test: + +```text +multiclass model + -> +checkpoint + -> +load pretrained backbone + -> +binary fine-tuning model + -> +one training step +``` + +Verify the new classifier updates correctly. + +Do not require real collider convergence. + +--- + +# 54. Legacy parity test + +Where a representative legacy checkpoint is available or can be constructed deterministically, verify the new compatibility loader can recover active model weights. + +Test at least: + +```text +module. prefix cleanup +_orig_mod. prefix cleanup +``` + +where actually supported by legacy behavior. + +Do not add large pretrained checkpoint binaries to tests. + +Prefer a tiny synthetic state dictionary using the same key conventions. + +--- + +# 55. No inference output implementation + +Do not implement: + +```text +evaluation_*.npz +ROOT score branches +prediction files +``` + +Those belong to Task 12. + +--- + +# 56. No ONNX export + +Do not implement ONNX export in this task. + +That belongs with inference/export work later. + +--- + +# 57. No Hydra/CLI wiring + +Do not implement: + +```bash +gnn4colliders train --resume ... +``` + +yet. + +Checkpoint APIs should be usable directly from Python. + +CLI/YAML wiring comes in Task 13. + +--- + +# 58. No distributed checkpointing + +Do not implement: + +```text +rank-aware saves +FSDP checkpoints +distributed checkpoint shards +NCCL barriers +``` + +Task 14 will handle DDP/HPC. + +Baseline checkpointing should be single-process. + +--- + +# 59. Documentation + +Update: + +```text +docs/architecture.md +``` + +to document the checkpoint boundary: + +```text +Trainer + in-memory lifecycle/state + +CheckpointManager + persistence and restoration + +LegacyCheckpointAdapter + compatibility only +``` + +Document the three distinct workflows: + +```text +resume +weight-only evaluation +pretrained transfer +``` + +Update: + +```text +docs/migration.md +``` + +once Task 11 passes validation. + +Document any unsupported legacy resume behavior explicitly. + +--- + +# 60. Public exports + +Expose the intended APIs. + +For example: + +```python +from gnn4colliders.training import ( + CheckpointManager, + load_model_weights, + restore_training_state, +) +``` + +and, if appropriate: + +```python +from gnn4colliders.models.root_gnn import ( + load_pretrained_edge_network, +) +``` + +Keep compatibility internals private unless users genuinely need them. + +--- + +# Validation + +Run focused checkpoint/training tests: + +```bash +uv run pytest tests/unit/training -v +``` + +Run model/task tests: + +```bash +uv run pytest tests/unit/models/root_gnn -v +uv run pytest tests/unit/tasks -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source is unchanged +* new checkpoints use explicit schema versioning +* resume and pretrained loading are distinct APIs +* classifier mismatch is only allowed intentionally for transfer learning +* no inference-output implementation was added +* no Hydra/CLI/DDP implementation was added +* no large checkpoint artifacts are tracked +* no unrelated files are included + +--- + +# Completion criteria + +Task 11 is complete when: + +1. New checkpoints have an explicit schema version. +2. Checkpoints are sufficiently self-describing to reconstruct model/task behavior. +3. Model weights save and restore correctly. +4. Optimizer state saves and restores correctly. +5. Scheduler state saves and restores correctly. +6. Early-stopping state saves and restores correctly. +7. Trainer epoch/global-step state saves and restores correctly. +8. Weight-only loading works independently of resume. +9. Resume training begins at the correct next epoch. +10. Resume preserves optimizer/scheduler progression. +11. Latest checkpoint selection is correct. +12. Best checkpoint selection is correct. +13. Atomic checkpoint writes are used where practical. +14. Legacy state-dict prefixes are handled through an isolated compatibility layer. +15. A legacy/pretrained ROOT-GNN can be loaded for transfer learning. +16. Pretrained backbone loading intentionally excludes/replaces the old classifier. +17. Frozen/unfrozen transfer behavior survives checkpoint loading. +18. Continuous vs save/resume integration behavior matches within justified tolerances. +19. Unit tests pass. +20. Integration tests pass. +21. Relevant parity tests pass. +22. Full tests pass. +23. No inference, ONNX, Hydra/CLI, or DDP implementation has been added. +24. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. checkpoint schema +4. checkpoint payload fields +5. model/task reconstruction metadata +6. TrainerState persistence +7. optimizer persistence +8. scheduler persistence +9. early-stopping persistence +10. RNG-state handling +11. filename/layout convention +12. latest selection behavior +13. best selection behavior +14. atomic-write strategy +15. resume API +16. weight-only API +17. pretrained-backbone API +18. legacy compatibility behavior +19. state-dict prefix handling +20. classifier mismatch handling for transfer +21. resume parity results +22. pretrained fine-tuning integration results +23. intentional differences from legacy format +24. unsupported legacy behavior, if any +25. validation commands and results + +After validation succeeds, create one Git commit containing only Task 11 changes. + +Use: + +```text +feat: implement checkpointing and resume support +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task12.md b/tasks/task12.md new file mode 100644 index 0000000000000000000000000000000000000000..9fbfb16f81464569e3e726bb9edffc555fb8e8fd --- /dev/null +++ b/tasks/task12.md @@ -0,0 +1,1286 @@ +# Task 12: Implement Evaluation, Prediction, and Output Serialization + +Implement the inference/evaluation layer that makes trained GNN4Colliders models usable outside the training loop. + +This task builds on: + +```text +Task 6: ROOT/Awkward ingestion and EventSample +Task 7: metadata-aware batching and GraphBatch +Task 8: ROOT-GNN model and fine-tuning +Task 9: task/loss/metric layer +Task 10: trainer/evaluation lifecycle +Task 11: checkpoint save/load and pretrained compatibility +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/models/ +src/gnn4colliders/tasks/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ + +tests/unit/ +tests/integration/ +tests/parity/ + +legacy/root_gnn_dgl/scripts/inference.py +legacy/root_gnn_dgl/scripts/training_script.py +legacy/root_gnn_dgl/models/loss.py +``` + +Focus on the active legacy inference/evaluation workflow. + +Do not implement Hydra/CLI wiring or DDP in this task. + +--- + +# Goal + +Implement a clean inference API that can: + +```text +checkpoint + -> +model/task reconstruction + -> +DataLoader + -> +prediction + -> +PredictionResult + -> +evaluation metrics and/or serialized outputs +``` + +The inference layer should support both: + +```text +binary fine-tuned models +multiclass pretrained models +``` + +It should preserve sample identity and metadata alignment throughout. + +--- + +# 1. Package structure + +Prefer a small structure such as: + +```text +src/gnn4colliders/inference/ + __init__.py + predictor.py + results.py + writers.py + root_writer.py +``` + +If fewer modules are clearer, use fewer. + +Possible responsibilities: + +```text +predictor.py + model execution and output accumulation + +results.py + prediction/evaluation result types + +writers.py + generic serialization such as NPZ + +root_writer.py + optional ROOT score writing +``` + +Do not build a large inference framework. + +--- + +# 2. Predictor responsibility + +The predictor should orchestrate: + +* model evaluation mode +* device placement +* inference-mode execution +* batch iteration +* logits collection +* task probability/score conversion +* prediction collection +* label collection when available +* metadata/sample ID collection +* optional metric computation + +The predictor should not own: + +* ROOT parsing +* graph construction +* checkpoint file selection policy +* task-specific mathematics +* output file-format details + +--- + +# 3. Predictor API + +Prefer a small explicit API. + +Conceptually: + +```python +predictor = Predictor( + model=model, + task=task, + device=device, +) + +result = predictor.predict(loader) +``` + +The exact API may differ. + +Keep it directly usable from Python. + +Do not wire it to Hydra or CLI yet. + +--- + +# 4. PredictionResult + +Introduce a clear result structure. + +Conceptually: + +```python +@dataclass(frozen=True) +class PredictionResult: + sample_ids: Sequence[str] + logits: torch.Tensor + scores: torch.Tensor + predictions: torch.Tensor + labels: torch.Tensor | None + metadata: ... +``` + +The exact fields may differ by task. + +Requirements: + +* preserve sample order +* preserve sample IDs +* keep raw logits available +* keep task-level scores/probabilities available +* keep hard predictions available where meaningful +* optionally preserve labels +* optionally preserve relevant event metadata + +Do not use positional tracking matrices. + +--- + +# 5. Score semantics + +Use Task 9 task methods to define scores. + +For binary classification, likely: + +```text +scores = sigmoid(logits) +``` + +For multiclass classification, characterize the active legacy behavior carefully. + +Determine whether inference output uses: + +```text +sigmoid per class +softmax +raw logits +``` + +Do not assume. + +Use the task layer as the single source of truth. + +--- + +# 6. Raw logits + +Always keep raw logits separate from scores. + +Conceptually: + +```text +logits +scores / probabilities +predictions +``` + +are distinct outputs. + +Do not overwrite logits after postprocessing. + +--- + +# 7. Prediction semantics + +Use Task 9 prediction behavior. + +For binary classification: + +```text +thresholded prediction +``` + +For multiclass: + +```text +argmax class prediction +``` + +where active behavior requires it. + +Do not duplicate task logic in predictor code. + +--- + +# 8. Labels + +If labels are present in the dataset, preserve them in `PredictionResult`. + +If the inference dataset is unlabeled, allow: + +```python +labels = None +``` + +Do not require labels just to perform prediction. + +--- + +# 9. Event metadata + +Preserve useful event metadata. + +At minimum maintain: + +```text +sample_id +``` + +and preserve additional requested fields where practical. + +For example: + +```text +fold +weight +event_number +run_number +source_file +``` + +Do not force every metadata field into output if it is irrelevant. + +Design the result structure so selected metadata can be retained cleanly. + +--- + +# 10. Sample order + +Prediction order must correspond exactly to DataLoader event order. + +Verify using `sample_id`. + +Do not sort outputs implicitly. + +If an output writer later needs original ROOT event order, make that requirement explicit. + +--- + +# 11. Device behavior + +Support explicit: + +```text +cpu +cuda +cuda:0 +``` + +Do not hardcode CUDA. + +Use: + +```python +model.eval() +torch.inference_mode() +``` + +or equivalent. + +Move output tensors to CPU before long-term accumulation unless there is a clear reason not to. + +Avoid retaining unnecessary GPU memory. + +--- + +# 12. Memory-conscious accumulation + +Do not retain autograd graphs. + +Prefer detached CPU tensors for accumulated outputs. + +For large datasets, structure prediction so future streaming output is possible. + +You do not need to implement streaming-to-disk now unless required by active behavior. + +Avoid building an API that inherently requires all outputs to remain on GPU. + +--- + +# 13. Evaluation metrics + +Provide a clean evaluation path. + +Conceptually: + +```python +result = predictor.predict(loader) + +metrics = task.metrics( + logits=result.logits, + targets=result.labels, + weights=result.weights, +) +``` + +or an equivalent helper. + +Use full-split accumulated metrics, especially for ROC AUC. + +Do not compute ROC AUC independently per mini-batch. + +--- + +# 14. EvaluationResult + +If useful, introduce a small evaluation result type. + +Conceptually: + +```python +@dataclass(frozen=True) +class EvaluationResult: + predictions: PredictionResult + metrics: Mapping[str, float] +``` + +Avoid duplicating tensors unnecessarily. + +Use a different design if simpler. + +--- + +# 15. Event weights in evaluation + +If weighted metrics require event weights, obtain them from named metadata: + +```text +metadata.weight +``` + +not legacy tracking columns. + +Preserve Task 9 semantics exactly. + +Do not reinterpret weights in the predictor. + +--- + +# 16. Checkpoint-based model loading + +Integrate with Task 11. + +Support a convenience workflow conceptually like: + +```python +model, task = load_model_and_task_for_inference(checkpoint) +``` + +or equivalent. + +Use Task 11's self-describing checkpoint metadata. + +Do not duplicate model reconstruction logic. + +--- + +# 17. Weight-only loading + +Inference should use the Task 11 weight-only loading path. + +Do not restore: + +```text +optimizer +scheduler +early stopping +``` + +for ordinary prediction/evaluation. + +Keep resume-training and inference paths distinct. + +--- + +# 18. Fine-tuned checkpoints + +Support loading a binary fine-tuned ROOT-GNN checkpoint. + +Verify: + +* correct backbone +* correct replacement classifier +* correct task type +* correct output dimension +* correct score semantics + +Add a dedicated integration test. + +--- + +# 19. Multiclass checkpoints + +Support loading active multiclass pretrained checkpoints for evaluation. + +Verify: + +* correct model reconstruction +* correct output size +* correct task semantics +* correct score/prediction behavior + +--- + +# 20. NPZ output + +Implement a simple NPZ writer. + +Prefer a function like: + +```python +write_npz(result, path) +``` + +The exact API may differ. + +Store stable, explicit field names. + +At minimum consider: + +```text +sample_id +logits +scores +predictions +labels +``` + +and selected metadata where available. + +Do not recreate opaque legacy positional arrays if named fields are clearer. + +--- + +# 21. NPZ compatibility + +Inspect the active legacy inference output. + +Legacy outputs historically include fields such as: + +```text +scores +labels +tracking_info +``` + +The new format should improve clarity. + +Prefer named fields such as: + +```text +scores +labels +sample_id +fold +weight +``` + +rather than a new `tracking_info` matrix. + +If legacy output compatibility is required, implement it as an optional compatibility writer rather than making it the primary format. + +Document the representation change. + +--- + +# 22. NPZ round-trip tests + +Write a tiny `PredictionResult` to NPZ and load it back. + +Verify: + +* sample IDs +* logits +* scores +* predictions +* labels +* selected metadata + +Use exact equality where possible. + +--- + +# 23. ROOT output writing + +Implement ROOT score writing if the active legacy workflow depends on it. + +Keep ROOT writing separate from model execution. + +The writer should accept a `PredictionResult` plus explicit input/output ROOT information. + +Conceptually: + +```python +write_root_scores( + result, + input_path=..., + output_path=..., + tree_name=..., +) +``` + +The exact API may differ. + +--- + +# 24. ROOT writer semantics + +Inspect the legacy inference implementation carefully. + +Characterize: + +* whether the input tree is cloned +* score branch names +* selection-pass branch behavior +* how multiple scores/classes are represented +* how event alignment is preserved +* behavior for events excluded by selection + +Do not assume. + +Reproduce only active required behavior. + +--- + +# 25. ROOT event alignment + +This is critical. + +Predictions must be written back to the correct source events. + +Use stable identity/provenance from Task 7 where possible, such as: + +```text +source file +tree +entry index +sample_id +``` + +Do not rely only on batch order if preprocessing selections can remove events. + +Add explicit alignment tests. + +--- + +# 26. Selection-pass behavior + +The legacy inference path may write a `selection_pass` branch. + +Characterize the active behavior. + +If the new pipeline filters events before inference, define clearly how scores map back to: + +```text +selected events +unselected events +``` + +Do not silently drop source events when writing a cloned ROOT tree unless that is the intended behavior. + +--- + +# 27. Score branch naming + +Use stable explicit branch names. + +For binary tasks, something like: + +```text +score +``` + +or task-specific names may be appropriate. + +For multiclass tasks, define a clear naming convention. + +For example: + +```text +score_class_0 +score_class_1 +... +``` + +or semantic class names if known. + +Do not make branch naming depend on hidden model internals. + +Document the chosen convention. + +--- + +# 28. Multi-checkpoint inference + +Inspect whether active legacy inference supports evaluating multiple checkpoints/models in one run. + +If this is actively required, implement a small explicit API. + +If not required immediately, defer it. + +Do not expand Task 12 unnecessarily. + +--- + +# 29. Ensemble behavior + +Do not implement ensemble averaging unless active workflows depend on it. + +If legacy inference accepts multiple checkpoints merely to produce separate scores, characterize that separately from true ensembling. + +Document deferred behavior. + +--- + +# 30. Evaluation output files + +If evaluation output historically uses names such as: + +```text +evaluation_.npz +``` + +do not hardcode this naming into predictor logic. + +Let the caller choose output path. + +Future CLI/configuration may choose naming conventions. + +--- + +# 31. Output provenance + +Consider writing compact provenance into NPZ or a sidecar. + +Useful fields may include: + +```text +checkpoint identifier +checkpoint schema version +model family +task type +feature schema version +graph schema version +``` + +Keep it small and human-understandable. + +Do not copy the entire checkpoint into the output. + +--- + +# 32. Inference compatibility checks + +Before prediction, verify relevant compatibility where practical: + +```text +checkpoint model family +task type +feature schema version +graph schema version +dataset representation +``` + +Use Task 11's compatibility metadata. + +Fail clearly on obvious mismatches. + +Do not silently run a checkpoint on incompatible feature/graph schemas. + +--- + +# 33. No optimizer/training state + +The predictor must not require: + +```text +optimizer +scheduler +early stopping +trainer state +``` + +Ordinary prediction should only require model/task/data/checkpoint weights. + +--- + +# 34. No model mutation + +Prediction should not alter learned weights. + +Add a test confirming parameters remain unchanged across inference. + +--- + +# 35. Repeatability + +Given: + +```text +same checkpoint +same data +same ordering +``` + +inference outputs should be deterministic on CPU for the supported deterministic path. + +Set the model to evaluation mode. + +Ensure dropout is disabled through `model.eval()`. + +Add a repeated-prediction test. + +--- + +# 36. Batch-size invariance + +Predictions should not depend on inference batch size. + +Add a test comparing: + +```text +batch_size = 1 +``` + +and: + +```text +batch_size > 1 +``` + +for the same ordered events. + +Compare per-sample logits/scores. + +This is particularly useful for validating DGL batching behavior. + +--- + +# 37. No training-time dropout differences + +Verify evaluation mode behaves correctly. + +A model with dropout should produce stable repeated inference outputs in `eval()`. + +Do not disable dropout manually in model code. + +--- + +# 38. Unit tests + +Add focused tests under: + +```text +tests/unit/inference/ +``` + +Suggested files: + +```text +test_predictor.py +test_results.py +test_npz_writer.py +test_root_writer.py +``` + +Use fewer files if clearer. + +--- + +# 39. Predictor unit tests + +Cover: + +```text +binary model +multiclass model +labels present +labels absent +metadata preservation +sample ID preservation +CPU execution +output shapes +``` + +Use tiny deterministic fixtures. + +--- + +# 40. Checkpoint inference integration test + +Add: + +```text +train or construct tiny model + -> +save Task 11 checkpoint + -> +fresh inference load + -> +predict same batch +``` + +Verify predictions match the original model. + +Do not require a long training run. + +--- + +# 41. Binary fine-tuning integration test + +Cover: + +```text +pretrained/fine-tuned binary checkpoint + -> +load + -> +Predictor + -> +scores + -> +predictions +``` + +Verify output shape: + +```text +[B, 1] +``` + +for logits/scores as appropriate. + +--- + +# 42. Multiclass integration test + +Cover: + +```text +multiclass checkpoint + -> +load + -> +Predictor +``` + +Verify: + +* logits shape +* score shape +* prediction shape +* class index semantics + +Use the active class count. + +--- + +# 43. Evaluation metric integration + +Using a tiny labeled split: + +```text +checkpoint + -> +Predictor + -> +PredictionResult + -> +task metrics +``` + +Verify Task 9 metrics are reproduced. + +Do not duplicate metric formulas in inference tests. + +--- + +# 44. NPZ integration test + +Run inference on a tiny deterministic dataset and write an NPZ. + +Load it back with NumPy. + +Verify sample/score/label alignment using `sample_id`. + +--- + +# 45. ROOT integration test + +If ROOT writing is implemented, create or reuse a tiny ROOT fixture. + +Run inference and write scores. + +Open the resulting ROOT file with Uproot. + +Verify: + +* original expected event count +* score branches exist +* scores map to correct entries +* selection-pass behavior +* unmodified source fields remain present if cloning/preserving is required + +Keep the fixture tiny. + +--- + +# 46. Legacy inference parity + +Where practical, compare the new inference outputs to the active legacy inference path. + +Use: + +* same tiny input +* equivalent model weights +* same checkpoint semantics +* same task postprocessing + +Compare: + +```text +scores +labels +event alignment +``` + +Do not require byte-identical output files. + +Compare semantics. + +--- + +# 47. Legacy NPZ compatibility + +If useful, provide a compatibility writer capable of producing the legacy-style fields. + +For example: + +```text +scores +labels +tracking_info +``` + +Only implement this if there is a real downstream consumer. + +Do not make legacy positional tracking output the new default. + +--- + +# 48. No ONNX export unless required now + +Do not implement ONNX export in the main Task 12 scope unless it is immediately required. + +Prefer a follow-up: + +```text +Task 12b: ONNX export +``` + +because export often requires graph/model wrappers and separate parity work. + +If existing requirements clearly demand ONNX now, document that and keep export isolated. + +--- + +# 49. No Hydra integration + +Do not implement YAML/Hydra configuration composition in this task. + +Use direct Python APIs. + +Hydra/CLI wiring belongs to Task 13. + +--- + +# 50. No CLI + +Do not implement: + +```bash +gnn4colliders predict +gnn4colliders evaluate +``` + +yet. + +Task 13 will expose the stable Python inference APIs through CLI commands. + +--- + +# 51. No DDP + +Do not add distributed inference yet. + +Task 14 will cover rank-aware distributed execution. + +Keep prediction single-process. + +--- + +# 52. No performance optimization + +Do not add: + +```text +torch.compile +AMP +custom CUDA +asynchronous prefetch +multi-process inference +``` + +yet. + +Establish correctness and output parity first. + +--- + +# 53. Public exports + +Expose intended inference APIs. + +For example: + +```python +from gnn4colliders.inference import ( + Predictor, + PredictionResult, + EvaluationResult, + write_npz, +) +``` + +Expose ROOT writers only if they are intended user-facing APIs. + +Keep internal accumulation helpers private. + +--- + +# 54. Documentation + +Update: + +```text +docs/architecture.md +``` + +to describe: + +```text +Checkpoint + -> +Model + Task + -> +Predictor + -> +PredictionResult + ├── evaluation metrics + ├── NPZ writer + └── ROOT writer +``` + +Document that: + +* raw logits remain available +* task postprocessing produces scores/predictions +* sample IDs preserve event identity +* named metadata replaces legacy tracking arrays +* output writers are separate from prediction execution + +Update: + +```text +docs/migration.md +``` + +when Task 12 passes validation. + +Document any legacy inference behavior intentionally deferred. + +--- + +# Validation + +Run focused inference tests: + +```bash +uv run pytest tests/unit/inference -v +``` + +Run task/model/checkpoint tests: + +```bash +uv run pytest tests/unit/tasks -v +uv run pytest tests/unit/models/root_gnn -v +uv run pytest tests/unit/training -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source remains unchanged +* no positional tracking API was reintroduced +* prediction preserves sample IDs +* output writers are separated from Predictor +* no optimizer/training state is required for inference +* no Hydra/CLI/DDP implementation was added +* no large generated prediction artifacts are tracked +* no unrelated changes are included + +--- + +# Completion criteria + +Task 12 is complete when: + +1. A clean `Predictor` API exists. +2. Binary fine-tuned checkpoints can be loaded and evaluated. +3. Multiclass checkpoints can be loaded and evaluated. +4. Prediction uses Task 9 score/prediction semantics. +5. Raw logits are preserved separately. +6. Sample IDs remain aligned with all outputs. +7. Named event metadata remains available where needed. +8. Labeled and unlabeled inference both work. +9. Full-split metrics can be computed for labeled data. +10. Prediction is deterministic in evaluation mode for deterministic CPU tests. +11. Prediction is invariant to inference batch size within justified numerical tolerance. +12. NPZ output works and round-trips. +13. NPZ fields are named clearly rather than using positional tracking arrays. +14. ROOT score writing works if required by the active workflow. +15. ROOT output preserves correct event alignment. +16. Selection-pass behavior is characterized if implemented. +17. Checkpoint compatibility checks are applied. +18. Predictor does not mutate model weights. +19. Unit tests pass. +20. Integration tests pass. +21. Relevant legacy parity tests pass. +22. Full tests pass. +23. No Hydra/CLI/DDP implementation has been added. +24. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. Predictor API +4. PredictionResult design +5. EvaluationResult design +6. binary score semantics +7. multiclass score semantics +8. prediction semantics +9. metadata/sample-ID handling +10. checkpoint-loading integration +11. compatibility checks +12. NPZ schema +13. NPZ round-trip results +14. ROOT output behavior +15. ROOT event-alignment strategy +16. selection-pass behavior +17. legacy inference parity results +18. batch-size invariance results +19. deterministic repeated-inference results +20. intentional differences from legacy output representation +21. behavior intentionally deferred +22. validation commands and results + +After validation succeeds, create one Git commit containing only Task 12 changes. + +Use: + +```text +feat: implement evaluation and inference outputs +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task13.md b/tasks/task13.md new file mode 100644 index 0000000000000000000000000000000000000000..d2c13fa32dc59f7002a37423bdcda2f1b86457e6 --- /dev/null +++ b/tasks/task13.md @@ -0,0 +1,1688 @@ +# Task 13: Implement Hydra Configuration and CLI Wiring + +Implement the user-facing configuration and command-line interface for GNN4Colliders. + +This task builds on the stable Python APIs created in previous tasks: + +```text +Task 6: ROOT/Awkward ingestion +Task 7: metadata-aware caching, splits, batching +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: classification tasks, losses, metrics +Task 10: training lifecycle +Task 11: checkpointing and resume +Task 12: evaluation, prediction, and output serialization +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +README.md +``` + +Then inspect: + +```text +configs/ + +src/gnn4colliders/data/ +src/gnn4colliders/models/ +src/gnn4colliders/tasks/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +src/gnn4colliders/cli/ + +tests/unit/ +tests/integration/ +tests/parity/ + +legacy/root_gnn_dgl/configs/ +legacy/root_gnn_dgl/scripts/ +``` + +Use the existing new Python APIs as the source of truth. + +Do not redesign lower-level components merely to make configuration convenient. + +--- + +# Goal + +Expose the GNN4Colliders workflows through: + +```text +YAML configuration + + +Hydra composition + + +single CLI +``` + +The intended user-facing interface is: + +```bash +gnn4colliders prepare +gnn4colliders train +gnn4colliders evaluate +gnn4colliders predict +``` + +with configuration overrides such as: + +```bash +gnn4colliders train \ + task=pretraining_multiclass \ + model=root_gnn/edge_network \ + trainer.max_epochs=50 +``` + +and: + +```bash +gnn4colliders train \ + task=tth_cp_finetune \ + model=root_gnn/fine_tuned_edge_network \ + checkpoint.pretrained=/path/to/pretrained.pt +``` + +The CLI must remain thin. + +All actual behavior should be implemented through the Python APIs created in Tasks 6–12. + +--- + +# 1. Configuration philosophy + +Configuration should describe experiment intent. + +Do not reproduce the legacy dynamic-import pattern: + +```yaml +module: ... +class: ... +args: ... +``` + +Prefer semantic configuration such as: + +```yaml +model: + name: edge_network + hidden_dim: 128 + processing_steps: 4 +``` + +The configuration should not require users to know internal Python module paths. + +--- + +# 2. Configuration groups + +Use Hydra configuration groups. + +Prefer a structure such as: + +```text +configs/ + config.yaml + + data/ + delphes.yaml + local_test.yaml + + model/ + root_gnn/ + edge_network.yaml + fine_tuned_edge_network.yaml + + task/ + pretraining_multiclass.yaml + binary_classification.yaml + tth_cp_finetune.yaml + + trainer/ + default.yaml + debug.yaml + + checkpoint/ + default.yaml + + inference/ + default.yaml + + environment/ + local.yaml + perlmutter.yaml +``` + +Adjust names if existing project conventions suggest better ones. + +Do not create config groups without a concrete use. + +--- + +# 3. Root config + +Create or update: + +```text +configs/config.yaml +``` + +Use Hydra defaults composition. + +Conceptually: + +```yaml +defaults: + - data: delphes + - model: root_gnn/edge_network + - task: pretraining_multiclass + - trainer: default + - checkpoint: default + - environment: local + - _self_ +``` + +Do not put large amounts of experiment-specific configuration in the root config. + +--- + +# 4. Data configuration + +Map the Task 6/7 APIs into explicit data config. + +Consider fields such as: + +```yaml +data: + files: [] + tree_name: Events + + batch_size: 32 + num_workers: 0 + shuffle: true + seed: 42 + + splits: + train_folds: [0, 1, 2] + validation_folds: [3] + test_folds: [4] + + cache: + enabled: true + directory: null +``` + +Use actual APIs implemented in the project. + +Do not invent fields unsupported by production code. + +--- + +# 5. Avoid machine-specific paths in experiment configs + +Do not hardcode paths such as: + +```text +/global/cfs/... +/pscratch/... +``` + +inside model/task configs. + +Machine/site-specific filesystem roots belong under: + +```text +environment/ +``` + +and may be composed into full paths by the application layer. + +Keep scientific experiment configuration portable. + +--- + +# 6. Environment configuration + +Provide minimal environment profiles. + +At minimum: + +```text +local +perlmutter +``` + +These may describe things such as: + +```yaml +environment: + device: cpu + data_root: null + output_root: outputs +``` + +or, for Perlmutter: + +```yaml +environment: + device: cuda + data_root: /global/cfs/... + output_root: /pscratch/... +``` + +Only include site-specific defaults that are genuinely useful. + +Avoid embedding usernames. + +Environment config should remain overridable. + +--- + +# 7. Model configuration + +Expose Task 8 model constructor parameters semantically. + +For example: + +```yaml +model: + family: root_gnn + name: edge_network + + node_input_dim: 7 + edge_input_dim: 3 + + hidden_dim: 128 + processing_steps: 4 + dropout: 0.1 + + out_size: 12 +``` + +Use actual Task 8 constructor parameters. + +Do not duplicate defaults in multiple locations unnecessarily. + +--- + +# 8. Fine-tuning model configuration + +Provide a config for transfer/fine-tuning. + +Conceptually: + +```yaml +model: + family: root_gnn + name: fine_tuned_edge_network + + pretrained_checkpoint: null + out_size: 1 + + freeze_backbone: true +``` + +Prefer checkpoint-related file paths under a checkpoint or transfer subsection if cleaner. + +The config should clearly distinguish: + +```text +resume training +``` + +from: + +```text +load pretrained backbone for a new task +``` + +Do not make those two workflows ambiguous. + +--- + +# 9. Task configuration + +Expose Task 9 semantics explicitly. + +Binary example: + +```yaml +task: + type: binary_classification + threshold: 0.5 + use_absolute_weights: false +``` + +Multiclass example: + +```yaml +task: + type: multiclass_classification + num_classes: 12 +``` + +Only expose real task parameters. + +Do not put optimizer or DataLoader settings under `task`. + +--- + +# 10. Trainer configuration + +Expose Task 10 trainer settings. + +Conceptually: + +```yaml +trainer: + max_epochs: 100 + device: ${environment.device} + seed: 42 + + optimizer: + name: adam + learning_rate: 1.0e-3 + weight_decay: 0.0 + + scheduler: + name: null + + early_stopping: + enabled: true + monitor: validation_loss + mode: min + patience: 10 +``` + +Map this directly into Task 10 APIs. + +Do not implement new trainer behavior merely because a config field exists. + +--- + +# 11. Checkpoint configuration + +Expose Task 11 checkpoint behavior. + +Conceptually: + +```yaml +checkpoint: + directory: ${environment.output_root}/checkpoints + + resume: null + pretrained: null + + save_every_epochs: 1 + save_best: true + save_last: true +``` + +Keep distinct: + +```text +resume +pretrained +``` + +Do not overload one checkpoint path field for both workflows. + +--- + +# 12. Inference configuration + +Expose Task 12 prediction/output behavior. + +For example: + +```yaml +inference: + checkpoint: null + output: predictions.npz + format: npz +``` + +ROOT-specific options may include: + +```yaml +inference: + format: root + input_root: null + output_root: null + tree_name: Events +``` + +Only expose output formats actually implemented. + +--- + +# 13. Structured configuration + +Use typed structured config where it improves validation and clarity. + +Prefer dataclasses for major config groups if practical: + +```python +@dataclass +class TrainerConfig: + max_epochs: int + seed: int + ... +``` + +Do not create a complex schema system. + +The goal is: + +* meaningful defaults +* type checking +* clear required values +* good Hydra error messages + +--- + +# 14. Validation + +Add configuration validation before expensive work begins. + +Examples: + +* `batch_size > 0` +* `max_epochs > 0` +* `out_size` compatible with task +* pretrained checkpoint required for transfer-learning workflow +* resume and pretrained not both used incorrectly +* required ROOT file paths exist where appropriate +* train/validation/test folds do not overlap if the project requires disjoint splits + +Fail early with clear messages. + +Do not defer obvious configuration errors until model execution. + +--- + +# 15. Cross-config consistency + +Validate relationships between config groups. + +Examples: + +```text +binary task + -> +model out_size = 1 +``` + +and: + +```text +multiclass task + -> +model out_size = num_classes +``` + +and: + +```text +fine-tuned model + -> +pretrained checkpoint supplied +``` + +Do not silently modify user configuration to make it valid. + +Raise clear errors. + +--- + +# 16. Model factory + +Implement a small model factory if one does not exist. + +Conceptually: + +```python +def build_model(config): + ... +``` + +It should map semantic config values such as: + +```text +root_gnn / edge_network +root_gnn / fine_tuned_edge_network +``` + +to the appropriate Task 8 constructors. + +Do not use arbitrary dynamic imports from YAML. + +Avoid a heavyweight registry unless there is a real need. + +--- + +# 17. Task factory + +Provide a small task factory. + +Conceptually: + +```python +def build_task(config): + ... +``` + +Map: + +```text +binary_classification +multiclass_classification +``` + +to Task 9 APIs. + +Do not duplicate task math. + +--- + +# 18. Optimizer/trainer construction + +Use Task 10 builders. + +The CLI/config layer may compose: + +```text +model +task +optimizer +scheduler +early stopping +trainer +``` + +but should not implement their internals. + +--- + +# 19. Data pipeline construction + +Use Tasks 6 and 7 APIs. + +The config/application layer may orchestrate: + +```text +ROOT dataset + -> +feature builder + -> +graph builder + -> +cache + -> +split + -> +DataLoader +``` + +Do not duplicate any of that implementation in CLI modules. + +--- + +# 20. Single CLI + +Expose a single project command: + +```bash +gnn4colliders +``` + +with subcommands. + +At minimum: + +```bash +gnn4colliders prepare +gnn4colliders train +gnn4colliders evaluate +gnn4colliders predict +``` + +Keep command names lowercase and simple. + +--- + +# 21. pyproject CLI entry point + +Configure: + +```toml +[project.scripts] +gnn4colliders = "gnn4colliders.cli:main" +``` + +or the equivalent final entry-point location. + +After installation: + +```bash +uv run gnn4colliders --help +``` + +must work. + +--- + +# 22. CLI structure + +Prefer a small layout such as: + +```text +src/gnn4colliders/cli/ + __init__.py + main.py + prepare.py + train.py + evaluate.py + predict.py +``` + +If fewer files are clearer, use fewer. + +Command modules should be thin. + +--- + +# 23. CLI responsibility + +CLI code may: + +* parse/select the subcommand +* initialize Hydra/config +* validate config +* construct application objects +* invoke stable Python APIs +* configure top-level logging +* return an exit status + +CLI code should not contain: + +* model math +* training loops +* graph construction logic +* loss formulas +* metric formulas +* checkpoint serialization +* ROOT parsing logic + +--- + +# 24. `prepare` command + +Implement: + +```bash +gnn4colliders prepare +``` + +for preprocessing/cache preparation. + +Conceptually: + +```text +config + -> +ROOT dataset + -> +features + -> +graphs + -> +cache +``` + +Use Task 6/7 APIs. + +Do not train a model. + +--- + +# 25. `train` command + +Implement: + +```bash +gnn4colliders train +``` + +Use: + +```text +config + -> +data loaders + -> +model + -> +task + -> +optimizer/scheduler + -> +Trainer + -> +training + -> +checkpoint manager +``` + +Support both: + +```text +training from scratch +resume training +fine-tuning from pretrained backbone +``` + +Keep these workflows explicit. + +--- + +# 26. Train-from-scratch workflow + +Example intended usage: + +```bash +gnn4colliders train \ + task=pretraining_multiclass \ + model=root_gnn/edge_network +``` + +The command should: + +1. construct model +2. construct task +3. construct data loaders +4. construct trainer +5. run fit +6. save checkpoints according to config + +--- + +# 27. Resume workflow + +Support something conceptually like: + +```bash +gnn4colliders train \ + checkpoint.resume=/path/to/checkpoint.pt +``` + +Use Task 11 resume logic. + +Do not accidentally treat this as transfer learning. + +Restore: + +* model +* optimizer +* scheduler +* trainer state +* early stopping state + +as defined by Task 11. + +--- + +# 28. Fine-tuning workflow + +Support: + +```bash +gnn4colliders train \ + model=root_gnn/fine_tuned_edge_network \ + checkpoint.pretrained=/path/to/pretrained.pt \ + task=binary_classification +``` + +Use Task 8/11 pretrained-backbone workflow. + +Do not restore the source optimizer by default. + +Construct a new optimizer for the fine-tuning task. + +--- + +# 29. `evaluate` command + +Implement: + +```bash +gnn4colliders evaluate \ + inference.checkpoint=/path/to/model.pt +``` + +Use Task 12 Predictor/Evaluation APIs. + +The command should: + +* load model/task +* construct requested split loader +* run prediction +* compute full-split metrics +* display/log metrics +* optionally write evaluation outputs if configured + +Do not implement evaluation math in CLI code. + +--- + +# 30. `predict` command + +Implement: + +```bash +gnn4colliders predict \ + inference.checkpoint=/path/to/model.pt \ + inference.output=predictions.npz +``` + +Use Task 12 APIs. + +Support implemented formats such as: + +```text +npz +root +``` + +where available. + +Do not require labels for prediction. + +--- + +# 31. Optional `export` command + +Do not implement: + +```bash +gnn4colliders export +``` + +unless ONNX/export functionality already exists from an earlier task. + +If export has not yet been implemented, leave it for a future focused task rather than adding a broken placeholder. + +--- + +# 32. Hydra usage + +Use Hydra for hierarchical config composition and overrides. + +Do not use Hydra to dynamically instantiate arbitrary Python module/class paths from user-controlled YAML. + +Prefer explicit factories or allow-listed structured construction. + +The configuration system should remain understandable from the YAML files. + +--- + +# 33. Hydra run-directory behavior + +Hydra may change the working directory or create run directories depending on configuration/version. + +Configure behavior deliberately. + +Do not allow implicit working-directory changes to break relative dataset paths. + +Prefer stable path handling using: + +```python +Path +``` + +and Hydra utilities where necessary. + +Document the chosen behavior. + +--- + +# 34. Output directory + +Define a predictable experiment output directory. + +Conceptually: + +```text +outputs/ + / + checkpoints/ + predictions/ + logs/ + resolved_config.yaml +``` + +Do not over-engineer experiment tracking. + +The output directory should be configurable. + +--- + +# 35. Save resolved configuration + +For training runs, save the fully resolved configuration alongside outputs. + +For example: + +```text +outputs//resolved_config.yaml +``` + +This is important for reproducibility. + +Do not depend only on the original fragmented Hydra config files. + +--- + +# 36. Run identity + +Introduce a simple run name or experiment name. + +For example: + +```yaml +experiment: + name: pretraining_multiclass +``` + +Do not require globally unique IDs yet. + +A timestamp may be appended if helpful, but tests should not depend on wall-clock naming. + +--- + +# 37. Reproducibility + +Ensure the configured seed reaches: + +* trainer/model initialization +* data shuffling +* any relevant preprocessing randomness + +Do not create separate hidden seeds in CLI code. + +Use the explicit Task 7/10 reproducibility APIs. + +--- + +# 38. Logging configuration + +Configure basic top-level Python logging. + +Allow a config such as: + +```yaml +logging: + level: INFO +``` + +or equivalent if worthwhile. + +Do not add a heavy experiment logging dependency. + +CLI commands should provide useful high-level information such as: + +```text +selected model +task +device +dataset size +checkpoint path +output path +``` + +Do not print huge config objects repeatedly. + +--- + +# 39. Error handling + +Catch only user-facing application/configuration errors where helpful. + +Do not broadly catch: + +```python +except Exception: +``` + +and hide tracebacks. + +For invalid config, missing checkpoint, or missing file, return a clear actionable error. + +Unexpected programming errors should remain visible during development. + +--- + +# 40. Dry-run / config inspection + +Add a lightweight way to inspect resolved configuration without running expensive training if Hydra provides this naturally. + +For example, standard Hydra behavior may already allow printing config. + +Do not build a separate complex dry-run framework unless useful. + +A `--help` path and resolved config visibility are sufficient. + +--- + +# 41. No interactive prompts + +CLI commands should be non-interactive. + +Do not ask users to confirm paths or overwrite behavior interactively. + +For destructive overwrite behavior, either: + +* refuse clearly +* require an explicit config/flag +* choose unique output paths + +This matters for HPC batch jobs. + +--- + +# 42. Configuration names + +Use names that correspond to scientific intent. + +Prefer: + +```text +task=pretraining_multiclass +task=tth_cp_finetune +``` + +over: + +```text +config_17 +experiment_v2_final +``` + +Keep standard configs understandable. + +--- + +# 43. Port active configs + +Port only a small set of active reference configurations from legacy. + +At minimum include representative configs for: + +```text +multiclass pretraining +binary baseline training +binary fine-tuning from pretrained ROOT-GNN +``` + +Use the active legacy configs as behavioral references. + +Do not migrate every historical YAML file. + +--- + +# 44. Preserve semantics, not syntax + +When porting legacy configs: + +* preserve relevant scientific parameters +* preserve data/task/model semantics +* use new config structure +* do not preserve legacy `module/class/args` syntax + +Document any active values that cannot yet be represented. + +--- + +# 45. Fine-tuning reference config + +Provide one complete fine-tuning config showing the intended workflow. + +Conceptually: + +```yaml +defaults: + - /data: delphes + - /model: root_gnn/fine_tuned_edge_network + - /trainer: default + - _self_ + +task: + type: binary_classification + +checkpoint: + pretrained: /path/to/pretrained.pt + +model: + freeze_backbone: true + out_size: 1 +``` + +Do not hardcode a real user path. + +--- + +# 46. Pretraining reference config + +Provide one complete multiclass pretraining config. + +Conceptually: + +```yaml +task: + type: multiclass_classification + num_classes: 12 + +model: + out_size: 12 +``` + +Use actual active architecture values where known. + +--- + +# 47. Debug config + +Add a minimal debug trainer/environment config if useful. + +For example: + +```yaml +trainer: + max_epochs: 1 +``` + +and tiny dataset/cache settings. + +This should make local smoke testing easy. + +Do not create fake production defaults. + +--- + +# 48. Testing configuration factories + +Add unit tests for: + +```text +model factory +task factory +trainer factory +data config validation +cross-config validation +``` + +Avoid testing Hydra internals. + +Test your own mapping/validation logic. + +--- + +# 49. CLI smoke tests + +Add tests for: + +```bash +gnn4colliders --help +gnn4colliders train --help +gnn4colliders predict --help +``` + +or equivalent depending on the CLI implementation. + +Keep them fast. + +--- + +# 50. End-to-end train CLI test + +Use tiny fixtures and a temporary output directory. + +Exercise: + +```text +CLI/config + -> +data + -> +model + -> +task + -> +one training epoch + -> +checkpoint +``` + +Do not require GPU. + +The test should prove that configuration correctly wires stable Python APIs. + +--- + +# 51. Fine-tuning CLI integration test + +Exercise: + +```text +pretraining checkpoint + -> +CLI fine-tuning configuration + -> +load pretrained backbone + -> +new classifier + -> +one training epoch +``` + +Verify a new checkpoint is produced. + +Keep it tiny and deterministic. + +--- + +# 52. Evaluation CLI integration test + +Exercise: + +```text +checkpoint + -> +evaluate command + -> +metrics +``` + +Verify the command completes successfully and produces expected structured output/log behavior. + +--- + +# 53. Prediction CLI integration test + +Exercise: + +```text +checkpoint + -> +predict command + -> +NPZ output +``` + +Load the output and verify expected fields/sample IDs. + +If ROOT writing is implemented, a separate small test may cover that. + +--- + +# 54. Resolved config test + +Verify a training run writes its resolved configuration into the run output directory. + +Load it back and confirm key values such as: + +```text +model +task +seed +batch size +optimizer +``` + +match the executed run. + +--- + +# 55. Path handling tests + +Test relative and absolute input/output paths where practical. + +Ensure Hydra working-directory behavior does not accidentally redirect paths. + +Use temporary directories. + +Do not rely on Perlmutter paths in tests. + +--- + +# 56. No DDP yet + +Do not implement: + +```text +torchrun +DistributedDataParallel +NCCL +DistributedSampler +rank-aware output +``` + +in this task. + +That belongs to Task 14. + +The single-process CLI should become the baseline that DDP wraps later. + +--- + +# 57. No Slurm integration yet + +Do not add job submission logic to the CLI. + +Slurm scripts/configuration belong to Task 14. + +Do not make the package depend on `sbatch` or Perlmutter. + +--- + +# 58. No performance tuning + +Do not add: + +```text +AMP +torch.compile +Numba +custom CUDA +prefetch optimization +``` + +in this task. + +The goal is stable user-facing composition. + +--- + +# 59. No new ML frameworks + +Do not introduce: + +```text +Lightning +Kedro +Click-based framework stacks unless genuinely needed +``` + +simply for CLI convenience. + +Prefer the existing chosen dependencies and Python standard library where reasonable. + +If a small CLI library is necessary, justify it. + +--- + +# 60. README update + +Update README with the intended user workflow. + +Include examples such as: + +```bash +uv sync --extra root-gnn +``` + +then: + +```bash +uv run gnn4colliders prepare ... +``` + +```bash +uv run gnn4colliders train \ + task=pretraining_multiclass +``` + +```bash +uv run gnn4colliders train \ + task=tth_cp_finetune \ + checkpoint.pretrained=/path/to/pretrained.pt +``` + +```bash +uv run gnn4colliders evaluate \ + inference.checkpoint=/path/to/model.pt +``` + +```bash +uv run gnn4colliders predict \ + inference.checkpoint=/path/to/model.pt \ + inference.output=predictions.npz +``` + +Clearly distinguish examples from commands guaranteed to work with real production data if the project is still under migration. + +--- + +# 61. Documentation + +Update: + +```text +docs/architecture.md +``` + +to document: + +```text +Hydra configs + -> +application factories + -> +stable Python APIs + -> +CLI +``` + +Make clear that YAML is configuration, not the implementation. + +Document the principle: + +```text +new experiment -> mostly YAML +new algorithm/behavior -> Python +``` + +Update: + +```text +docs/migration.md +``` + +once Task 13 passes validation. + +--- + +# 62. AGENTS.md + +Add concise rules if not already present: + +```text +CLI modules must remain thin. +Hydra configs describe intent. +Do not expose arbitrary Python module/class imports through YAML. +Application factories map semantic config to explicit supported implementations. +``` + +Do not duplicate large documentation sections. + +--- + +# 63. Public config/application APIs + +If useful, expose a small application construction layer such as: + +```text +src/gnn4colliders/config/ + schema.py + validation.py + factories.py +``` + +or: + +```text +src/gnn4colliders/app/ +``` + +Choose the smallest structure that keeps CLI modules thin. + +Do not turn this into a dependency-injection framework. + +--- + +# Validation + +Validate configuration composition: + +```bash +uv run gnn4colliders --help +``` + +and relevant Hydra config inspection. + +Run focused tests: + +```bash +uv run pytest tests/unit -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Manually smoke-test, using tiny fixtures where possible: + +```bash +uv run gnn4colliders prepare ... +uv run gnn4colliders train ... trainer.max_epochs=1 +uv run gnn4colliders evaluate ... +uv run gnn4colliders predict ... +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source remains unchanged +* CLI modules contain orchestration only +* no task/model mathematics was duplicated +* no arbitrary module/class dynamic imports were introduced +* no Perlmutter-specific paths are hardcoded into general configs +* resume and pretrained transfer remain distinct +* resolved config is saved with runs +* no DDP/Slurm implementation was added +* no unrelated files are included + +--- + +# Completion criteria + +Task 13 is complete when: + +1. Hydra composition works. +2. Config groups exist for data, model, task, trainer, checkpoint, inference, and environment where needed. +3. Configs describe intent rather than Python import paths. +4. Structured validation catches obvious invalid configurations. +5. Binary vs multiclass model/task consistency is validated. +6. Resume and pretrained transfer are represented distinctly. +7. A single `gnn4colliders` CLI exists. +8. `prepare` works through stable Python APIs. +9. `train` works through stable Python APIs. +10. Training from scratch works from configuration. +11. Resume training works from configuration. +12. Fine-tuning from a pretrained ROOT-GNN checkpoint works from configuration. +13. `evaluate` works from a checkpoint. +14. `predict` works from a checkpoint. +15. NPZ prediction output is reachable through CLI. +16. ROOT output is reachable if Task 12 implemented it. +17. The fully resolved run config is saved. +18. Representative active legacy configs have new equivalents. +19. Debug/smoke configuration exists if useful. +20. CLI/config unit tests pass. +21. End-to-end CLI integration tests pass. +22. Full tests pass. +23. No DDP/Slurm/performance optimization has been added. +24. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. config directory structure +4. root Hydra composition +5. structured config types +6. validation rules +7. model factory design +8. task factory design +9. data/trainer construction strategy +10. CLI structure +11. `prepare` workflow +12. `train` workflow +13. resume workflow +14. pretrained fine-tuning workflow +15. `evaluate` workflow +16. `predict` workflow +17. output-directory convention +18. resolved-config behavior +19. active legacy configs ported +20. CLI unit-test results +21. integration-test results +22. full-suite results +23. intentional differences from legacy configuration +24. deferred behavior +25. validation commands and results + +After validation succeeds, create one Git commit containing only Task 13 changes. + +Use: + +```text +feat: add Hydra configuration and CLI workflows +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task14.md b/tasks/task14.md new file mode 100644 index 0000000000000000000000000000000000000000..daab54e9604686d3e3553f93a14e5810a49524f4 --- /dev/null +++ b/tasks/task14.md @@ -0,0 +1,1974 @@ +# Task 14: Implement Distributed Training, Perlmutter Execution, and Slurm Launching + +Implement the distributed/HPC execution layer for GNN4Colliders, with Perlmutter as the first supported production environment. + +This task builds on the stable single-process stack: + +```text +Task 7: deterministic data loading, splitting, caching, and batching +Task 8: ROOT-GNN model and fine-tuning +Task 9: task/loss/metric layer +Task 10: Trainer +Task 11: checkpoints and resume +Task 12: inference/evaluation +Task 13: Hydra configuration and CLI +``` + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +configs/environment/ +configs/trainer/ + +src/gnn4colliders/data/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +src/gnn4colliders/cli/ + +scripts/ +tests/unit/ +tests/integration/ +tests/parity/ + +legacy/root_gnn_dgl/scripts/training_script.py +legacy/root_gnn_dgl/ +``` + +Inspect existing Perlmutter/Slurm scripts in the repository or legacy tree where available. + +Treat the existing single-process implementation as the baseline behavior. + +Do not redesign model, task, data, checkpoint, or CLI APIs merely to support distributed execution. + +--- + +# Goal + +Support production execution of GNN4Colliders using PyTorch `DistributedDataParallel` while preserving the same model/task/trainer semantics as the single-process path. + +The intended architecture is: + +```text +Slurm + ↓ +torchrun / distributed launcher + ↓ +one process per GPU + ↓ +initialize process group + ↓ +DistributedSampler / rank-local DataLoader + ↓ +DistributedDataParallel(model) + ↓ +Trainer + ↓ +rank-aware metric aggregation + ↓ +rank-0 logging/checkpointing +``` + +The same training configuration should work in: + +```text +single-process CPU +single-process GPU +multi-GPU single-node +multi-node GPU +``` + +with only execution/environment configuration changing. + +--- + +# 1. Distributed execution package + +Add a small distributed utility layer. + +Prefer something such as: + +```text +src/gnn4colliders/distributed/ + __init__.py + context.py + launch.py + collectives.py +``` + +or: + +```text +src/gnn4colliders/training/ + distributed.py +``` + +Choose the smallest organization that keeps DDP-specific behavior out of the core trainer. + +Possible responsibilities: + +```text +context.py + rank/world-size/device/process-group context + +collectives.py + metric/output aggregation helpers + +launch.py + distributed initialization/finalization +``` + +Do not create a generic distributed framework. + +--- + +# 2. Preserve single-process behavior + +The current single-process trainer must remain fully functional. + +Do not require a process group for normal execution. + +Prefer a distributed context that can represent: + +```text +distributed = False +rank = 0 +world_size = 1 +local_rank = 0 +``` + +without initializing `torch.distributed`. + +Tests should continue running normally on CPU. + +--- + +# 3. Distributed context + +Introduce an explicit context object if useful. + +Conceptually: + +```python +@dataclass(frozen=True) +class DistributedContext: + enabled: bool + rank: int + local_rank: int + world_size: int + device: torch.device + + @property + def is_main_process(self) -> bool: + return self.rank == 0 +``` + +The exact API may differ. + +Avoid scattering direct environment-variable reads throughout the codebase. + +--- + +# 4. Initialization + +Implement process-group initialization from standard launcher environment variables. + +Support standard variables such as: + +```text +RANK +LOCAL_RANK +WORLD_SIZE +MASTER_ADDR +MASTER_PORT +``` + +Prefer compatibility with: + +```bash +torchrun ... +``` + +Do not invent a private launcher protocol. + +For GPU training, use the appropriate PyTorch backend expected by the environment. + +For CPU tests, allow an appropriate CPU backend when distributed tests are run locally. + +--- + +# 5. Finalization + +Cleanly destroy the process group at application shutdown when initialized. + +Avoid leaving distributed cleanup scattered through command functions. + +Use `try/finally` or a context-manager pattern where appropriate. + +--- + +# 6. Device assignment + +Each rank must explicitly use its assigned device. + +Conceptually: + +```python +torch.cuda.set_device(local_rank) +device = torch.device("cuda", local_rank) +``` + +where required. + +Do not hardcode: + +```text +cuda:0 +``` + +for every rank. + +Do not rely on implicit default CUDA device behavior. + +--- + +# 7. Model wrapping + +Wrap the model in: + +```python +DistributedDataParallel +``` + +only when distributed execution is enabled. + +Keep model construction itself unchanged. + +Conceptually: + +```text +model + ↓ +model.to(device) + ↓ +DDP(model, ...) +``` + +The task object should not need to know whether the model is wrapped. + +--- + +# 8. Avoid DDP-specific model code + +Do not place rank checks or collective operations inside: + +```text +EdgeNetwork +FineTunedEdgeNetwork +``` + +Models remain ordinary PyTorch modules. + +Distributed behavior belongs outside model definitions. + +--- + +# 9. State-dict compatibility + +Ensure Task 11 checkpointing interacts cleanly with DDP. + +New checkpoints should save the underlying model state without permanently introducing: + +```text +module. +``` + +prefixes. + +Prefer saving: + +```python +ddp_model.module.state_dict() +``` + +or an equivalent normalized helper. + +The resulting checkpoint should load into: + +```text +single-process model +DDP model +``` + +without manual user intervention. + +--- + +# 10. Checkpoint restore under DDP + +Support restoring a Task 11 checkpoint before or after DDP wrapping according to whichever design is cleanest. + +Choose one canonical order and document it. + +Avoid separate checkpoint formats for distributed and non-distributed training. + +--- + +# 11. Rank-0 checkpoint writing + +Only the main process should write ordinary checkpoint files. + +Prevent every rank from creating: + +```text +epoch_0010.pt +``` + +simultaneously. + +Use: + +```text +rank == 0 +``` + +or the distributed context property. + +Other ranks should continue participating in training without writing duplicate artifacts. + +--- + +# 12. Checkpoint synchronization + +Where required, synchronize ranks around checkpoint boundaries. + +Do not add barriers unnecessarily. + +Use them only where correctness requires processes to remain aligned. + +Avoid excessive filesystem synchronization on HPC systems. + +--- + +# 13. Rank-aware logging + +Only rank 0 should produce normal user-facing epoch logs by default. + +Avoid: + +```text +8 ranks × identical INFO messages +``` + +for every epoch. + +Errors and truly rank-specific diagnostics may still include rank-prefixed logging. + +Add rank information to distributed debug logs where useful. + +--- + +# 14. Hydra output behavior under multiple ranks + +Task 13 introduced Hydra and output directories. + +Ensure multiple ranks do not independently create conflicting run directories or resolved-config files. + +Only the main rank should write: + +```text +resolved_config.yaml +standard run metadata +ordinary logs +``` + +unless a per-rank artifact is explicitly intended. + +All ranks must agree on the same run/output location. + +--- + +# 15. Data distribution + +Use a rank-aware sampling strategy for training datasets. + +Prefer PyTorch: + +```python +DistributedSampler +``` + +where compatible with the Task 7 dataset/DataLoader architecture. + +Each training event should normally be processed by one rank per epoch. + +Do not load the complete training dataset independently on every GPU unless necessary. + +--- + +# 16. Dataset vs pre-batched semantics + +Inspect Task 7 carefully. + +Determine whether the DataLoader operates over: + +```text +individual GraphSample objects +``` + +or: + +```text +pre-batched objects +``` + +Apply distributed sampling at the correct level. + +Do not accidentally split individual graphs inside an already constructed pre-batch in a way that changes training semantics. + +Document the chosen distribution boundary. + +--- + +# 17. DistributedSampler epoch handling + +For shuffled distributed training, call: + +```python +sampler.set_epoch(epoch) +``` + +at the appropriate time. + +This is required for deterministic but epoch-varying shuffle behavior. + +Integrate it cleanly with Task 10's epoch loop. + +Do not hide the epoch number in global state. + +--- + +# 18. Deterministic shuffling + +Preserve Task 7 reproducibility semantics. + +For the same: + +```text +seed +world size +epoch +dataset +``` + +rank-local sample assignment should be reproducible. + +Do not independently reseed each rank with arbitrary wall-clock values. + +--- + +# 19. Rank-specific seeds + +Where process-local randomness must differ between ranks, derive it deterministically. + +Conceptually: + +```python +rank_seed = base_seed + rank +``` + +or a similarly explicit scheme. + +Do not use this to alter dataset split membership. + +Document exactly which randomness is: + +```text +shared across ranks +rank-specific +``` + +--- + +# 20. Training loss under DDP + +Understand DDP gradient semantics. + +Do not manually all-reduce gradients already handled by DDP. + +Task 9 loss mathematics should remain unchanged. + +Each rank computes its local batch loss. + +DDP handles gradient synchronization during backward. + +--- + +# 21. Reported training loss + +Epoch-level reported training loss must represent the global dataset, not rank 0's local shard only. + +Aggregate the necessary numerator/denominator across ranks. + +Do not simply average rank-level averages unless that is mathematically correct for the existing Task 10 loss aggregation semantics. + +Use explicit reductions. + +--- + +# 22. Metric aggregation + +Task 9 includes metrics such as ROC AUC that require full-split outputs. + +Implement distributed metric gathering carefully. + +Conceptually: + +```text +rank 0 outputs ┐ +rank 1 outputs ├── gather +rank 2 outputs ┤ +... ┘ + ↓ + global metric +``` + +Do not average per-rank ROC AUC values. + +--- + +# 23. Global-output gathering + +Provide helpers for gathering: + +```text +logits +targets +weights +sample IDs if needed +``` + +across ranks. + +For small/normal evaluation workloads, gathering to rank 0 is acceptable. + +Keep the API structured so future large-scale streaming metrics remain possible. + +--- + +# 24. Tensor gathering with unequal lengths + +Different ranks may process different numbers of samples. + +Do not assume identical tensor lengths. + +Implement robust gathering for variable-size tensors. + +Possible strategy: + +```text +gather lengths +pad tensors +all_gather +trim +concatenate +``` + +or use appropriate PyTorch object/tensor gathering utilities if justified. + +Avoid inefficient Python-object collectives for large numeric tensors when a tensor approach is straightforward. + +--- + +# 25. Sample ID gathering + +If sample IDs are required for evaluation/inference alignment, gather them in a way that preserves correspondence with numeric outputs. + +Do not separately sort one field and not the others. + +Maintain: + +```text +sample_id[i] +logits[i] +target[i] +weight[i] +``` + +alignment. + +--- + +# 26. Global ordering + +Distributed evaluation may produce rank-sharded ordering. + +Decide whether the final gathered result should preserve: + +```text +distributed sampler order +``` + +or reconstruct: + +```text +original dataset order +``` + +For inference outputs, original sample identity/order may matter. + +Use stable `sample_id` / source-entry information rather than assuming rank concatenation corresponds to source order. + +Document the behavior. + +--- + +# 27. Distributed validation + +Validation should run across ranks rather than duplicating the full validation set on each GPU, unless there is a strong reason otherwise. + +Use an appropriate distributed sampler. + +Aggregate full-split metrics globally. + +Rank 0 should report the final result. + +--- + +# 28. Distributed test/evaluation + +Support the same execution mechanics for: + +```text +evaluate +``` + +where practical. + +Task 12 prediction semantics must remain unchanged. + +Do not implement distributed ROOT writing in the same first step unless necessary. + +--- + +# 29. Distributed prediction outputs + +For NPZ prediction under DDP, prefer: + +```text +gather complete PredictionResult to rank 0 + ↓ +rank 0 writes one output file +``` + +for the baseline implementation. + +Do not allow every rank to overwrite the same NPZ path. + +If datasets are too large for this later, sharded outputs can be added as a future enhancement. + +--- + +# 30. ROOT inference output under DDP + +Do not concurrently write the same ROOT output file from multiple ranks. + +For this task, choose one of: + +```text +A. gather outputs to rank 0 and write there +``` + +or: + +```text +B. explicitly disallow distributed ROOT writing for now +``` + +Prefer correctness over ambitious parallel I/O. + +Document the limitation if deferred. + +--- + +# 31. Early stopping under DDP + +Task 10 early stopping must use globally aggregated validation metrics. + +The decision should be made consistently across ranks. + +Prefer: + +```text +rank 0 computes decision + ↓ +broadcast stop flag +``` + +or compute identically on all ranks from globally shared metrics. + +Do not allow different ranks to stop on different epochs. + +--- + +# 32. Scheduler behavior under DDP + +Scheduler progression must remain consistent across ranks. + +For validation-metric-based schedulers, all ranks should use the same globally aggregated metric. + +Add a deterministic test where feasible. + +--- + +# 33. Global step semantics + +Define `global_step` consistently. + +Prefer that one synchronized optimizer update across all ranks increments: + +```text +global_step += 1 +``` + +once. + +Do not multiply `global_step` by `world_size`. + +Preserve Task 10/11 resume semantics. + +--- + +# 34. Resume under DDP + +Support: + +```text +multi-GPU run + -> +save checkpoint on rank 0 + -> +restart + -> +all ranks load same checkpoint + -> +resume at correct epoch/global step +``` + +Use the same Task 11 checkpoint format. + +Do not create one model checkpoint per rank. + +--- + +# 35. RNG resume under DDP + +If Task 11 saves RNG states, determine how distributed RNG restoration should work. + +At minimum: + +* model/trainer checkpoint state should resume correctly +* each rank should restore or deterministically reconstruct its expected local RNG state + +If exact per-rank RNG checkpointing is not implemented yet, document the limitation. + +Do not silently claim exact distributed resume reproducibility. + +--- + +# 36. DDP unused parameters + +Do not enable: + +```python +find_unused_parameters=True +``` + +by default merely to suppress errors. + +Determine whether the active ROOT-GNN/fine-tuning model actually requires it. + +For frozen parameters, normal `requires_grad=False` behavior should generally suffice. + +Use the simplest correct DDP configuration. + +--- + +# 37. Gradient accumulation + +Do not add gradient accumulation unless already supported and required. + +Distributed training should first preserve the current: + +```text +one batch -> one optimizer step +``` + +semantics. + +--- + +# 38. Effective global batch size + +Document the relationship: + +```text +global batch size += +per-rank batch size × world size +``` + +for ordinary DDP. + +Do not silently reinterpret an existing `batch_size` config without documentation. + +Choose whether config means: + +```text +per-rank batch size +``` + +or: + +```text +global batch size +``` + +Prefer a clear explicit field/name. + +--- + +# 39. Batch-size config semantics + +I recommend treating: + +```yaml +data: + batch_size: 32 +``` + +as: + +```text +per-process / per-GPU batch size +``` + +unless existing Task 13 semantics make another interpretation clearly preferable. + +Then optionally expose/log: + +```text +effective global batch size = batch_size × world_size +``` + +Do not silently divide a requested global batch across ranks. + +Document the choice. + +--- + +# 40. CLI distributed mode + +Extend the Task 13 CLI/configuration with explicit distributed settings. + +For example: + +```yaml +distributed: + enabled: false + backend: nccl +``` + +or infer enablement from launcher environment while retaining explicit config validation. + +Prefer standard launcher usage such as: + +```bash +torchrun ... -m gnn4colliders.cli ... +``` + +or: + +```bash +torchrun ... $(which gnn4colliders) train ... +``` + +depending on what is robust with the project entry point. + +Do not create a custom multiprocessing launcher unless necessary. + +--- + +# 41. Launcher detection + +If launched through `torchrun`, detect distributed environment cleanly. + +Avoid requiring users to manually specify: + +```text +rank +local_rank +world_size +``` + +in YAML. + +Those are runtime execution values, not scientific configuration. + +--- + +# 42. Single-process CLI remains unchanged + +These should continue to work: + +```bash +uv run gnn4colliders train ... +uv run gnn4colliders evaluate ... +uv run gnn4colliders predict ... +``` + +without `torchrun`. + +Do not force distributed syntax on normal local users. + +--- + +# 43. Environment config + +Extend: + +```text +configs/environment/perlmutter.yaml +``` + +only with settings appropriate for application execution. + +Keep Slurm allocation parameters separate from scientific model/task configuration where practical. + +For example: + +```yaml +environment: + name: perlmutter + device: cuda +``` + +Do not embed usernames or personal scratch paths. + +--- + +# 44. Distributed config group + +Consider: + +```text +configs/distributed/ + single.yaml + ddp.yaml +``` + +For example: + +```yaml +distributed: + enabled: true + backend: nccl +``` + +Keep it minimal. + +Launcher-provided rank/world-size values should not live in YAML. + +--- + +# 45. Slurm scripts + +Add production-oriented Slurm scripts under: + +```text +scripts/slurm/ +``` + +Prefer templates/examples such as: + +```text +train_single_gpu.sh +train_multi_gpu.sh +train_multi_node.sh +evaluate.sh +``` + +Only add scripts that are useful now. + +Avoid unnecessary duplication by sharing environment/setup snippets where practical. + +--- + +# 46. Slurm script philosophy + +Slurm scripts should: + +* request resources +* activate/load the execution environment +* invoke the GNN4Colliders CLI +* launch `torchrun` when distributed +* forward configuration overrides + +They should not contain: + +* model definitions +* training logic +* dataset preprocessing logic +* hardcoded experiment parameters that belong in Hydra + +--- + +# 47. Do not hardcode private user paths + +Slurm examples must not hardcode a specific user's: + +```text +/home/... +/pscratch/sd//... +``` + +Use environment variables, placeholders, or Hydra overrides. + +For example: + +```bash +OUTPUT_ROOT="${PSCRATCH}/gnn4colliders" +``` + +only if that environment variable is reliably present in the target environment. + +Otherwise use clearly documented placeholders. + +--- + +# 48. Perlmutter modules/environment + +Do not guess site-specific module versions. + +Use the project's existing documented environment/setup if available. + +If repository docs already define how `uv`, Python, CUDA, PyTorch, and DGL are made available on Perlmutter, reuse that. + +Do not bake transient site configuration into Python source. + +--- + +# 49. Slurm distributed launch + +Use the normal Slurm + PyTorch distributed pattern appropriate for the environment. + +The script should clearly establish: + +```text +number of nodes +GPUs per node +processes per node +master address +master port +``` + +through standard tooling. + +Avoid manually assigning ranks in shell loops if `srun`/`torchrun` can manage them correctly. + +--- + +# 50. Single-node multi-GPU launch + +Provide a tested/example path conceptually like: + +```bash +torchrun \ + --standalone \ + --nproc_per_node=4 \ + ... train ... +``` + +for a single allocated node. + +The exact command should match the chosen CLI packaging. + +--- + +# 51. Multi-node launch + +Provide a clear Slurm-compatible multi-node launch example. + +Keep node/rank derivation in the Slurm/launcher layer. + +Do not put Slurm-specific parsing inside core training modules unless unavoidable. + +--- + +# 52. Main process semantics + +Create a single helper/property for: + +```python +is_main_process +``` + +Use it for: + +* checkpoint writes +* resolved config writes +* standard logging +* final NPZ writing +* concise user-facing summaries + +Do not repeat `rank == 0` logic everywhere if a context abstraction already exists. + +--- + +# 53. Filesystem race prevention + +Audit all run-output creation from Task 13. + +Ensure distributed execution does not race while creating: + +```text +run directories +checkpoint directories +resolved configs +logs +prediction outputs +``` + +Rank 0 should generally create/write them. + +Other ranks may synchronize afterward if necessary. + +--- + +# 54. Cache behavior under multiple ranks + +Audit Task 7 caching. + +Do not allow multiple ranks to independently create the same cache file unsafely. + +Choose a safe policy, such as: + +```text +cache already exists: + all ranks read + +cache must be created: + rank 0 builds + synchronize + all ranks read +``` + +or another robust method. + +Do not implement concurrent writes to the same cache artifact. + +--- + +# 55. `prepare` command under distributed launch + +The `prepare` command does not need DDP. + +If a user accidentally launches it under multiple ranks, either: + +```text +only rank 0 performs preparation +``` + +or reject the configuration clearly. + +Avoid multiple ranks building the same cache. + +--- + +# 56. Dataloader workers + +Ensure DataLoader worker configuration behaves sensibly under DDP. + +Remember that: + +```text +num_workers +``` + +is usually per rank. + +Do not automatically multiply or divide it without explicit semantics. + +Document: + +```text +total loader workers ≈ num_workers × world_size +``` + +for distributed runs. + +--- + +# 57. Worker seeding + +Ensure worker seeds remain deterministic and distinct. + +Use the existing Task 7 DataLoader worker initialization/generator approach. + +Do not add ad hoc `np.random.seed()` calls inside dataset methods. + +--- + +# 58. Padding/prebatch semantics + +Audit Task 7 padding/prebatch behavior under distributed sampling. + +Ensure a rank-local sample distribution does not invalidate: + +```text +fixed-step padding +node/edge padding +prebatch grouping +``` + +where those behaviors are active. + +Add tests around the actual active mode. + +Do not resurrect unused legacy padding modes simply for DDP. + +--- + +# 59. Rank imbalance + +Distributed samplers may pad/drop samples to create evenly sized shards. + +Characterize this explicitly. + +For training, decide whether: + +```text +drop_last +padding duplicate samples +``` + +matches desired semantics. + +For evaluation, ensure duplicate sampler padding does not cause metrics/prediction outputs to count repeated samples. + +This is especially important. + +--- + +# 60. Evaluation duplicate handling + +If a distributed evaluation sampler pads the dataset by repeating entries, remove duplicate padded samples before global metrics/output. + +Prefer using sample identity/index information to ensure each real event appears once. + +Do not silently bias metrics with repeated final events. + +--- + +# 61. Distributed sampler abstraction + +If standard `DistributedSampler` causes undesirable evaluation duplication, implement a small evaluation sampler that partitions indices without duplication. + +Keep it simple and well-tested. + +Do not write a general sampler framework. + +--- + +# 62. Metric aggregation tests + +Add deterministic tests proving distributed metric aggregation reproduces single-process results. + +Use a small fixture split across multiple CPU processes. + +Compare: + +```text +single-process loss/metrics +``` + +against: + +```text +distributed global loss/metrics +``` + +within justified tolerance. + +--- + +# 63. CPU distributed tests + +Do not require GPUs for the main automated test suite. + +Use: + +```text +torch.distributed +``` + +with a CPU-compatible backend for a small number of spawned local processes. + +Keep tests small and robust. + +--- + +# 64. Distributed test launcher + +For tests, use PyTorch multiprocessing utilities or subprocesses. + +Avoid requiring Slurm. + +The distributed training code should be testable locally. + +--- + +# 65. DDP one-step parity + +Add an integration test: + +```text +same initial model +same total dataset +same effective global batch +``` + +Compare one update from: + +```text +single-process reference +``` + +against: + +```text +2-rank CPU DDP +``` + +where mathematically equivalent. + +Be careful about loss reduction and batch-size semantics. + +Use a tiny deterministic model/data fixture. + +--- + +# 66. Global metric parity + +Explicitly compare: + +```text +loss +accuracy +ROC AUC +``` + +between single-process evaluation and distributed evaluation. + +ROC AUC must be computed from globally gathered samples. + +--- + +# 67. DDP checkpoint test + +Run a tiny distributed training step. + +Verify: + +* only rank 0 creates checkpoint +* checkpoint contains normalized non-DDP model keys +* checkpoint loads in a fresh single-process model +* outputs match + +--- + +# 68. Distributed resume test + +Where practical, test: + +```text +distributed training + -> +checkpoint + -> +new distributed processes + -> +resume +``` + +At minimum verify: + +```text +epoch +global_step +model +optimizer +scheduler +``` + +restore consistently. + +Keep this test small. + +If full multi-process resume is too fragile for routine unit CI, implement a focused integration test or document an opt-in test marker. + +--- + +# 69. Fine-tuning DDP test + +Verify: + +```text +FineTunedEdgeNetwork +freeze_backbone=True +``` + +runs correctly under DDP. + +Only classifier parameters should be trainable. + +Do not enable `find_unused_parameters` solely because the backbone is frozen unless actually required. + +--- + +# 70. Unfrozen fine-tuning DDP test + +Verify an unfrozen transferred backbone receives synchronized gradients across ranks. + +A small parameter comparison after one optimizer step is sufficient. + +--- + +# 71. Logging test + +Verify rank 0 produces the intended standard output/logging while nonzero ranks do not duplicate ordinary epoch messages. + +Do not make tests overly coupled to exact prose strings. + +--- + +# 72. Slurm scripts are examples, not hidden configuration + +Keep resource specifications visible and editable. + +For example: + +```text +nodes +GPUs per node +CPUs per task +wall time +account/constraint/qos if applicable +``` + +should be ordinary Slurm directives or documented variables. + +Do not hide them inside Python configuration. + +--- + +# 73. Account and queue settings + +Do not hardcode a specific Berkeley Lab project/account unless the repository has an explicitly agreed default. + +Use placeholders or environment variables. + +Avoid making example scripts unusable for collaborators. + +--- + +# 74. Signal/preemption handling + +Do not implement elaborate Slurm preemption/requeue handling unless required now. + +However, keep checkpoint APIs compatible with future signal-triggered save logic. + +If existing legacy jobs rely heavily on preemption signals, characterize and document this as a follow-up. + +--- + +# 75. Distributed inference scalability + +For the initial implementation, gathering prediction results to rank 0 is acceptable for moderate datasets. + +Document that truly large distributed inference may later need: + +```text +rank-sharded output +streaming merge +``` + +Do not build that now unless required. + +--- + +# 76. Failure semantics + +If one rank encounters a fatal exception, allow the distributed job to fail clearly. + +Do not catch and suppress rank-local exceptions in ways that leave other ranks hanging indefinitely. + +Use normal PyTorch distributed failure behavior where possible. + +--- + +# 77. Timeout configuration + +Use reasonable process-group initialization defaults. + +Do not hardcode extremely short timeouts that fail on HPC startup. + +Do not expose dozens of low-level distributed tuning options prematurely. + +If timeout is configurable, keep it simple. + +--- + +# 78. No FSDP + +Do not implement: + +```text +FullyShardedDataParallel +ZeRO +DeepSpeed +``` + +in this task. + +ROOT-GNN should first have a correct DDP implementation. + +--- + +# 79. No model parallelism + +Do not implement model/tensor/pipeline parallelism. + +DDP data parallelism is the scope. + +--- + +# 80. No AMP yet unless already separately implemented + +Do not add mixed precision merely because the code is now on GPUs. + +Keep numerical behavior comparable to the existing single-process baseline. + +AMP belongs to the later performance task. + +--- + +# 81. No `torch.compile` yet + +Do not introduce compilation in this task. + +Distributed correctness should be established first. + +--- + +# 82. Public API + +Expose only small useful distributed helpers. + +For example: + +```python +from gnn4colliders.distributed import ( + DistributedContext, + initialize_distributed, +) +``` + +Do not make every internal collective helper public. + +Most users should interact through the CLI. + +--- + +# 83. CLI examples + +Update README with clear examples. + +Single GPU: + +```bash +uv run gnn4colliders train \ + environment=perlmutter \ + task=pretraining_multiclass +``` + +Single-node multi-GPU, conceptually: + +```bash +torchrun \ + --standalone \ + --nproc_per_node=4 \ + ... train \ + distributed=ddp \ + environment=perlmutter +``` + +Use the exact final working command syntax. + +Also point users to: + +```text +scripts/slurm/ +``` + +for batch-job examples. + +--- + +# 84. Document batch-size semantics + +README/docs must explicitly state whether: + +```yaml +data: + batch_size: +``` + +means: + +```text +per GPU +``` + +or: + +```text +global +``` + +Also state how effective global batch size changes with world size. + +This matters scientifically. + +--- + +# 85. Document reproducibility semantics + +Document: + +* base seed +* rank seed derivation +* DataLoader worker seeding +* sampler `set_epoch` +* what reproducibility is guaranteed on CPU +* what GPU nondeterminism may remain + +Do not claim exact deterministic GPU training unless verified. + +--- + +# 86. Architecture documentation + +Update: + +```text +docs/architecture.md +``` + +with: + +```text +Hydra / CLI + ↓ +DistributedContext + ↓ +DataLoader + distributed sampler + ↓ +DDP model + ↓ +Trainer + ↓ +global metric aggregation + ↓ +rank-0 checkpoints/outputs +``` + +Make clear that DDP is an execution concern rather than part of model/task semantics. + +--- + +# 87. Migration documentation + +Update: + +```text +docs/migration.md +``` + +with: + +* which legacy distributed behaviors are reproduced +* which legacy globals/sys.path/environment assumptions were removed +* checkpoint compatibility under DDP +* metric aggregation behavior +* remaining HPC gaps + +--- + +# 88. AGENTS.md + +Add concise durable rules if absent: + +```text +Distributed execution must preserve single-process semantics. +Only rank 0 writes shared run artifacts. +Do not put distributed collectives inside model/task math. +Global metrics must be computed from globally aggregated data. +Tests must not require Slurm or GPUs. +``` + +--- + +# 89. Unit tests + +Add focused tests under something like: + +```text +tests/unit/distributed/ +``` + +Cover: + +```text +context construction +environment parsing +main-process detection +rank/device mapping +collective helpers +variable-length gathering +``` + +Do not require CUDA. + +--- + +# 90. Integration tests + +Add distributed CPU integration tests for: + +```text +training +validation +metric aggregation +checkpoint saving +fine-tuning +``` + +Use two ranks where possible. + +Keep fixtures tiny. + +Mark genuinely expensive tests appropriately if the project has marker conventions. + +--- + +# 91. Slurm smoke validation + +Do not require Slurm in automated tests. + +Validate shell scripts statically where possible. + +At minimum: + +```bash +bash -n scripts/slurm/*.sh +``` + +should pass. + +Avoid executing `sbatch` during tests. + +--- + +# 92. Import hygiene + +Core package import should not fail merely because Slurm commands/environment are absent. + +Slurm support belongs in shell scripts/execution helpers, not import-time checks. + +--- + +# 93. Distributed dependencies + +Prefer PyTorch's built-in distributed stack. + +Do not add Horovod, DeepSpeed, Accelerate, or MPI Python dependencies. + +Add no dependency unless actually required. + +--- + +# 94. Validation + +Run single-process tests first: + +```bash +uv run pytest tests/unit -v +uv run pytest tests/integration -v +``` + +Run focused distributed tests: + +```bash +uv run pytest tests/unit/distributed -v +``` + +and whatever integration marker/path is introduced, for example: + +```bash +uv run pytest tests/integration -k distributed -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Validate Slurm shell syntax: + +```bash +bash -n scripts/slurm/train_single_gpu.sh +bash -n scripts/slurm/train_multi_gpu.sh +``` + +and any other scripts added. + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source remains unchanged +* normal single-process workflows still work +* no DDP-specific behavior leaked into model/task mathematics +* only rank 0 writes shared artifacts +* DDP checkpoints contain normalized model state keys +* distributed metrics represent the complete split +* evaluation does not duplicate padded samples +* Slurm scripts contain no personal hardcoded paths/accounts +* no FSDP/AMP/torch.compile/performance tuning was added +* no unrelated changes are included + +--- + +# Completion criteria + +Task 14 is complete when: + +1. Single-process execution remains unchanged. +2. A clean distributed context/init/finalize layer exists. +3. Standard `torchrun` environment variables are supported. +4. Each rank uses the correct local device. +5. ROOT-GNN can be wrapped in DDP without model-specific distributed code. +6. Training data is partitioned correctly across ranks. +7. Distributed sampler epoch seeding is correct. +8. Global training loss reporting is correct. +9. Validation metrics are globally aggregated. +10. ROC AUC is computed globally rather than averaged per rank. +11. Variable rank-local sample counts are handled. +12. Distributed evaluation does not double-count sampler-padding duplicates. +13. Sample/output alignment survives distributed gathering. +14. Early-stopping decisions are consistent across all ranks. +15. Scheduler behavior is consistent across ranks. +16. `global_step` semantics remain correct. +17. Only rank 0 writes checkpoints. +18. DDP checkpoints load cleanly in single-process mode. +19. Resume works under distributed execution to the supported level. +20. Fine-tuning works with frozen and unfrozen backbones under DDP. +21. Only rank 0 writes normal resolved configs/logs/prediction outputs. +22. Cache creation is safe under distributed launches. +23. Single-node multi-GPU Slurm execution is documented. +24. Multi-node execution is documented or implemented as planned. +25. Perlmutter environment config exists without personal paths. +26. Slurm scripts are thin launch wrappers. +27. CPU-based distributed tests pass. +28. Single-process tests still pass. +29. Full test suite passes. +30. No FSDP, AMP, torch.compile, or unrelated performance work has been added. +31. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. distributed package/API design +4. process-group initialization behavior +5. rank/local-rank/world-size handling +6. device assignment behavior +7. DDP wrapping strategy +8. DataLoader/sampler distribution boundary +9. per-rank batch-size semantics +10. deterministic sampler/seed behavior +11. training-loss aggregation strategy +12. metric gathering strategy +13. variable-length gather implementation +14. evaluation duplicate handling +15. sample-ID/output ordering behavior +16. early-stopping synchronization +17. scheduler synchronization +18. checkpoint rank behavior +19. DDP state-dict normalization +20. resume behavior +21. fine-tuning DDP behavior +22. rank-aware logging/output behavior +23. cache race handling +24. single-node launch method +25. multi-node launch method +26. Slurm scripts added +27. Perlmutter-specific assumptions +28. single-process parity results +29. distributed metric parity results +30. checkpoint compatibility results +31. distributed tests executed +32. limitations/deferred distributed behavior +33. validation commands and results + +After validation succeeds, create one Git commit containing only Task 14 changes. + +Use: + +```text +feat: add distributed and Perlmutter execution +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task15.md b/tasks/task15.md new file mode 100644 index 0000000000000000000000000000000000000000..1c8f067a5b6fa17fbbf2a39509c719d048a8aeb1 --- /dev/null +++ b/tasks/task15.md @@ -0,0 +1,2030 @@ +# Task 15: Profile and Optimize GNN4Colliders Performance + +Implement a focused, measurement-driven performance optimization pass for GNN4Colliders. + +This task builds on the stable and parity-tested implementation from: + +```text +Task 7: metadata-aware caching, splits, batching, and GraphBatch +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: classification tasks, losses, metrics, and output handling +Task 10: training and evaluation lifecycle +Task 11: checkpointing and resume +Task 12: evaluation, prediction, and output serialization +Task 13: Hydra configuration and CLI +Task 14: DDP, Perlmutter, and Slurm execution +``` + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/features/ +src/gnn4colliders/graphs/ +src/gnn4colliders/models/root_gnn/ +src/gnn4colliders/tasks/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +src/gnn4colliders/distributed/ +src/gnn4colliders/cli/ + +configs/ +scripts/slurm/ +tests/ +``` + +Also inspect any existing benchmark/profiling utilities in the repository. + +Treat the currently validated scientific behavior as fixed. + +Do not change model mathematics, task semantics, data selection, feature definitions, graph topology, event weighting, checkpoint semantics, or inference outputs merely to obtain better benchmark numbers. + +--- + +# Goal + +Measure where time and memory are spent in the end-to-end ROOT-GNN workflow, identify the dominant bottlenecks, and implement only optimizations justified by measurements. + +Profile the pipeline: + +```text +ROOT / Awkward I/O + -> +EventSample construction + -> +shared feature construction + -> +graph construction + -> +cache read/write + -> +sampling / batching + -> +GraphBatch + -> +host-to-device transfer + -> +ROOT-GNN forward + -> +task loss + -> +backward + -> +optimizer step + -> +DDP synchronization where enabled +``` + +The task should answer: + +```text +Where is wall-clock time spent? +Where is CPU time spent? +Where is GPU time spent? +Where are synchronization stalls? +Is the GPU waiting for data? +What consumes the most CPU/GPU memory? +Which optimizations materially improve throughput? +Which attempted optimizations do not help? +Do optimized paths preserve existing parity? +``` + +Performance changes must be evidence-driven. + +--- + +# 1. Add a benchmark/profiling area + +Add a dedicated benchmark area. + +Prefer: + +```text +benchmarks/ + README.md + benchmark_preprocessing.py + benchmark_dataloader.py + benchmark_training.py + benchmark_inference.py +``` + +or a similarly small structure. + +Do not put manual performance benchmarks in the normal unit-test suite. + +Benchmark programs should be runnable directly from the repository. + +--- + +# 2. Benchmark reproducibility + +Benchmarks must be reproducible. + +Use: + +* explicit seeds +* fixed benchmark inputs +* fixed iteration counts +* warmup iterations where needed +* explicit device +* explicit batch size +* explicit worker count +* explicit graph/event sizes + +Do not rely on wall-clock-generated random data. + +Record enough information in benchmark output to interpret results. + +--- + +# 3. Benchmark output + +Prefer concise structured benchmark output. + +At minimum report relevant values such as: + +```text +events / second +graphs / second +batches / second +training steps / second +milliseconds / batch +seconds / epoch +``` + +Where relevant also report: + +```text +mean +median +min/max +standard deviation +``` + +Do not report excessive precision. + +--- + +# 4. Benchmark metadata + +Each benchmark run should report or record relevant execution metadata. + +For example: + +```text +Python version +PyTorch version +DGL version +device +GPU model if available +world size +batch size +num_workers +graph/event size +dtype +``` + +Keep this lightweight. + +Do not add a benchmark database or tracking service. + +--- + +# 5. Establish a baseline first + +Before changing production performance code, establish baseline measurements. + +Record baseline numbers in: + +```text +benchmarks/README.md +``` + +or a concise performance document. + +Do not rewrite large parts of the pipeline before measuring the current implementation. + +--- + +# 6. Separate setup from steady-state timing + +Do not mix one-time setup costs into steady-state throughput unless intentionally measuring startup. + +For example distinguish: + +```text +dataset construction +cache construction +first batch +steady-state cached batches +model initialization +first CUDA execution +steady-state forward/backward +``` + +CUDA initialization and kernel compilation/warmup should not distort the steady-state benchmark. + +--- + +# 7. Correct CUDA timing + +For CUDA benchmarks, account for asynchronous execution. + +Use appropriate synchronization around timed regions, such as: + +```python +torch.cuda.synchronize() +``` + +or CUDA events. + +Do not use naïve wall-clock timing around asynchronous CUDA operations and treat it as accurate GPU execution time. + +--- + +# 8. Profiling tools + +Use standard tooling first. + +Prefer: + +```text +cProfile / pstats +PyTorch profiler +torch.utils.benchmark +CUDA memory statistics +``` + +where appropriate. + +Do not add a heavyweight profiling dependency unless necessary. + +--- + +# 9. PyTorch profiler + +Add a documented way to profile a short training section with: + +```python +torch.profiler +``` + +Capture enough iterations to distinguish: + +```text +CPU preprocessing +host-to-device copies +DGL operations +matrix operations +backward +optimizer work +distributed synchronization +``` + +Do not enable profiler overhead by default during ordinary training. + +--- + +# 10. Profile feature construction + +Measure shared feature construction independently. + +Profile: + +* event/object transformation +* energy calculation +* scaling +* stacking/concatenation +* dtype conversions +* avoidable allocations + +Do not change feature values or ordering. + +--- + +# 11. Feature-construction optimization + +Only after profiling, consider safe improvements such as: + +```text +reducing repeated conversions +reducing unnecessary copies +vectorizing Python loops +preallocating outputs +avoiding repeated scale-array construction +``` + +Preserve exact schema and ordering. + +Add parity tests for any changed implementation. + +--- + +# 12. Profile graph construction + +Measure graph construction separately. + +Include: + +```text +edge-index construction +fully connected no-self-loop topology +edge feature calculation +phi wrapping +DGL graph creation +ndata/edata population +``` + +Test representative node counts. + +Remember that fully connected graph edge count scales as: + +```text +N * (N - 1) +``` + +Do not accidentally change topology while optimizing. + +--- + +# 13. Edge-index construction + +Inspect whether edge index generation is repeated unnecessarily. + +If safe, consider efficient deterministic generation of: + +```text +src +dst +``` + +for the fully connected directed no-self-loop graph. + +Preserve the validated edge ordering if parity tests depend on it. + +Do not substitute an ordering with the same mathematical graph unless ordering is explicitly known to be irrelevant. + +--- + +# 14. Edge-feature optimization + +Profile: + +```text +deta +dphi +dR +``` + +calculation. + +Look for: + +```text +repeated indexing +temporary arrays +Python loops +unnecessary device/host conversions +``` + +Optimize only if measurable. + +Preserve phi wrapping exactly. + +--- + +# 15. Graph construction caching + +If graph topology is repeatedly constructed for identical node counts, evaluate whether caching reusable topology indices provides a measurable benefit. + +For example, a bounded cache keyed by: + +```text +number of nodes +device +``` + +may be useful. + +Do not implement unbounded caching. + +Do not cache event-specific edge features as if they were topology. + +Only keep this optimization if benchmarks show a real benefit. + +--- + +# 16. Profile cache behavior + +Measure: + +```text +cold cache creation +warm cache loading +cache deserialization +filesystem throughput +``` + +Distinguish CPU processing cost from filesystem cost. + +Do not infer cache performance from total training time alone. + +--- + +# 17. Cache format changes require evidence + +Do not replace the Task 7 cache format solely because another format seems theoretically faster. + +If cache serialization/deserialization is a proven bottleneck, evaluate a narrowly scoped improvement. + +Any format change must preserve: + +```text +schema versioning +preprocessing fingerprinting +metadata +sample identity +compatibility validation +``` + +Do not break existing caches silently. + +If a format changes, bump the cache schema version. + +--- + +# 18. Profile DataLoader throughput + +Measure DataLoader performance independently from model execution. + +Benchmark combinations of: + +```text +batch_size +num_workers +persistent_workers +prefetch_factor +pin_memory +``` + +where relevant. + +Use a bounded, meaningful search. + +Do not brute-force dozens of configurations without a hypothesis. + +--- + +# 19. DataLoader worker count + +Determine whether additional workers improve actual throughput. + +On Perlmutter this may depend strongly on: + +```text +CPU allocation +filesystem +batching implementation +cache format +``` + +Do not make a large `num_workers` value the global default based on a laptop or a single node. + +Keep defaults conservative unless benchmark evidence supports changing them. + +--- + +# 20. Persistent workers + +Evaluate: + +```python +persistent_workers=True +``` + +for multi-epoch training when `num_workers > 0`. + +Keep it only if: + +* compatible with the dataset implementation +* does not introduce stale state +* improves measured epoch throughput + +Do not enable it when `num_workers=0`. + +--- + +# 21. Prefetching + +Evaluate DataLoader prefetch behavior. + +Do not set very large prefetch values by default. + +Measure both throughput and host-memory cost. + +Expose performance-related loader settings through configuration only when useful. + +--- + +# 22. Pinned memory + +For CUDA workflows, benchmark: + +```python +pin_memory=True +``` + +where compatible with the custom batch structure. + +Do not assume DGL graph objects behave identically to ordinary tensor-only batches. + +Test actual host-to-device timing. + +--- + +# 23. Non-blocking device transfers + +If pinned memory is effective, evaluate: + +```python +tensor.to(device, non_blocking=True) +``` + +for ordinary tensors. + +Do not mark operations non-blocking without satisfying the required memory conditions. + +Keep transfer logic centralized in `GraphBatch.to(...)` or the existing device-transfer boundary. + +--- + +# 24. Profile GraphBatch transfer + +Measure separately: + +```text +DGL graph transfer +labels transfer +global-feature transfer +weight transfer +other numeric metadata transfer +``` + +Do not transfer strings such as `sample_id` to GPU. + +Look for repeated transfers of data that is never consumed by the model/task. + +--- + +# 25. Avoid unnecessary GPU metadata + +Audit which `GraphBatch` fields need to reside on GPU. + +Typically: + +```text +graph features +global features +labels during training/evaluation +weights required by the loss +``` + +need device access. + +Metadata such as: + +```text +sample_id +source file +tree name +entry index +fold +``` + +should remain on CPU unless actually required by GPU computation. + +Do not discard metadata; simply avoid pointless GPU transfer. + +--- + +# 26. Profile model forward + +Profile `EdgeNetwork` forward separately from the DataLoader. + +Measure: + +```text +node encoder +edge encoder +global encoder +each message-passing step +global pooling/update +decoder +classifier +``` + +where profiler attribution allows. + +Determine whether time is dominated by: + +```text +DGL message passing +MLPs +graph pooling +memory movement +Python overhead +``` + +Do not optimize blindly. + +--- + +# 27. DGL graph-local state + +Audit Task 8's use of: + +```python +graph.local_scope() +``` + +and temporary `ndata` / `edata`. + +Ensure temporary graph-state safety is not creating obviously unnecessary copying. + +Do not remove `local_scope()` merely for speed if doing so leaks mutation between forward calls. + +Correctness remains higher priority. + +--- + +# 28. Reduce repeated graph-data lookups + +If profiling shows meaningful Python/DGL lookup overhead, safely reduce repeated access to: + +```text +graph.ndata[...] +graph.edata[...] +``` + +inside tight loops. + +Do not duplicate large tensors merely to save trivial dictionary lookup time. + +Require benchmark evidence. + +--- + +# 29. MLP optimization + +Profile the Task 8 MLP blocks. + +Consider only semantics-preserving changes such as: + +```text +removing avoidable Python overhead +avoiding repeated module construction +using efficient contiguous tensors where needed +``` + +Do not change: + +```text +layer ordering +activation functions +LayerNorm placement +dropout semantics +hidden dimensions +``` + +without explicitly leaving Task 15 scope. + +--- + +# 30. Tensor contiguity + +Profile whether non-contiguous tensors cause material overhead in hot paths. + +Use `.contiguous()` only where it measurably helps or is required. + +Do not scatter unnecessary contiguous copies through the model. + +--- + +# 31. In-place operations + +Do not introduce in-place operations merely for theoretical memory savings if they make autograd or parity fragile. + +Only use in-place operations where clearly safe and beneficial. + +Add tests around any altered computation path. + +--- + +# 32. Profile backward + +Measure forward and backward separately. + +Determine whether training is dominated by: + +```text +forward graph operations +backward graph operations +MLP gradients +optimizer update +DDP all-reduce +``` + +This matters before selecting optimization targets. + +--- + +# 33. Optimizer overhead + +Measure optimizer step overhead. + +Do not replace the scientifically/configurationally active optimizer merely for speed. + +If PyTorch exposes a semantics-compatible implementation option such as: + +```text +foreach +fused +``` + +benchmark it separately. + +Only use it if compatibility and platform support are clear. + +--- + +# 34. `zero_grad(set_to_none=True)` + +Benchmark and consider: + +```python +optimizer.zero_grad(set_to_none=True) +``` + +if not already used. + +This is often a safe performance improvement, but verify: + +* existing code does not rely on zero tensors instead of `None` +* training parity remains intact + +Document the choice. + +--- + +# 35. Mixed precision + +Now that correctness and DDP are established, evaluate mixed precision as an optional performance feature. + +Do not make it mandatory. + +For modern supported CUDA hardware, evaluate: + +```text +bf16 +``` + +before automatically choosing fp16, where the platform supports it. + +The actual supported precision mode should be determined from the current execution environment. + +--- + +# 36. AMP configuration + +If mixed precision is implemented, expose it explicitly. + +Conceptually: + +```yaml +trainer: + precision: float32 +``` + +with supported alternatives such as: + +```text +bfloat16 +float16 +``` + +only if tested. + +Do not silently enable AMP based on CUDA availability. + +--- + +# 37. AMP task/loss behavior + +Keep numerically sensitive operations in suitable precision. + +Do not force Task 9 metric computation into reduced precision. + +Metrics should continue to operate on stable detached outputs. + +Check weighted losses carefully under negative or large event weights. + +--- + +# 38. AMP parity + +For optional mixed precision, do not require bitwise parity with float32. + +Instead validate: + +```text +finite loss +finite gradients +reasonable numerical agreement +stable task metrics +``` + +using explicitly justified tolerances. + +Float32 remains the reference correctness path. + +--- + +# 39. AMP performance threshold + +Do not keep a mixed-precision implementation merely because it technically works. + +Record whether it improves: + +```text +throughput +GPU memory +time per epoch +``` + +on the target GPU environment. + +Document the measured benefit. + +--- + +# 40. `torch.compile` + +Evaluate `torch.compile` only after establishing ordinary eager-mode profiles. + +Do not restructure the entire model around compilation. + +Treat compilation as optional. + +Benchmark: + +```text +first-call compile cost +steady-state throughput +compatibility with DGL +compatibility with DDP +``` + +If the active DGL graph workflow does not benefit or causes graph breaks, document that and do not force it. + +--- + +# 41. Compilation configuration + +If retained, expose compilation explicitly. + +For example: + +```yaml +trainer: + compile: + enabled: false +``` + +Do not enable it by default without clear target-environment evidence. + +--- + +# 42. No custom CUDA kernels in this task + +Do not write custom CUDA/C++ extensions. + +Do not add Triton kernels. + +Do not add CuPy implementations. + +Use existing PyTorch/DGL primitives first. + +Custom kernels require a separate task with dedicated correctness/performance justification. + +--- + +# 43. No Numba unless profiling proves need + +Do not add Numba merely because preprocessing is CPU-side. + +Use NumPy/Awkward/vectorized operations first. + +Only introduce a new compiled dependency if there is a demonstrated bottleneck that cannot be addressed cleanly otherwise. + +Prefer dependency restraint. + +--- + +# 44. DDP profiling + +Profile distributed training separately from single-GPU training. + +Measure: + +```text +computation +all-reduce / communication +data loading +rank imbalance +synchronization +``` + +Use PyTorch profiler or other standard mechanisms where practical. + +Do not assume poor scaling is caused by NCCL before measuring. + +--- + +# 45. DDP scaling efficiency + +For available Perlmutter hardware, benchmark at least conceptually: + +```text +1 GPU +2 GPUs +4 GPUs +``` + +and, if practical: + +```text +multiple nodes +``` + +Report: + +```text +throughput +speedup +scaling efficiency +``` + +For example: + +```text +speedup(N) = throughput(N) / throughput(1) +efficiency(N) = speedup(N) / N +``` + +Do not treat perfect linear scaling as a correctness requirement. + +--- + +# 46. Fair DDP comparisons + +Be explicit about whether scaling benchmarks use: + +```text +fixed per-GPU batch size +``` + +or: + +```text +fixed global batch size +``` + +Do not compare inconsistent workloads without saying so. + +Prefer fixed per-GPU batch for throughput scaling measurements, while documenting the resulting larger global batch. + +--- + +# 47. Rank imbalance + +Profile rank step times. + +Determine whether some ranks systematically receive more expensive graph batches because of varying node/edge counts. + +If graph-size imbalance is material, document it. + +Do not immediately redesign the sampler in this task unless a simple semantics-preserving improvement is available. + +--- + +# 48. Graph-size-aware batching + +If profiling proves that highly variable graph size is a major performance problem, evaluate a simple size-aware batching strategy. + +Examples might include grouping events by approximate: + +```text +node count +edge count +``` + +before batching. + +However, this can change sample ordering and stochastic training behavior. + +Do not make it the default unless: + +* semantics are understood +* reproducibility remains explicit +* parity expectations are updated appropriately +* throughput improves materially + +Prefer documenting this as a follow-up if it becomes a significant algorithmic/data-loader change. + +--- + +# 49. Padding efficiency + +If Task 7 retains an active padding mode, measure: + +```text +real nodes vs padded nodes +real edges vs padded edges +``` + +and quantify wasted work. + +Do not remove a compatibility-required padding mode merely because it is inefficient. + +If a more efficient mode can be optional, benchmark and document it separately. + +--- + +# 50. Memory profiling + +Measure memory at important stages. + +For CPU where practical: + +```text +dataset/cache loading +batch construction +output accumulation +``` + +For CUDA measure: + +```text +allocated memory +reserved memory +peak memory +``` + +using PyTorch-supported APIs. + +Record peak training and inference memory for representative benchmarks. + +--- + +# 51. Detect obvious retained tensors + +Inspect the training and inference loops for accidentally retained computation graphs. + +Examples: + +```text +storing non-detached loss tensors +storing GPU logits for entire epochs +keeping graph references in history +``` + +Fix such issues if found. + +Add regression tests when practical. + +--- + +# 52. Epoch metric accumulation memory + +Task 10/14 may collect logits/targets for full-split metrics. + +Profile the memory cost. + +Keep the current correct behavior as baseline. + +If this becomes a bottleneck, consider moving accumulated tensors to CPU as soon as possible. + +Do not redesign ROC AUC as an approximate streaming metric in this task. + +--- + +# 53. Inference memory + +Profile `PredictionResult` accumulation. + +Ensure predictions are detached and moved to CPU. + +Do not keep all results on GPU. + +If large-result memory is identified as a serious issue, document streaming/sharded output as a future task rather than over-expanding Task 15. + +--- + +# 54. Inference throughput + +Benchmark: + +```text +batch size +DataLoader workers +device transfer +forward +postprocessing +output accumulation +``` + +separately. + +Do not include ROOT/NPZ serialization time in model throughput numbers unless explicitly measuring end-to-end inference. + +--- + +# 55. Output serialization benchmark + +Separately measure: + +```text +NPZ writing +ROOT score writing +``` + +for representative output sizes. + +Do not optimize writers unless serialization is a meaningful bottleneck in actual inference. + +--- + +# 56. CPU-to-GPU overlap + +If profiling shows the GPU waiting significantly for data and pinned/non-blocking transfers are functioning, evaluate whether the existing DataLoader naturally overlaps preprocessing with GPU execution. + +Do not create a custom CUDA stream/prefetch framework unless simpler DataLoader improvements are insufficient. + +If advanced prefetching appears necessary, document it as a follow-up. + +--- + +# 57. Avoid premature micro-optimizations + +Do not spend significant code complexity optimizing components contributing negligible runtime. + +Prioritize the top measured bottlenecks. + +Prefer: + +```text +10% simple improvement +``` + +over: + +```text +1% improvement with substantial architectural complexity +``` + +unless the latter is scientifically/operationally important. + +--- + +# 58. Benchmark before/after every retained optimization + +For each production optimization retained in the final diff, record: + +```text +baseline +optimized +relative improvement +benchmark scenario +device/environment +``` + +Do not claim improvements without before/after numbers. + +--- + +# 59. Remove unsuccessful experiments + +Do not leave speculative optimization code disabled throughout the repository. + +If an experiment does not provide a meaningful benefit: + +* revert it +* document the benchmark result briefly if useful + +Keep production code simple. + +--- + +# 60. Performance configuration + +Only add config options for optimizations that are actually retained. + +Possible examples: + +```yaml +data: + pin_memory: true + persistent_workers: true + prefetch_factor: 2 + +trainer: + precision: float32 + + compile: + enabled: false +``` + +Do not turn every internal implementation detail into YAML. + +--- + +# 61. Conservative defaults + +Do not change baseline defaults merely because a particular Perlmutter benchmark was faster. + +Defaults should remain: + +```text +portable +predictable +correct +``` + +Performance profiles may override settings for the target machine. + +--- + +# 62. Perlmutter performance profile + +If useful, extend: + +```text +configs/environment/perlmutter.yaml +``` + +or add a dedicated performance-oriented trainer/data config. + +For example, a profile could set appropriate: + +```text +num_workers +pin_memory +persistent_workers +precision +``` + +based on actual measurements. + +Do not hardcode values without benchmark evidence. + +--- + +# 63. Debug/local profile + +Keep local/debug configs lightweight. + +Do not apply aggressive Perlmutter worker/GPU settings to ordinary development. + +--- + +# 64. Scientific parity after optimization + +Run existing parity tests after every meaningful production-code optimization. + +At minimum preserve: + +```text +feature values +edge topology/features +model float32 outputs +loss semantics +metric semantics +checkpoint semantics +inference sample alignment +``` + +Performance changes must not invalidate these contracts. + +--- + +# 65. Deterministic model parity + +For float32 eager-mode optimizations, compare fixed-weight model outputs before and after changes. + +Use established parity tolerances. + +Do not weaken tolerances merely to allow an unnecessary optimization. + +--- + +# 66. Training-step parity + +Where optimizer/training internals change, such as: + +```text +zero_grad(set_to_none=True) +fused/foreach optimizer path +``` + +test at least one deterministic training step against the baseline. + +Confirm parameter updates remain equivalent within justified tolerances. + +--- + +# 67. Batch-size invariance + +Performance changes to batching/device movement must not break Task 12 batch-size invariance. + +Run the existing inference tests. + +--- + +# 68. DDP correctness after optimization + +Run Task 14 distributed tests after retained changes. + +Do not optimize single-GPU behavior at the cost of incorrect distributed training. + +In particular verify: + +```text +global loss +metrics +checkpoint writing +fine-tuning +``` + +still work under DDP. + +--- + +# 69. No algorithmic model redesign + +Do not change: + +```text +message-passing equations +number of processing steps +aggregation function +pooling +hidden dimensions +classifier structure +``` + +for performance. + +Those would be new model experiments, not implementation optimization. + +--- + +# 70. No scientific preprocessing redesign + +Do not change: + +```text +selected objects +feature definitions +feature ordering +energy calculation +phi wrapping +graph connectivity +fold/split semantics +event weights +``` + +for speed. + +--- + +# 71. Benchmark command examples + +Document commands in: + +```text +benchmarks/README.md +``` + +For example: + +```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_training.py --device cuda +uv run python benchmarks/benchmark_inference.py --device cuda +``` + +Use the actual final argument interface. + +--- + +# 72. Optional CLI benchmark command + +Do not add: + +```bash +gnn4colliders benchmark +``` + +unless it provides clear value. + +Standalone benchmark scripts are sufficient for Task 15. + +Keep the user-facing scientific CLI focused. + +--- + +# 73. Perlmutter benchmark scripts + +If helpful, add a small Slurm benchmark launcher under: + +```text +scripts/slurm/ +``` + +such as: + +```text +benchmark_single_gpu.sh +benchmark_multi_gpu.sh +``` + +Keep them thin. + +Do not duplicate benchmark logic in shell. + +--- + +# 74. Benchmark artifacts + +Do not commit large profiler traces or benchmark output files. + +Add generated profiling artifacts to `.gitignore` where appropriate. + +Examples: + +```text +*.pt.trace.json +profiles/ +benchmark-results/ +``` + +Do not ignore source benchmark scripts. + +--- + +# 75. Unit tests + +Performance scripts themselves do not need extensive unit tests. + +However, add regression tests for production-code changes. + +Examples: + +```text +optimized graph topology cache returns correct indices +non-blocking batch transfer preserves fields +AMP configuration validation works +compile configuration does not affect default path +``` + +Focus testing on correctness, not timing thresholds. + +--- + +# 76. Do not add timing assertions to ordinary tests + +Do not add brittle tests such as: + +```python +assert runtime < 0.1 +``` + +to the unit/integration suite. + +Performance depends on hardware and CI load. + +Benchmarks measure speed; tests verify semantics. + +--- + +# 77. Benchmark smoke tests + +If useful, add only lightweight import/smoke coverage verifying benchmark scripts can initialize their argument/config setup. + +Do not run meaningful performance workloads in CI. + +--- + +# 78. Documentation + +Update: + +```text +docs/architecture.md +``` + +only where retained optimizations alter implementation architecture. + +For example document: + +```text +pinned/non-blocking transfer boundary +optional precision mode +optional compile mode +topology cache +``` + +Do not fill architecture documentation with benchmark tables. + +--- + +# 79. Performance documentation + +Add: + +```text +docs/performance.md +``` + +or use: + +```text +benchmarks/README.md +``` + +to record: + +* benchmark methodology +* representative hardware +* baseline results +* optimized results +* retained optimizations +* rejected optimizations +* known bottlenecks +* recommended Perlmutter settings + +Keep results clearly tied to hardware/configuration. + +--- + +# 80. README update + +Add a short performance section to README. + +Point to the detailed benchmark documentation. + +Include practical recommended invocation/config examples where validated. + +Do not advertise performance numbers without stating the hardware/configuration used. + +--- + +# 81. AGENTS.md + +Add concise durable performance rules if absent: + +```text +Profile before optimizing. +Do not change scientific semantics for performance. +Benchmark retained optimizations before and after. +Keep performance features optional unless portability is proven. +Do not add timing thresholds to ordinary tests. +``` + +Avoid duplicating detailed benchmark documentation. + +--- + +# 82. Dependency restraint + +Prefer existing dependencies and Python/PyTorch tooling. + +Do not add: + +```text +Numba +CuPy +Triton +DeepSpeed +custom CUDA extensions +external profilers as mandatory dependencies +``` + +without strong measured justification. + +If an external profiling tool is useful manually, document it rather than making it a runtime dependency. + +--- + +# 83. Benchmark ROOT-GNN fine-tuning + +Include at least one representative fine-tuning benchmark. + +Measure: + +```text +frozen backbone +unfrozen backbone +``` + +if both are real workflows. + +This may reveal substantially different backward/optimizer costs. + +--- + +# 84. Benchmark multiclass pretraining + +Include the active multiclass pretraining path. + +Use the standard active output size and representative model configuration. + +Do not benchmark only the smaller binary fine-tuning model. + +--- + +# 85. Frozen-backbone optimization + +When the fine-tuning backbone is frozen, verify unnecessary autograd work is not being performed. + +Parameters with: + +```python +requires_grad = False +``` + +should not accumulate gradients. + +Do not add special-case detached forward logic unless it is correct and measurably beneficial. + +Remember that gradients may still be required through backbone activations depending on what is trainable downstream; reason carefully before detaching anything. + +--- + +# 86. Inference-mode correctness + +Ensure inference uses: + +```python +torch.inference_mode() +``` + +where already appropriate. + +If Task 12 uses `no_grad()` and changing to `inference_mode()` is compatible, benchmark the difference. + +Keep the change only if safe. + +--- + +# 87. Avoid repeated model/device setup + +Audit CLI/training/inference paths for repeated: + +```text +model.to(device) +checkpoint reload +task reconstruction +``` + +inside batch/epoch loops. + +Fix obvious repeated setup if found. + +Add regression coverage where useful. + +--- + +# 88. Avoid repeated configuration parsing + +Configuration resolution should happen at application startup, not per batch. + +Do not micro-optimize Hydra itself unless profiling somehow proves it is in the training hot path, which it should not be. + +--- + +# 89. Avoid repeated filesystem checks in hot loops + +Audit for repeated: + +```text +Path.exists() +cache metadata reads +checkpoint discovery +``` + +inside batch loops. + +Move invariant filesystem work outside hot execution loops. + +Do not weaken correctness checks at application boundaries. + +--- + +# 90. Performance logging + +If useful, add optional per-epoch throughput reporting to the trainer. + +For example: + +```text +samples/sec +batches/sec +epoch duration +``` + +Keep this lightweight. + +Do not synchronize CUDA unnecessarily every batch just to produce timing logs. + +Use coarse epoch-level wall time for ordinary logging. + +--- + +# 91. Distributed throughput logging + +Under DDP, report global throughput correctly. + +For example: + +```text +global samples processed / epoch wall-clock time +``` + +Do not multiply throughput incorrectly if sampler padding/duplicate handling is present. + +Only rank 0 should report ordinary performance summaries. + +--- + +# 92. Warmup exclusions + +For GPU benchmark measurements, exclude initial warmup where appropriate. + +For actual end-to-end startup benchmark, include it explicitly as a separate number. + +Do not mix the two. + +--- + +# 93. Synchronization overhead + +Avoid introducing unnecessary: + +```python +torch.cuda.synchronize() +dist.barrier() +``` + +into production hot loops. + +Synchronization may be needed for profiling but should not remain in normal execution unless required for correctness. + +--- + +# 94. Memory cleanup + +Do not add routine: + +```python +torch.cuda.empty_cache() +``` + +inside training loops as a supposed optimization. + +Use it only for explicit diagnostic/phase-boundary reasons if needed. + +Let PyTorch's allocator manage normal execution. + +--- + +# 95. Benchmark reporting + +At the end of the task, summarize results in a compact table. + +For each retained change report something like: + +```text +Optimization | Scenario | Before | After | Improvement | Memory impact +``` + +Also list tested changes that were rejected. + +Do not cherry-pick only successful results. + +--- + +# 96. Prioritize simplicity + +If two implementations provide similar performance, keep the simpler implementation. + +Do not sacrifice maintainability for marginal gains. + +This repository is intended to support future model families such as ROOT-Transformer, so shared performance infrastructure should remain architecture-neutral where reasonable. + +--- + +# 97. Future model compatibility + +Do not introduce graph-specific assumptions into shared: + +```text +data +training +distributed +inference +``` + +layers solely for ROOT-GNN optimization. + +Graph-specific optimizations belong under: + +```text +graphs/ +models/root_gnn/ +``` + +where possible. + +Keep the future flow viable: + +```text +EventSample + -> +shared features + ├── GraphSample -> ROOT-GNN + └── SequenceSample -> ROOT-Transformer +``` + +--- + +# 98. Validation + +Run the correctness suite before and after performance changes: + +```bash +uv run pytest +``` + +Run focused parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run distributed tests: + +```bash +uv run pytest tests/unit/distributed -v +uv run pytest tests/integration -k distributed -v +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests benchmarks +uv run ruff format --check src tests benchmarks +``` + +Run representative benchmarks. + +At minimum: + +```bash +uv run python benchmarks/benchmark_preprocessing.py +uv run python benchmarks/benchmark_dataloader.py +uv run python benchmarks/benchmark_training.py --device cpu +``` + +On a CUDA environment also run the appropriate: + +```bash +uv run python benchmarks/benchmark_training.py --device cuda +uv run python benchmarks/benchmark_inference.py --device cuda +``` + +On Perlmutter, run the validated single-GPU benchmark and multi-GPU benchmark where available. + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source remains unchanged +* scientific parity is preserved +* no experimental architecture changes were introduced +* no unmeasured speculative optimization remains +* optional performance paths do not change default semantics +* no profiler output artifacts are accidentally tracked +* no personal Perlmutter paths/accounts were added +* no unrelated changes are included + +--- + +# Completion criteria + +Task 15 is complete when: + +1. Reproducible benchmark scripts exist. +2. Baseline preprocessing performance is measured. +3. Baseline DataLoader performance is measured. +4. Baseline model inference performance is measured. +5. Baseline training-step performance is measured. +6. At least representative CPU behavior is characterized. +7. Representative GPU behavior is characterized where the environment permits. +8. PyTorch profiling identifies the main runtime bottlenecks. +9. Peak memory behavior is measured for representative training/inference. +10. Retained optimizations have before/after benchmark evidence. +11. Unsuccessful optimization experiments are removed from production code. +12. Feature-construction parity remains intact. +13. Graph topology/edge-feature parity remains intact. +14. Model float32 parity remains intact. +15. Loss/metric semantics remain intact. +16. Inference output/sample alignment remains intact. +17. Checkpoint compatibility remains intact. +18. Single-process training remains correct. +19. DDP training remains correct. +20. Distributed global metrics remain correct. +21. Optional DataLoader performance improvements are implemented if beneficial. +22. Optional host-to-device transfer improvements are implemented if beneficial. +23. Optional optimizer implementation improvements are implemented only if semantics remain valid. +24. Mixed precision is characterized and optionally supported if beneficial. +25. `torch.compile` is characterized and retained only if beneficial and compatible. +26. Recommended Perlmutter performance settings are documented. +27. Performance results clearly identify hardware and configuration. +28. Full tests pass. +29. Lint/format checks pass. +30. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. benchmark structure +4. benchmark methodology +5. benchmark hardware/environment +6. baseline feature-construction performance +7. baseline graph-construction performance +8. baseline cache performance +9. baseline DataLoader throughput +10. baseline CPU training performance +11. baseline GPU training performance, if measured +12. baseline inference performance +13. baseline peak memory usage +14. profiler-identified bottlenecks +15. retained feature/data optimizations +16. retained graph optimizations +17. retained DataLoader optimizations +18. retained device-transfer optimizations +19. retained model/training optimizations +20. AMP/bfloat16 results +21. `torch.compile` results +22. single-GPU before/after results +23. DDP scaling results +24. memory before/after results +25. rejected optimization experiments and why +26. parity/correctness results +27. recommended local defaults +28. recommended Perlmutter settings +29. remaining bottlenecks +30. suggested future performance work +31. validation commands and results + +After validation succeeds, create one Git commit containing only Task 15 changes. + +Use: + +```text +perf: profile and optimize ROOT-GNN execution +``` + +Before committing, inspect the final diff and ensure no unrelated files or generated profiling artifacts are included. diff --git a/tasks/task16.md b/tasks/task16.md new file mode 100644 index 0000000000000000000000000000000000000000..9f5f600c27d1875c456162b7eecd8a6891f7c838 --- /dev/null +++ b/tasks/task16.md @@ -0,0 +1,2072 @@ +# Task 16: Documentation, Examples, and Migration Closure + +Complete the ROOT-GNN rewrite by making GNN4Colliders understandable, reproducible, and usable without relying on knowledge of the legacy repository. + +This task should not introduce new scientific functionality. + +It should consolidate the architecture, document the supported workflows, provide end-to-end examples, close migration gaps, and clearly separate: + +```text +supported behavior +intentional redesigns +legacy-only behavior +deferred work +``` + +This task builds on: + +```text +Task 6: ROOT/Awkward ingestion +Task 7: metadata-aware caching, splits, batching +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: classification tasks, losses, metrics +Task 10: training lifecycle +Task 11: checkpointing and resume +Task 12: evaluation, prediction, and outputs +Task 13: Hydra configuration and CLI +Task 14: DDP, Perlmutter, and Slurm execution +Task 15: profiling and performance optimization +``` + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +docs/performance.md +``` + +if `docs/performance.md` exists. + +Then inspect: + +```text +configs/ +benchmarks/ +scripts/ +src/gnn4colliders/ +tests/ +legacy/ +``` + +Also inspect the active example/reference configurations introduced in Tasks 13–15. + +Do not assume documentation is correct merely because it already exists. Verify examples against the actual current APIs and CLI. + +--- + +# Goal + +Finish the ROOT-GNN rewrite as a coherent v1-quality workflow. + +A new user should be able to understand and execute: + +```text +install + ↓ +prepare data + ↓ +pretrain ROOT-GNN + ↓ +resume training + ↓ +fine-tune from pretrained backbone + ↓ +evaluate + ↓ +predict + ↓ +run on Perlmutter +``` + +without reading legacy implementation code. + +The final documentation should explain the architecture clearly enough that a future ROOT-Transformer implementation can reuse the shared infrastructure correctly. + +--- + +# 1. Treat this as migration closure + +The primary purpose of this task is not to add more features. + +It is to answer: + +```text +What is implemented? +What is parity-tested? +What changed intentionally? +What remains legacy-only? +What is deferred? +How does a user run the new stack? +``` + +Do not expand scope into new model families or major infrastructure. + +--- + +# 2. Audit documentation against production code + +Before editing documentation, verify the current behavior of: + +```text +CLI commands +Hydra config names +model constructors +task names +checkpoint workflows +inference formats +distributed launch +benchmark commands +``` + +Update docs to match code. + +Do not document aspirational commands that do not work. + +--- + +# 3. README should become the primary entry point + +Rewrite or substantially improve: + +```text +README.md +``` + +so a new user can understand the project quickly. + +Prefer a structure such as: + +```text +GNN4Colliders +Overview +Installation +Quick start +Core concepts +Pretraining +Fine-tuning +Evaluation +Prediction +Checkpoint/resume +Perlmutter +Configuration +Development +Architecture +Migration status +``` + +Keep the README useful but not encyclopedic. + +Link to detailed docs for deeper material. + +--- + +# 4. Project description + +Clearly describe GNN4Colliders as a collider-ML toolkit rather than a permanently GNN-only package. + +For example, communicate the conceptual architecture: + +```text +ROOT + ↓ +EventSample + ↓ +shared collider features + ├── GraphSample -> ROOT-GNN + └── future SequenceSample -> ROOT-Transformer +``` + +Do not imply graph-specific infrastructure is the universal shared representation. + +--- + +# 5. Explain package naming + +Clarify: + +```text +repository: GNN4Colliders +Python package: gnn4colliders +model family: ROOT-GNN +config/module identifier: root_gnn +``` + +Use consistent naming throughout all docs. + +Do not alternate between legacy names casually. + +--- + +# 6. Installation documentation + +Document the canonical development installation. + +Use the actual current project setup. + +For example, if still correct: + +```bash +uv sync --extra root-gnn +``` + +Document: + +```text +supported Python version +core install +ROOT-GNN optional dependencies +development/test dependencies +``` + +Do not document dependency combinations that are no longer validated. + +--- + +# 7. Python version + +Ensure all documentation consistently states the canonical supported Python version/range from `pyproject.toml`. + +Do not leave stale references to legacy Python 3.8 or earlier Task 2 experiments. + +--- + +# 8. DGL / PyTorch environment + +Document the validated ROOT-GNN environment clearly. + +If DGL installation still requires a nonstandard wheel source or environment-specific setup, explain it accurately. + +Do not include obsolete installation workarounds from earlier migration debugging. + +Keep environment-specific caveats concise. + +--- + +# 9. Quick-start workflow + +Add a minimal quick start that demonstrates the current stable flow. + +Conceptually: + +```bash +uv sync --extra root-gnn + +uv run gnn4colliders prepare ... + +uv run gnn4colliders train \ + task=pretraining_multiclass \ + model=root_gnn/edge_network + +uv run gnn4colliders evaluate \ + inference.checkpoint=/path/to/checkpoint.pt +``` + +Use the exact current CLI syntax. + +The example should be executable with documented sample/tiny data where feasible. + +--- + +# 10. Document `prepare` + +Explain: + +```text +ROOT input + -> +EventSample + -> +features + -> +graphs + -> +cache +``` + +Document: + +```text +what is cached +where cache is written +how cache compatibility is checked +what invalidates a cache +``` + +Do not expose implementation details that users do not need. + +--- + +# 11. Document metadata semantics + +Clearly explain that the new architecture does not use public positional: + +```text +tracking[:, N] +``` + +fields. + +Document named metadata concepts such as: + +```text +fold +weight +sample_id +``` + +and any active additional fields. + +Explain that legacy mappings are handled only at compatibility boundaries. + +--- + +# 12. Document `EventSample` + +Describe the architecture-neutral event representation. + +Conceptually: + +```text +EventSample + objects + label + global_features + metadata +``` + +Use the actual current API. + +Make clear this object is shared infrastructure and is not ROOT-GNN-specific. + +--- + +# 13. Document graph construction + +Explain ROOT-GNN graph semantics: + +```text +directed fully connected graph +no self-loops +N * (N - 1) edges +``` + +Document active node feature schema: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +and edge features: + +```text +deta +dphi +dR +``` + +Do not reproduce unnecessary implementation code in the README. + +Link to architecture docs for details. + +--- + +# 14. Document feature semantics + +Document scientifically relevant compatibility behavior such as: + +```text +energy calculation +phi wrapping +feature scaling +object ordering +``` + +where needed for reproducibility. + +Clearly distinguish user-visible scientific contracts from internal implementation details. + +--- + +# 15. Document ROOT-GNN + +Describe: + +```text +EdgeNetwork +``` + +at a conceptual level: + +```text +node/edge/global encoders + -> +iterative edge/node/global message passing + -> +graph representation + -> +classifier + -> +raw logits +``` + +State explicitly that the model returns raw logits rather than sigmoid/softmax outputs. + +--- + +# 16. Document transfer/fine-tuning + +This is a first-class workflow and must be documented clearly. + +Explain: + +```text +pretrained multiclass ROOT-GNN + ↓ +load reusable backbone + ↓ +replace source classifier + ↓ +binary/target-task classifier + ↓ +fine-tune +``` + +Document: + +```text +freeze_backbone=true +freeze_backbone=false +``` + +and their meaning. + +Do not describe transfer learning as resume training. + +--- + +# 17. Distinguish resume vs transfer learning + +Create a prominent explanation. + +```text +Resume: + continue the same training run/task + restore optimizer/scheduler/trainer state + +Transfer: + reuse pretrained model/backbone + create a new task/head + create a new optimizer +``` + +This distinction should appear in both README and checkpoint documentation. + +--- + +# 18. Pretraining example + +Provide one complete example for multiclass pretraining. + +Use the actual current config names. + +Conceptually: + +```bash +uv run gnn4colliders train \ + model=root_gnn/edge_network \ + task=pretraining_multiclass +``` + +Include the important overrides a user is likely to need: + +```text +data path +output path +batch size +epochs +``` + +Keep the example readable. + +--- + +# 19. Fine-tuning example + +Provide one complete transfer-learning example. + +Conceptually: + +```bash +uv run gnn4colliders train \ + model=root_gnn/fine_tuned_edge_network \ + task=binary_classification \ + checkpoint.pretrained=/path/to/pretrained.pt +``` + +Use the actual final configuration layout. + +Also show how to choose: + +```text +freeze_backbone=true +``` + +or: + +```text +freeze_backbone=false +``` + +--- + +# 20. Resume example + +Provide a separate example: + +```bash +uv run gnn4colliders train \ + checkpoint.resume=/path/to/checkpoint.pt +``` + +Use the actual CLI. + +Explain exactly what state is restored. + +--- + +# 21. Evaluation example + +Document: + +```bash +uv run gnn4colliders evaluate ... +``` + +Explain: + +```text +which split is evaluated +how metrics are computed +where event weights are used +how full-split ROC AUC is handled +``` + +Do not duplicate Task 9 formulas unnecessarily in README. + +Link to more detailed docs. + +--- + +# 22. Prediction example + +Document: + +```bash +uv run gnn4colliders predict ... +``` + +Cover: + +```text +NPZ output +ROOT output if supported +labeled vs unlabeled datasets +sample_id preservation +``` + +Explain that inference output uses named fields rather than legacy positional tracking arrays. + +--- + +# 23. Output schema documentation + +Document the primary NPZ output schema. + +For example, if implemented: + +```text +sample_id +logits +scores +predictions +labels +weight +fold +``` + +Use only actual current fields. + +Clearly distinguish required and optional fields. + +--- + +# 24. ROOT output documentation + +If ROOT score writing exists, document: + +```text +tree handling +score branch naming +selection_pass behavior +event alignment +``` + +Do not leave this behavior discoverable only through tests. + +If ROOT writing was intentionally deferred, say so clearly instead. + +--- + +# 25. Checkpoint documentation + +Create or improve a detailed checkpoint section. + +Document: + +```text +new checkpoint schema +schema version +model metadata +task metadata +trainer state +optimizer state +scheduler state +early stopping +``` + +Explain that checkpoints are self-describing to the supported extent. + +--- + +# 26. Legacy checkpoint compatibility + +Document exactly what historical checkpoints are supported. + +For example: + +```text +model-weight loading: supported +pretrained-backbone loading: supported +full legacy optimizer resume: supported / unsupported +module. prefix normalization: supported +_orig_mod. normalization: supported +``` + +Use actual implemented behavior. + +Do not claim compatibility that has not been tested. + +--- + +# 27. Configuration documentation + +Add a dedicated configuration guide if helpful, for example: + +```text +docs/configuration.md +``` + +Explain Hydra config groups: + +```text +data +model +task +trainer +checkpoint +inference +environment +distributed +``` + +Use actual current groups. + +--- + +# 28. Explain configuration philosophy + +Document the rule: + +```text +new experiment -> configuration/YAML +new algorithm or behavior -> Python +``` + +Explain that config files use semantic names rather than arbitrary module/class import paths. + +This is an important departure from the legacy system. + +--- + +# 29. Hydra override examples + +Provide concise examples such as: + +```bash +uv run gnn4colliders train \ + trainer.max_epochs=50 \ + data.batch_size=64 +``` + +and relevant model/task/environment overrides. + +Avoid an enormous override catalog. + +--- + +# 30. Resolved configuration + +Document that runs persist their fully resolved config. + +Explain where it is written. + +This should be presented as part of experiment reproducibility. + +--- + +# 31. Environment profiles + +Document: + +```text +environment=local +environment=perlmutter +``` + +or the actual final names. + +Explain what environment profiles control and what they should not contain. + +Do not encourage users to put scientific model settings in environment configs. + +--- + +# 32. Perlmutter guide + +Add or improve: + +```text +docs/perlmutter.md +``` + +if a dedicated guide is useful. + +Cover: + +```text +environment setup +uv environment +single-GPU run +multi-GPU run +Slurm scripts +output paths +common debugging checks +``` + +Use only validated current commands. + +--- + +# 33. Slurm examples + +Document the scripts under: + +```text +scripts/slurm/ +``` + +Explain which script to use for: + +```text +single GPU +single-node multi-GPU +multi-node +evaluation +benchmarking +``` + +if those scripts exist. + +Do not duplicate the entire script contents in README. + +--- + +# 34. DDP semantics + +Document distributed behavior important to users: + +```text +batch_size is per GPU/process +effective global batch size = batch_size * world_size +num_workers is per process +only rank 0 writes shared artifacts +``` + +Also explain global metric aggregation at a conceptual level. + +--- + +# 35. Reproducibility documentation + +Document the current reproducibility model. + +Cover: + +```text +base seed +DataLoader seed +sampler set_epoch behavior +rank-specific seed derivation +model initialization +CPU determinism +GPU nondeterminism caveats +``` + +Do not promise bitwise GPU determinism unless actually validated. + +--- + +# 36. Performance guide + +Link to: + +```text +docs/performance.md +``` + +or: + +```text +benchmarks/README.md +``` + +Document validated recommendations for: + +```text +num_workers +pin_memory +persistent_workers +precision +DDP +``` + +where Task 15 produced evidence. + +Do not turn hardware-specific benchmark results into universal defaults. + +--- + +# 37. Architecture documentation + +Review and consolidate: + +```text +docs/architecture.md +``` + +It should clearly describe the current architecture. + +At minimum cover: + +```text +ROOT/Awkward I/O +EventSample +shared features +GraphSample +cache +GraphBatch +ROOT-GNN +task layer +trainer +checkpointing +inference +configuration/CLI +distributed execution +``` + +Show responsibility boundaries. + +--- + +# 38. Architecture diagram + +Add a concise textual/Markdown diagram. + +For example: + +```text +ROOT + ↓ +EventSample + ↓ +shared collider features + ↓ +GraphSample + ↓ +cache / split / batch + ↓ +GraphBatch + ↓ +ROOT-GNN + ↓ +logits + ↓ +Task + ├── loss + ├── predictions + └── metrics + ↓ +Trainer / Predictor +``` + +Also show future extensibility: + +```text +shared features + ├── GraphSample -> ROOT-GNN + └── SequenceSample -> ROOT-Transformer +``` + +Do not make ROOT-Transformer appear implemented if it is not. + +--- + +# 39. Responsibility table + +Consider adding a concise table in architecture docs. + +For example: + +```text +Layer Responsibility +data ROOT/Awkward ingestion and samples +features shared collider feature construction +graphs graph representation/topology +models/root_gnn ROOT-GNN architecture +tasks losses/predictions/metrics +training lifecycle/checkpoints +inference prediction/output +distributed DDP execution utilities +cli/config user-facing composition +``` + +Use current code structure. + +--- + +# 40. Migration document + +Turn: + +```text +docs/migration.md +``` + +into a clear migration-status document. + +For each major legacy area, mark: + +```text +migrated +compatibility adapter +intentionally redesigned +deferred +not supported +``` + +Avoid vague status such as "mostly done." + +--- + +# 41. Legacy migration matrix + +Add a matrix for active legacy behavior. + +Include areas such as: + +```text +ROOT ingestion +node features +edge construction +cache +folds +weights +batching +padding +Edge_Network +fine-tuning +loss +metrics +training +early stopping +checkpoints +inference +ROOT outputs +DDP +Slurm +ONNX +``` + +For each, document current status. + +--- + +# 42. Intentional redesigns + +Explicitly document important departures from legacy architecture. + +At minimum include: + +```text +tracking[:, 0] -> metadata.fold +tracking[:, 1] -> metadata.weight +tracking arrays -> named metadata +dynamic YAML module/class imports -> semantic Hydra configs +large training script -> layered Trainer/Task/Checkpoint APIs +persistent graph mutation -> scoped temporary graph state +model constructor global RNG mutation -> explicit external seeding +``` + +Only list changes actually implemented. + +--- + +# 43. Preserve behavior vs preserve bugs + +Migration docs should explain that compatibility targets externally relevant scientific behavior. + +Do not imply every legacy implementation quirk is intentionally retained. + +Classify known legacy quirks as appropriate: + +```text +compatibility requirement +observed behavior +intentional cleanup +known legacy bug +deferred investigation +``` + +--- + +# 44. Outstanding ambiguities + +Review earlier characterization work and document unresolved issues. + +Examples may include: + +```text +negative event weights +rare empty graphs +padding edge cases +legacy validation/test naming +historical checkpoint variants +``` + +Do not silently erase unresolved ambiguities from the migration record. + +If later tasks resolved them, update their status. + +--- + +# 45. Example directory + +Add a small examples area if useful: + +```text +examples/ + README.md + pretraining/ + finetuning/ + inference/ +``` + +Prefer configuration examples over duplicate Python scripts. + +Do not create a parallel application framework under `examples/`. + +--- + +# 46. Example configs + +Provide complete runnable or nearly runnable examples for: + +```text +multiclass pretraining +binary fine-tuning +resume +evaluation +prediction +``` + +Use placeholders only for values that necessarily depend on local data paths. + +Clearly mark placeholders. + +--- + +# 47. Tiny smoke workflow + +If the repository has tiny ROOT fixtures or can safely provide a tiny generated fixture workflow, document a smoke-test path. + +The goal is something like: + +```text +prepare tiny data +train one epoch +save checkpoint +evaluate +predict +``` + +that runs quickly on CPU. + +Prefer reusing test fixtures/tooling rather than committing large binary datasets. + +--- + +# 48. Do not misuse test-only fixtures + +If tests generate their own ROOT fixtures dynamically, do not tell users to depend on internal pytest-only APIs. + +If useful, extract a small supported example-data generator into an appropriate script. + +Only do this if it materially improves onboarding. + +--- + +# 49. Developer guide + +Add or improve a development section. + +Document: + +```bash +uv sync --extra root-gnn +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +and relevant benchmark commands. + +Explain the repository rule: + +```text +one logical task -> one coherent commit +``` + +if this is an established project convention. + +--- + +# 50. Testing guide + +Document the purpose of: + +```text +tests/unit/ +tests/integration/ +tests/parity/ +``` + +Explain that parity tests protect migration behavior. + +Do not list every test file. + +--- + +# 51. Parity philosophy + +Document the rule: + +```text +legacy executable behavior is characterized first +new implementation is compared against deterministic fixtures +``` + +Explain that parity tests are especially important for: + +```text +features +graph topology +model forward +loss semantics +metrics +checkpoint loading +inference outputs +``` + +--- + +# 52. Benchmark guide + +Document how to run Task 15 benchmarks. + +For example: + +```bash +uv run python benchmarks/benchmark_preprocessing.py +uv run python benchmarks/benchmark_training.py --device cuda +``` + +Use actual working commands. + +Explain that benchmark results depend on hardware and configuration. + +--- + +# 53. CLI help audit + +Run: + +```bash +uv run gnn4colliders --help +uv run gnn4colliders prepare --help +uv run gnn4colliders train --help +uv run gnn4colliders evaluate --help +uv run gnn4colliders predict --help +``` + +Review for confusing/stale names. + +Minor help-text improvements are in scope. + +Do not redesign the CLI. + +--- + +# 54. Config naming audit + +Review config names for consistency. + +Remove or rename obviously stale/experimental config names only if doing so does not unnecessarily break established use. + +Prefer: + +```text +pretraining_multiclass +binary_classification +fine_tuning +``` + +or the actual established semantic naming. + +Document aliases if compatibility requires them. + +--- + +# 55. Dead documentation cleanup + +Remove stale documentation that refers to: + +```text +legacy positional tracking APIs +old script entry points +dynamic import YAML +obsolete dependency versions +abandoned architecture plans +``` + +unless explicitly preserved in migration-history sections. + +Do not leave contradictory instructions in different docs. + +--- + +# 56. Docstrings + +Add or improve docstrings for important public APIs where absent. + +Prioritize: + +```text +EventSample +GraphSample +GraphBatch +EdgeNetwork +FineTunedEdgeNetwork +classification tasks +Trainer +CheckpointManager +Predictor +DistributedContext +``` + +Keep docstrings concise. + +Do not write essays inside source files. + +--- + +# 57. Type/API clarity + +While reviewing docs, minor fixes to public annotations or obvious naming inconsistencies are acceptable. + +Do not refactor major APIs in this documentation task. + +If a serious design flaw is discovered, document it as follow-up rather than destabilizing the finished migration. + +--- + +# 58. Public API examples + +Ensure documented imports work. + +For example: + +```python +from gnn4colliders.models.root_gnn import EdgeNetwork +from gnn4colliders.tasks import BinaryClassificationTask +from gnn4colliders.training import Trainer +from gnn4colliders.inference import Predictor +``` + +Use only exports that actually exist. + +Do not teach users to import private modules. + +--- + +# 59. Python API example + +Add one minimal Python example for users who do not want the CLI. + +Conceptually: + +```python +model = ... +task = ... +trainer = ... +``` + +Keep it short. + +The primary supported user workflow can remain CLI/Hydra. + +--- + +# 60. CLI remains primary experiment interface + +Document the intended split: + +```text +CLI/Hydra: + normal experiments + +Python APIs: + libraries, testing, custom workflows +``` + +Do not force all users into notebooks or scripts. + +--- + +# 61. Notebooks + +Do not create a large notebook suite in this task. + +If existing notebooks are stale, either: + +```text +update them +mark them exploratory +remove references to them +``` + +as appropriate. + +Notebooks should not be the canonical production workflow. + +--- + +# 62. Legacy directory documentation + +Explain the purpose of: + +```text +legacy/ +``` + +Clearly state that it is a frozen behavioral reference and not production implementation. + +Do not tell users to add `legacy/` to `PYTHONPATH`. + +--- + +# 63. Legacy old rewrite attempt + +If: + +```text +legacy/physicsnemo/ +``` + +or an equivalent abandoned rewrite remains, document clearly that it is not the behavioral target. + +Do not accidentally present it as a supported backend. + +--- + +# 64. Avoid deleting legacy in Task 16 + +Do not delete the legacy implementation yet. + +It remains useful for: + +```text +parity tests +historical reference +checkpoint compatibility investigation +``` + +Legacy cleanup/deletion belongs to a later dedicated task. + +--- + +# 65. Create a migration completion checklist + +At the end of `docs/migration.md`, include a concise checklist. + +For example: + +```text +[x] active data path +[x] graph construction +[x] EdgeNetwork +[x] transfer learning +[x] task losses/metrics +[x] trainer +[x] checkpoints +[x] inference +[x] CLI +[x] DDP +[x] profiling +[ ] ONNX export +[ ] ROOT-Transformer +``` + +Use actual status. + +Do not mark something complete solely because there is placeholder code. + +--- + +# 66. Define ROOT-GNN v1 completion + +Add a concise statement describing what constitutes the completed ROOT-GNN rewrite. + +For example: + +```text +ROOT-GNN v1 is complete when the supported new stack can: +- prepare active datasets +- reproduce validated legacy feature/graph/model/task behavior +- train from scratch +- fine-tune pretrained backbones +- resume checkpoints +- evaluate/predict +- run single-process and DDP on Perlmutter +``` + +Tailor this to actual implemented functionality. + +--- + +# 67. Known limitations + +Add a clear known-limitations section. + +Possible examples, if still true: + +```text +ONNX export deferred +distributed ROOT writing limited +full historical optimizer-resume compatibility incomplete +very large inference uses in-memory result accumulation +ROOT-Transformer not yet implemented +``` + +Do not hide limitations. + +--- + +# 68. Future roadmap + +Keep the roadmap short. + +After ROOT-GNN v1, likely items include: + +```text +ONNX/export +legacy cleanup +packaging/release hardening +ROOT-Transformer +large-scale inference improvements +``` + +Do not turn README into a speculative roadmap document. + +--- + +# 69. Validate every command in documentation + +For every command intended to be runnable, execute it or at least its safe smoke equivalent. + +Especially validate: + +```text +installation +--help +config composition +tiny prepare +tiny train +tiny fine-tune +tiny evaluate +tiny predict +benchmark command syntax +``` + +Do not ship copied commands that have not been checked. + +--- + +# 70. Documentation tests + +If helpful, add lightweight tests to protect key docs/API snippets. + +For example: + +```text +CLI help smoke tests +example config composition tests +public import tests +``` + +Do not introduce a heavyweight documentation-testing framework unless already used. + +--- + +# 71. Config smoke tests + +Validate all reference configs compose successfully. + +For each committed reference configuration, ensure Hydra can resolve it without performing expensive training. + +Fail on stale field names. + +--- + +# 72. Example smoke tests + +If example configs/scripts are added, add small tests ensuring they do not rot. + +Keep these tests fast. + +Do not require production ROOT files. + +--- + +# 73. Documentation consistency + +Use consistent terminology everywhere. + +Prefer: + +```text +training +validation +test +``` + +rather than reproducing confusing legacy naming. + +Prefer: + +```text +metadata.weight +metadata.fold +``` + +rather than tracking columns. + +Prefer: + +```text +fine-tuning +``` + +or one consistent spelling throughout. + +--- + +# 74. Scientific terminology + +Preserve established collider terminology accurately. + +Do not rename scientifically meaningful fields merely for stylistic reasons. + +Where names have historical meanings, explain them. + +--- + +# 75. Code comments + +Do not compensate for unclear docs by adding excessive inline comments everywhere. + +Use source comments only for non-obvious implementation constraints. + +Documentation belongs primarily in README/docs and concise public docstrings. + +--- + +# 76. No new major dependencies + +This task should add no major runtime dependency. + +Documentation tooling dependencies are also unnecessary unless the project already has a docs build system. + +Plain Markdown is sufficient. + +--- + +# 77. Do not build a documentation website + +Do not introduce: + +```text +Sphinx +MkDocs +Docusaurus +``` + +unless already present and clearly desired. + +The task is content completion, not documentation-platform migration. + +--- + +# 78. Review `.gitignore` + +Ensure generated artifacts from documented workflows are not accidentally committed. + +Check handling of: + +```text +outputs/ +checkpoints/ +predictions/ +profiles/ +benchmark-results/ +__pycache__/ +*.egg-info/ +``` + +Only change `.gitignore` where appropriate. + +Do not ignore source example configs or benchmark scripts. + +--- + +# 79. Repository cleanliness + +Remove accidental generated development artifacts if they are tracked or unignored, such as: + +```text +__pycache__/ +src/gnn4colliders.egg-info/ +temporary outputs +profiling traces +``` + +Do not delete user data. + +Only clean repository-generated artifacts. + +--- + +# 80. Architecture-neutral audit + +Use this task to verify the shared design still supports a future non-graph model family. + +Review especially: + +```text +data/ +features/ +training/ +tasks/ +inference/ +distributed/ +config/ +``` + +for unnecessary ROOT-GNN-specific naming or assumptions. + +Do not perform major refactors. + +Document any remaining coupling as follow-up work. + +--- + +# 81. ROOT-Transformer readiness note + +Add a concise architecture note for future development. + +Conceptually: + +```text +ROOT + -> +EventSample + -> +shared feature representation + ├── graph adapter -> ROOT-GNN + └── future sequence adapter -> ROOT-Transformer +``` + +Identify the intended extension points. + +Do not implement the transformer here. + +--- + +# 82. Migration debt audit + +Search for remaining new-code references to legacy concepts such as: + +```text +tracking +tracking_info +sys.path +legacy module/class dynamic import +hardcoded ROOT-GNN assumptions in shared code +``` + +Classify each result: + +```text +required compatibility +test-only +documentation/history +migration debt +``` + +Fix only obvious low-risk leftovers. + +Document the rest. + +--- + +# 83. TODO/FIXME audit + +Search: + +```text +TODO +FIXME +HACK +XXX +``` + +through production code. + +For each: + +* resolve it if trivial and in scope +* convert it into a clear documented limitation +* remove stale notes +* leave substantial future work for later tasks + +Do not turn Task 16 into an unbounded cleanup sprint. + +--- + +# 84. Warning/error message audit + +Review major user-facing failures. + +Examples: + +```text +invalid config +missing input files +incompatible cache +incompatible checkpoint +wrong model/task combination +unsupported output format +distributed misconfiguration +``` + +Improve obviously cryptic messages where low-risk. + +Do not redesign exception architecture. + +--- + +# 85. Final end-to-end smoke workflow + +Create one documented validation workflow that exercises the new system from start to finish using tiny inputs. + +Conceptually: + +```text +prepare + ↓ +train multiclass model briefly + ↓ +save checkpoint + ↓ +load pretrained backbone + ↓ +fine-tune binary model briefly + ↓ +evaluate + ↓ +predict NPZ +``` + +This does not need scientific convergence. + +Its purpose is integration confidence. + +--- + +# 86. Do not use legacy code in the smoke workflow + +The final new-stack smoke workflow must run without importing production logic from: + +```text +legacy/ +``` + +Parity tests may continue to use legacy. + +The actual user workflow must not. + +--- + +# 87. Optional end-to-end script + +If useful, add something like: + +```text +scripts/dev/smoke_end_to_end.sh +``` + +or: + +```text +scripts/dev/smoke_end_to_end.py +``` + +Keep it thin and built entirely on the public CLI. + +Do not duplicate application logic. + +--- + +# 88. Smoke script requirements + +A smoke script should: + +* use temporary/tiny data +* avoid hardcoded user paths +* fail on command errors +* be readable +* not require Slurm +* not require GPU + +Do not make it part of normal package runtime. + +--- + +# 89. Release-readiness checklist + +Add a concise internal checklist to docs if useful. + +For example: + +```text +tests pass +parity tests pass +CLI examples validated +reference configs compose +README current +migration matrix current +no generated artifacts tracked +``` + +Do not introduce formal release tooling yet. + +--- + +# 90. Testing + +Run: + +```bash +uv run pytest +``` + +Run focused parity: + +```bash +uv run pytest tests/parity -v +``` + +Run distributed tests: + +```bash +uv run pytest tests/unit/distributed -v +uv run pytest tests/integration -k distributed -v +``` + +Run lint/format: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +Validate public imports. + +Validate CLI: + +```bash +uv run gnn4colliders --help +uv run gnn4colliders prepare --help +uv run gnn4colliders train --help +uv run gnn4colliders evaluate --help +uv run gnn4colliders predict --help +``` + +Validate all committed reference configs compose. + +Validate the documented tiny end-to-end workflow. + +If Slurm scripts are present: + +```bash +bash -n scripts/slurm/*.sh +``` + +or the appropriate individual scripts. + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* no scientific behavior changed unexpectedly +* no production imports from legacy were introduced +* docs match current CLI/config APIs +* no stale tracking-array usage is presented as public API +* no generated artifacts are tracked +* no major new dependency was added +* no unrelated refactor is included + +--- + +# Completion criteria + +Task 16 is complete when: + +1. README accurately describes the current project. +2. Installation instructions are validated. +3. Project/package/model naming is consistent. +4. The quick-start workflow is documented. +5. Data preparation is documented. +6. `EventSample` and named metadata concepts are documented. +7. ROOT-GNN graph semantics are documented. +8. ROOT-GNN model behavior is documented. +9. Transfer/fine-tuning is documented as a first-class workflow. +10. Resume and transfer learning are clearly distinguished. +11. Pretraining example is validated. +12. Fine-tuning example is validated. +13. Resume example is validated. +14. Evaluation example is validated. +15. Prediction example is validated. +16. Checkpoint format/workflows are documented. +17. Supported legacy checkpoint compatibility is documented accurately. +18. Hydra configuration groups and overrides are documented. +19. Environment/Perlmutter execution is documented. +20. DDP batch-size and worker semantics are documented. +21. Reproducibility behavior is documented accurately. +22. Performance guidance from Task 15 is documented. +23. `docs/architecture.md` reflects the final implementation. +24. `docs/migration.md` contains a clear status matrix. +25. Intentional redesigns from legacy are documented. +26. Remaining ambiguities/limitations are explicit. +27. Future ROOT-Transformer extension points are documented without implementing it. +28. Reference configs compose successfully. +29. Public imports shown in docs work. +30. CLI help commands work. +31. Tiny end-to-end new-stack smoke workflow succeeds. +32. Full unit/integration/parity suite passes. +33. Distributed tests still pass. +34. Lint/format checks pass. +35. Legacy code remains frozen. +36. No major new feature or dependency has been introduced. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. README structure +4. installation instructions finalized +5. quick-start workflow +6. pretraining example +7. fine-tuning example +8. resume example +9. evaluation example +10. prediction example +11. output schema documentation +12. checkpoint documentation +13. Hydra/configuration documentation +14. Perlmutter/DDP documentation +15. reproducibility documentation +16. performance documentation +17. architecture-document updates +18. migration matrix summary +19. intentional legacy redesigns documented +20. unresolved migration items +21. architecture-neutrality audit findings +22. ROOT-Transformer extension points +23. stale/dead documentation removed +24. generated repository artifacts cleaned +25. TODO/FIXME audit results +26. public import validation +27. CLI help validation +28. reference config validation +29. end-to-end smoke-test results +30. full test results +31. lint/format results +32. any follow-up tasks identified + +After validation succeeds, create one Git commit containing only Task 16 changes. + +Use: + +```text +docs: finalize ROOT-GNN documentation and migration +``` + +Before committing, inspect the final diff and ensure no unrelated implementation changes or generated artifacts are included. diff --git a/tasks/task17.md b/tasks/task17.md new file mode 100644 index 0000000000000000000000000000000000000000..1cd0fa4038b84005b6d3fdd95da286546ecbb00e --- /dev/null +++ b/tasks/task17.md @@ -0,0 +1,1899 @@ +# Task 17: Implement ONNX Export and Export-Parity Validation + +Implement a clean model-export path for GNN4Colliders, focused first on exporting supported ROOT-GNN inference models to ONNX and validating exported-model parity against native PyTorch inference. + +This task builds on: + +```text id="70ycbf" +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: classification tasks and output semantics +Task 11: checkpoint loading +Task 12: evaluation and inference +Task 13: Hydra configuration and CLI +Task 14: distributed execution +Task 16: documentation and migration closure +``` + +Before making changes, read: + +```text id="qt0f76" +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text id="g7fohc" +src/gnn4colliders/models/root_gnn/ +src/gnn4colliders/tasks/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +src/gnn4colliders/cli/ +src/gnn4colliders/graphs/ + +configs/ +tests/unit/ +tests/integration/ +tests/parity/ + +legacy/root_gnn_dgl/ +``` + +Also inspect whether the legacy repository contains an active ONNX/export script and characterize only the export behavior actually used. + +Do not redesign the ROOT-GNN architecture merely to make ONNX export easier. + +--- + +# Goal + +Provide a robust export path: + +```text id="ww6x60" +checkpoint + ↓ +reconstruct ROOT-GNN model + ↓ +adapt graph/model inputs into exportable representation + ↓ +ONNX export + ↓ +ONNX Runtime inference + ↓ +compare against PyTorch +``` + +Support at minimum: + +```text id="de5e55" +ROOT-GNN pretrained multiclass model +ROOT-GNN fine-tuned binary model +``` + +The export path must preserve inference semantics. + +Do not treat successful file creation as sufficient. + +Export is only considered successful when ONNX Runtime outputs match native PyTorch outputs within justified tolerances. + +--- + +# 1. Add an export package + +Prefer a small structure such as: + +```text id="f5m7kt" +src/gnn4colliders/export/ + __init__.py + onnx.py + adapters.py +``` + +or a similarly compact organization. + +Possible responsibilities: + +```text id="re1n1w" +onnx.py + export/load/validate ONNX model + +adapters.py + convert native ROOT-GNN graph inputs into an export-friendly tensor interface +``` + +Do not create a generic deployment framework. + +--- + +# 2. Keep export separate from normal model execution + +The normal ROOT-GNN model should continue consuming its established graph/batch representation. + +Do not rewrite production inference around ONNX constraints. + +If an export-specific wrapper is needed, isolate it. + +Conceptually: + +```text id="pnbm0r" +native EdgeNetwork + ↓ +ExportAdapter + ↓ +tensor-only forward signature + ↓ +ONNX +``` + +The adapter exists for export only. + +--- + +# 3. Characterize DGL export constraints first + +Before writing an adapter, determine whether the current DGL model can be exported directly. + +Do not assume DGL graph operations are ONNX-exportable. + +Perform a small proof-of-concept. + +If direct export fails or produces unsupported operators, document that finding and implement a tensor-based export adapter. + +Do not keep a fragile direct-DGL export path merely because a file can be generated. + +--- + +# 4. Tensor-based graph representation + +If needed, define an explicit export representation. + +Conceptually: + +```text id="2dqa9s" +node_features +edge_features +edge_src +edge_dst +global_features +graph_index / batch_index +``` + +Use the minimum tensors required to reproduce Task 8 message passing. + +The export representation must preserve: + +```text id="we9d1w" +node ordering +edge ordering +graph membership +global-feature association +``` + +Do not encode hidden positional assumptions beyond those already validated. + +--- + +# 5. Export adapter API + +Prefer something conceptually like: + +```python id="kz3uy0" +adapter = RootGNNExportAdapter(model) + +logits = adapter( + node_features, + edge_features, + edge_src, + edge_dst, + graph_index, + global_features, +) +``` + +The exact signature may differ. + +Keep it explicit and tensor-only where practical. + +Do not expose `EventMetadata` to the exported model. + +--- + +# 6. Metadata is not model input + +Do not export: + +```text id="dbv0zk" +fold +weight +sample_id +source_file +``` + +as model inputs unless the model actually consumes them. + +Task/event metadata remains outside the ONNX computational graph. + +--- + +# 7. Preserve raw-logit output + +The exported model should return the same raw logits as native PyTorch. + +Do not include: + +```text id="oip3wc" +sigmoid +softmax +thresholding +argmax +``` + +inside the ONNX model unless the established production model itself includes them. + +Task 9 remains the owner of postprocessing semantics. + +--- + +# 8. Export pretrained multiclass model + +Support exporting the active multiclass ROOT-GNN model. + +Verify: + +```text id="kkvpgz" +output shape = [batch_size, num_classes] +``` + +Use the active model config from the checkpoint. + +Do not hardcode the historical class count into generic export code. + +--- + +# 9. Export fine-tuned binary model + +Support exporting: + +```text id="cr9zyu" +FineTunedEdgeNetwork +``` + +Verify: + +```text id="qq26fr" +output shape = [batch_size, 1] +``` + +The exported model must include the transferred backbone plus the replacement classifier. + +Do not require the source pretraining classifier. + +--- + +# 10. Checkpoint integration + +Use Task 11 model reconstruction/loading APIs. + +Conceptually: + +```text id="15bqxe" +checkpoint + ↓ +reconstruct model + ↓ +load weights + ↓ +model.eval() + ↓ +export +``` + +Do not duplicate checkpoint parsing in export code. + +--- + +# 11. Export from model object + +Also support exporting an already-constructed model where practical. + +Conceptually: + +```python id="d7kmlu" +export_root_gnn_onnx( + model=model, + example_batch=batch, + output_path=path, +) +``` + +This makes tests easier and avoids coupling export strictly to disk checkpoints. + +--- + +# 12. Export from checkpoint convenience path + +Provide a higher-level convenience API if useful: + +```python id="j37f3n" +export_checkpoint_to_onnx( + checkpoint_path, + output_path, + example_batch=..., +) +``` + +This should reuse the lower-level export path. + +Avoid two independent implementations. + +--- + +# 13. Example input construction + +ONNX export usually requires representative input tensors. + +Provide an explicit helper that converts: + +```text id="pkyvh2" +GraphBatch +``` + +into: + +```text id="1bg1fm" +ExportInputs +``` + +or equivalent. + +Do not make the export code reconstruct graphs from ROOT files. + +The export boundary begins from already-prepared graph/batch data. + +--- + +# 14. Export input dataclass + +If useful, add: + +```python id="rocyja" +@dataclass(frozen=True) +class RootGNNExportInputs: + node_features: torch.Tensor + edge_features: torch.Tensor + edge_src: torch.Tensor + edge_dst: torch.Tensor + graph_index: torch.Tensor + global_features: torch.Tensor | None +``` + +Adjust fields to the actual algorithm. + +Keep this internal unless it has real user value. + +--- + +# 15. Batching semantics + +Support batched graphs if practical. + +The adapter must reproduce DGL batching semantics exactly. + +Explicitly handle: + +```text id="i3bbf1" +which node belongs to which graph +which edge belongs to which graph +per-graph pooling +global features per graph +``` + +Do not infer graph boundaries from edge ordering alone if a robust explicit representation exists. + +--- + +# 16. Graph membership representation + +Use a stable tensor representation for graph membership. + +For example: + +```text id="hi8tqm" +node_batch_index +edge_batch_index +``` + +or offsets such as: + +```text id="s3b23o" +node_splits +edge_splits +``` + +Choose the representation that best matches export/runtime compatibility. + +Document it. + +--- + +# 17. Dynamic graph sizes + +Support dynamic: + +```text id="w3f4xn" +number of nodes +number of edges +batch size +``` + +where the chosen ONNX exporter/runtime supports it reliably. + +Use dynamic axes or symbolic dimensions as appropriate. + +Do not lock the exported model to one tiny example graph unless unavoidable. + +--- + +# 18. Dynamic-axis validation + +Test the same ONNX file on at least: + +```text id="3g043o" +one graph size used for export +a different node/edge count +a different batch size +``` + +where supported. + +Successful inference on the export example alone is not enough. + +--- + +# 19. Edge ordering + +Preserve the graph edge order validated in Task 5. + +The tensor adapter should derive: + +```text id="9obebd" +edge_src +edge_dst +edge_features +``` + +consistently. + +Do not reorder edges during export conversion unless proven numerically irrelevant and explicitly documented. + +--- + +# 20. Node aggregation semantics + +Reproduce the Task 8 aggregation exactly. + +If native code uses DGL sum aggregation, the export adapter must implement equivalent tensor operations. + +Potential mechanisms include: + +```text id="akrzpu" +index_add +scatter-like tensor operations +segment reductions +``` + +Use operators that export cleanly to ONNX. + +Do not change sum aggregation to mean/max. + +--- + +# 21. Global pooling semantics + +Reproduce native: + +```text id="65w0yl" +mean pooling of node representations +mean pooling of edge representations +global update +``` + +or the actual Task 8 behavior exactly. + +Pay special attention to varying graph sizes. + +Do not compute a mean over the entire batch. + +Pooling must remain per graph. + +--- + +# 22. Empty graphs + +Characterize whether empty-node or empty-edge graphs are supported by the current production pipeline. + +If they are not valid inputs, validate and reject them before export/inference. + +If they are supported, add explicit tests. + +Do not invent unsupported empty-graph semantics solely for ONNX. + +--- + +# 23. Single-node graph + +Task 5 established that a single-node fully connected graph has: + +```text id="wsdhnc" +1 node +0 edges +``` + +Ensure the export adapter either supports this correctly or rejects it with a clear documented limitation. + +Do not assume every graph has at least one edge. + +--- + +# 24. Exportable operations + +Prefer common ONNX-supported tensor operations. + +Do not introduce custom ONNX ops for the first implementation. + +Avoid deployment environments that require custom DGL operators. + +The exported graph should be as portable as practical. + +--- + +# 25. ONNX opset + +Choose an explicit supported ONNX opset. + +Do not leave it implicit. + +Select the lowest practical opset that supports the required tensor operations and current exporter/runtime. + +Document the chosen opset and why. + +--- + +# 26. Exporter API + +Use the current supported PyTorch ONNX export path available in the project environment. + +Do not depend on obsolete exporter APIs if the current PyTorch stack provides a stable newer path. + +Keep exporter-specific logic localized in: + +```text id="lb6wc8" +export/onnx.py +``` + +--- + +# 27. ONNX dependency management + +Add ONNX tooling as an optional export/development dependency rather than forcing every GNN4Colliders installation to include it. + +Prefer an extra such as: + +```text id="deu91j" +export +``` + +or: + +```text id="5em87q" +onnx +``` + +Use the naming convention that best matches the existing `pyproject.toml`. + +Possible dependencies: + +```text id="e8gnwn" +onnx +onnxruntime +``` + +Add only what is actually required. + +--- + +# 28. CPU ONNX Runtime validation + +The automated test suite should validate exported models using CPU ONNX Runtime. + +Do not require CUDA ONNX Runtime in normal tests. + +CUDA Runtime benchmarking may be optional/manual. + +--- + +# 29. ONNX model validation + +After export, run ONNX structural validation where supported. + +For example: + +```text id="34m6h9" +load model +checker validation +``` + +Fail clearly if the graph is invalid. + +Do not assume a successfully written file is valid. + +--- + +# 30. PyTorch vs ONNX parity + +This is the primary correctness requirement. + +For the same deterministic inputs: + +```text id="jx7lus" +PyTorch model + -> +reference logits + +ONNX Runtime + -> +exported logits +``` + +Compare outputs. + +Use appropriate tolerances. + +For float32 CPU inference, start with tight tolerances and loosen only if justified by operator differences. + +--- + +# 31. Multiclass parity test + +Add an export parity test for active multiclass pretraining. + +Use: + +```text id="eclkeo" +fixed model weights +deterministic graph inputs +multiple graphs if possible +``` + +Verify ONNX logits match PyTorch. + +--- + +# 32. Fine-tuning parity test + +Add parity for the binary fine-tuned model. + +Verify transferred backbone plus new classifier matches native PyTorch. + +Test both construction from: + +```text id="bo2ura" +live model +checkpoint +``` + +where practical. + +--- + +# 33. Multiple processing steps + +Ensure parity covers: + +```text id="06jwtq" +processing_steps > 1 +``` + +so the export graph actually exercises repeated message passing. + +Do not validate only a trivial zero/one-step configuration. + +--- + +# 34. Global-feature parity + +If global features are active, export and test them. + +Test: + +```text id="mj6a6v" +with global features +``` + +and, if supported by the current model: + +```text id="uv9157" +without global features +``` + +Do not invent dummy globals to avoid handling the real interface. + +--- + +# 35. Dropout/eval semantics + +Export only in evaluation mode. + +Ensure: + +```python id="rhp02x" +model.eval() +``` + +before export. + +Dropout must behave as inference, not training. + +Add a repeated native/export parity test if useful. + +--- + +# 36. Do not export training + +Do not export: + +```text id="jzct8f" +loss +backward +optimizer +scheduler +training loop +``` + +The ONNX graph is inference-only. + +--- + +# 37. Task postprocessing stays outside ONNX + +Do not include Task 9 metrics or weighting. + +For a binary exported model: + +```text id="z73v2h" +ONNX -> logits +Python task -> sigmoid / threshold +``` + +For multiclass: + +```text id="yyqjek" +ONNX -> logits +Python task -> scores / argmax +``` + +This preserves separation of concerns. + +--- + +# 38. Optional postprocessing helper + +It is acceptable to provide an inference helper that takes ONNX logits and applies the existing Task 9 Python task logic. + +Do not create a second independent definition of score semantics. + +--- + +# 39. ONNX inference wrapper + +Add a lightweight runtime wrapper if useful. + +Conceptually: + +```python id="zis9l9" +runner = OnnxPredictor(path) + +logits = runner.predict(export_inputs) +``` + +Keep it small. + +Do not duplicate the full Task 12 `Predictor`. + +--- + +# 40. Runtime wrapper responsibilities + +An ONNX runtime helper may own: + +```text id="9r6aqc" +loading session +mapping tensor names +NumPy conversion +running inference +returning logits +``` + +It should not own: + +```text id="0fjm8y" +ROOT reading +graph construction +task metrics +checkpoint handling +``` + +--- + +# 41. Input/output names + +Use stable human-readable ONNX tensor names. + +For example: + +```text id="yaw4v1" +node_features +edge_features +edge_src +edge_dst +node_batch +edge_batch +global_features +logits +``` + +Use the actual final representation. + +Do not expose meaningless generated names when explicit names are easy to provide. + +--- + +# 42. Export metadata + +Store useful metadata alongside the exported ONNX model. + +Possible approaches: + +```text id="btxb1f" +ONNX model metadata +sidecar JSON +``` + +At minimum consider: + +```text id="lmndxr" +model family +checkpoint schema version +feature schema version +graph schema version +model config +task type +ONNX opset +GNN4Colliders version +``` + +Keep it compact. + +--- + +# 43. Sidecar metadata + +If ONNX's metadata API is insufficient for structured values, write: + +```text id="2j9vch" +model.onnx +model.onnx.json +``` + +or equivalent. + +Do not serialize large model tensors twice. + +--- + +# 44. Export provenance + +If exporting from a checkpoint, record a checkpoint identifier/path or digest as informational provenance. + +Do not make absolute local filesystem paths mandatory for model use. + +A basename or optional hash may be preferable. + +--- + +# 45. Schema compatibility + +Export metadata should include: + +```text id="btmtr1" +feature schema version +graph schema version +``` + +so future consumers can detect incompatible preprocessing. + +Do not expect ONNX itself to construct collider features from raw ROOT branches. + +--- + +# 46. Exported-model contract + +Document clearly that the exported ROOT-GNN expects already-constructed numerical graph tensors. + +The deployment flow is: + +```text id="tfxfjz" +raw event + ↓ +GNN4Colliders preprocessing + ↓ +export tensor representation + ↓ +ONNX model +``` + +unless a future task exports preprocessing too. + +Do not imply `model.onnx` accepts ROOT files. + +--- + +# 47. No preprocessing export + +Do not export: + +```text id="s6l78d" +ROOT reading +Awkward transforms +object selection +feature scaling from raw branches +graph construction from raw physics objects +``` + +into ONNX in this task. + +Keep the scope to the trained neural network computation. + +--- + +# 48. CLI `export` command + +Now add: + +```bash id="fs047y" +gnn4colliders export +``` + +through the Task 13 CLI. + +Keep it thin. + +Conceptually: + +```bash id="zqlqrr" +uv run gnn4colliders export \ + export.checkpoint=/path/to/model.pt \ + export.output=model.onnx +``` + +Use the final Hydra/config conventions. + +--- + +# 49. Export config group + +Add: + +```text id="8e3w0v" +configs/export/ + onnx.yaml +``` + +or equivalent. + +Possible fields: + +```yaml id="o9pmef" +export: + format: onnx + checkpoint: null + output: model.onnx + opset: ... +``` + +Only expose meaningful stable options. + +--- + +# 50. Example input for CLI export + +The CLI needs a representative graph/batch to establish export shapes. + +Choose a clean explicit approach. + +Possibilities include: + +```text id="ef7hnb" +use a configured dataset and first batch +use a saved export-input fixture +accept dimensions plus a generated valid dummy graph +``` + +Prefer using the actual configured preprocessing/data pipeline when practical, because it validates the real input representation. + +Do not silently export using a tiny shape that cannot generalize. + +--- + +# 51. `export` command responsibility + +The CLI should: + +```text id="u7ykk2" +resolve config +load checkpoint +construct/reuse example batch +convert to export inputs +export ONNX +validate ONNX +optionally run parity check +write metadata +``` + +Actual export math stays in `gnn4colliders.export`. + +--- + +# 52. CLI parity validation + +By default, after exporting, run at least one PyTorch-vs-ONNX comparison unless this would be prohibitively expensive. + +Prefer failing export if parity validation fails. + +If a skip option exists, make it explicit. + +Do not silently produce an unvalidated export. + +--- + +# 53. Output overwrite behavior + +Do not overwrite an existing ONNX file unexpectedly. + +Use the project’s existing output policy. + +Either: + +```text id="h8k3q5" +refuse +require explicit overwrite +write unique path +``` + +Keep this non-interactive for batch jobs. + +--- + +# 54. Export on CPU + +The baseline export workflow should work on CPU. + +A user should not need a GPU merely to convert a checkpoint to ONNX. + +Map the model/example tensors appropriately. + +--- + +# 55. Export GPU-trained checkpoints on CPU + +Checkpoints created on CUDA should be exportable on a CPU machine using Task 11 `map_location` behavior where architecture dependencies permit. + +Add a test where practical. + +--- + +# 56. ONNX Runtime device independence + +Automated parity should use CPU ONNX Runtime. + +Do not assume a CUDA execution provider. + +Document optional CUDA runtime use separately only if validated. + +--- + +# 57. Dynamic shapes vs portability + +Prefer an export that supports realistic variable graph sizes. + +However, if specific ONNX operators/exporter limitations prevent fully dynamic shapes, document exact constraints. + +Do not pretend a fixed-shape model is dynamic. + +If necessary, support a clearly named fixed-shape mode only as a fallback. + +--- + +# 58. Avoid padding redesign + +Do not reintroduce legacy fixed: + +```text id="ygwyq3" +16000 nodes +104000 edges +``` + +padding merely to simplify export. + +Use the clean new graph representation. + +Only use fixed padding if a real downstream deployment target requires it and this is explicitly documented. + +--- + +# 59. ONNX graph inspection + +Inspect the exported graph. + +Verify it does not accidentally include: + +```text id="bmbynj" +training-only branches +constant example-specific graph topology +hardcoded example batch size +hardcoded example node count +``` + +unless those dimensions are intentionally fixed. + +Add tests for dynamic behavior rather than relying only on manual inspection. + +--- + +# 60. Operator compatibility + +Record the major ONNX operators needed by the exported model. + +Do not manually optimize the ONNX graph in this task. + +If an operator is poorly supported by the target runtime, adjust the adapter using equivalent standard tensor operations where possible. + +--- + +# 61. Avoid custom scatter dependencies if possible + +If aggregation requires scatter behavior, prefer ONNX-exportable native PyTorch operations. + +Do not add PyTorch Scatter solely for export unless necessary. + +Keep runtime dependencies minimal. + +--- + +# 62. Performance is secondary to parity + +Benchmark ONNX inference if useful, but do not optimize the adapter solely for speed in this task. + +The primary goal is: + +```text id="ph78rw" +portable export + correctness +``` + +not beating native DGL inference. + +--- + +# 63. Optional ONNX benchmark + +If straightforward, add a small benchmark: + +```text id="8qp54e" +benchmarks/benchmark_onnx.py +``` + +Compare: + +```text id="ucg05t" +PyTorch CPU inference +ONNX Runtime CPU inference +``` + +for representative inputs. + +Do not make performance claims without measurements. + +--- + +# 64. Unit tests + +Add focused tests under: + +```text id="bu7qgo" +tests/unit/export/ +``` + +Suggested coverage: + +```text id="sjk6ld" +test_export_inputs.py +test_export_adapter.py +test_onnx_export.py +test_onnx_runtime.py +``` + +Use fewer files if clearer. + +--- + +# 65. Export-input conversion tests + +Verify conversion from `GraphBatch` preserves: + +```text id="7oxvoc" +node features +edge features +edge src/dst +graph membership +global features +graph count +``` + +Compare directly against the native DGL graph. + +--- + +# 66. Adapter native parity + +Before ONNX export, compare: + +```text id="r946qe" +native DGL EdgeNetwork +``` + +against: + +```text id="13amoi" +tensor-only ExportAdapter +``` + +using the same weights and graph. + +This isolates adapter correctness from ONNX exporter issues. + +Require tight parity. + +--- + +# 67. Adapter multiclass test + +Test multiclass native-vs-adapter parity. + +Use multiple processing steps and multiple graphs. + +--- + +# 68. Adapter fine-tuning test + +Test fine-tuned binary native-vs-adapter parity. + +Verify the new classifier is represented correctly. + +--- + +# 69. ONNX file smoke test + +Export a tiny valid model. + +Verify: + +```text id="yd7du1" +file exists +model loads +ONNX checker succeeds +runtime session initializes +``` + +Do not stop there; run numerical parity too. + +--- + +# 70. ONNX numerical parity test + +Compare native PyTorch and ONNX Runtime logits. + +Use deterministic fixed inputs. + +Test at least: + +```text id="4umw7v" +one graph +multiple graphs +different graph size +``` + +where dynamic export supports them. + +--- + +# 71. Batch-size dynamic test + +If batch dimension is declared dynamic, test: + +```text id="ws79d7" +batch size 1 +batch size > 1 +``` + +with the same exported model. + +--- + +# 72. Node-count dynamic test + +If node count is dynamic, test a different number of nodes from the export example. + +Because edges scale as: + +```text id="46vfip" +N * (N - 1) +``` + +this also validates dynamic edge count. + +--- + +# 73. Single-node test + +Where supported, validate: + +```text id="l6zhtf" +N = 1 +E = 0 +``` + +through: + +```text id="b5vlcx" +native model +adapter +ONNX Runtime +``` + +If ONNX cannot support this case cleanly, validate the rejection path and document it. + +--- + +# 74. Checkpoint export integration test + +Add: + +```text id="jvc91c" +model + -> +Task 11 checkpoint + -> +fresh export load + -> +ONNX + -> +runtime logits +``` + +Compare against the original model. + +This proves the user-facing workflow, not just direct model export. + +--- + +# 75. Fine-tuned checkpoint export integration + +Explicitly test: + +```text id="4zjsai" +multiclass pretrained backbone + -> +fine-tuned binary checkpoint + -> +export + -> +ONNX binary logits +``` + +This is a required real workflow. + +--- + +# 76. CLI export integration test + +Exercise the final CLI with temporary files. + +Conceptually: + +```text id="z95kom" +checkpoint + -> +gnn4colliders export + -> +model.onnx + -> +ONNX Runtime validation +``` + +Keep the model/data tiny. + +Do not require GPU. + +--- + +# 77. Metadata sidecar test + +If metadata is written, verify: + +```text id="a8l6zc" +model family +model/task config +schema versions +opset +``` + +are present and parseable. + +Do not test ephemeral timestamp strings too strictly. + +--- + +# 78. Invalid checkpoint behavior + +Test: + +```text id="qxah3i" +wrong model family +unsupported checkpoint schema +missing model state +incompatible ROOT-GNN config +``` + +Raise clear export errors. + +Do not emit partial ONNX files after validation failure. + +--- + +# 79. Atomic export write + +Where practical, export to a temporary file and only move into the final path after: + +```text id="mk0s8v" +export succeeds +ONNX validation succeeds +parity validation succeeds +``` + +Avoid leaving a seemingly valid final artifact after failure. + +--- + +# 80. Cleanup temporary files + +Ensure failed export attempts clean up temporary ONNX/metadata artifacts where practical. + +Do not leave test litter in output directories. + +--- + +# 81. Documentation + +Add: + +```text id="8xjvnm" +docs/export.md +``` + +or an equivalent clear section. + +Document: + +```text id="xj73lt" +supported models +export command +tensor input contract +dynamic dimensions +raw-logit output +metadata sidecar +ONNX Runtime validation +known limitations +``` + +--- + +# 82. README export section + +Add a concise example: + +```bash id="fcxzsb" +uv run gnn4colliders export \ + export.checkpoint=/path/to/model.pt \ + export.output=model.onnx +``` + +Use the actual final syntax. + +Link to detailed export documentation. + +--- + +# 83. Architecture documentation + +Update: + +```text id="qahvah" +docs/architecture.md +``` + +with the export boundary: + +```text id="uydv0k" +GraphBatch + -> +ROOT-GNN ExportAdapter + -> +tensor-only representation + -> +ONNX model +``` + +Make clear this is an inference/deployment adapter, not the primary graph representation. + +--- + +# 84. Migration documentation + +Update: + +```text id="2wglae" +docs/migration.md +``` + +to mark ONNX/export status accurately. + +If legacy export behavior cannot be reproduced exactly, describe the new supported contract. + +Do not mark legacy deployment parity complete unless validated. + +--- + +# 85. Legacy export parity + +If the legacy repository has an active ONNX export path, characterize: + +```text id="sb3fad" +input representation +output semantics +fixed/dynamic shapes +supported model type +``` + +Compare only externally relevant behavior. + +Do not preserve awkward legacy export implementation solely for implementation parity. + +--- + +# 86. No TensorRT in this task + +Do not implement: + +```text id="xd00dm" +TensorRT +Torch-TensorRT +OpenVINO +``` + +This task is ONNX only. + +A portable ONNX artifact is the foundation for later deployment backends. + +--- + +# 87. No preprocessing deployment framework + +Do not build C++/CUDA ROOT preprocessing for ONNX deployment. + +The exported model contract begins with processed graph tensors. + +Keep scope bounded. + +--- + +# 88. No quantization + +Do not add: + +```text id="12o86l" +INT8 quantization +dynamic quantization +QAT +``` + +in this task. + +Quantization changes numerical behavior and deserves its own dedicated task if needed. + +--- + +# 89. No model simplification dependency + +Do not add ONNX graph-simplifier tooling as a required dependency. + +Only use it if a proven exporter issue requires it. + +Prefer standard ONNX output first. + +--- + +# 90. No distributed export + +Only rank 0 should perform export if the CLI is accidentally launched under a distributed environment. + +Do not export one ONNX file per DDP rank. + +Export operates on the normalized underlying model. + +--- + +# 91. DDP checkpoint export + +Explicitly verify a checkpoint produced during DDP training can be exported. + +Task 14 should already normalize state dicts. + +Add an integration test or reuse an existing normalized DDP checkpoint fixture. + +--- + +# 92. Export config validation + +Validate: + +```text id="onq10b" +checkpoint exists +format is supported +output extension matches ONNX +opset is supported +model family supports export +required example/input source is available +``` + +Fail before expensive work where possible. + +--- + +# 93. Public API + +Expose a small intended API. + +For example: + +```python id="2e4a4t" +from gnn4colliders.export import ( + export_root_gnn_onnx, + RootGNNExportAdapter, +) +``` + +Only expose the adapter if users genuinely need it. + +Keep internal tensor conversion helpers private where practical. + +--- + +# 94. Optional extra installation docs + +If ONNX dependencies are optional, document: + +```bash id="5iiijm" +uv sync --extra root-gnn --extra onnx +``` + +or the actual selected extra name. + +Do not require ONNX dependencies for users who only train ROOT-GNN. + +--- + +# 95. Import hygiene + +Installing only the normal ROOT-GNN extra should not make: + +```python id="z1mofz" +import gnn4colliders +``` + +fail because ONNX Runtime is absent. + +Keep optional dependency imports inside export functionality. + +Raise a clear actionable error if export is requested without the required extra. + +--- + +# 96. Optional dependency error + +Prefer an error conceptually like: + +```text id="20w0as" +ONNX export requires the 'onnx' optional dependency. +Install with: uv sync --extra root-gnn --extra onnx +``` + +Use the actual package-extra name. + +Do not emit a cryptic `ModuleNotFoundError` where a clear project error is easy. + +--- + +# 97. Validate no native behavior regression + +After implementing export, run the full existing test suite. + +ONNX support must not change: + +```text id="qr6203" +training +checkpoint loading +fine-tuning +inference +DDP +``` + +The native path remains authoritative. + +--- + +# 98. Validation + +Run focused export tests: + +```bash id="wd1h8a" +uv run pytest tests/unit/export -v +``` + +Run export integration tests: + +```bash id="93g5qh" +uv run pytest tests/integration -k "export or onnx" -v +``` + +Run model parity tests: + +```bash id="p5c8fy" +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash id="lix9ck" +uv run pytest +``` + +Run lint/format checks: + +```bash id="1o5l3y" +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Validate CLI: + +```bash id="0qr5zt" +uv run gnn4colliders export --help +``` + +Perform a manual tiny export: + +```bash id="g1ur2s" +uv run gnn4colliders export \ + export.checkpoint=/path/to/test/checkpoint.pt \ + export.output=/tmp/root_gnn.onnx +``` + +Then validate the generated ONNX model with the implemented runtime/parity utility. + +Inspect: + +```bash id="fskv6x" +git status +git diff +``` + +Verify: + +* native ROOT-GNN model mathematics remain unchanged +* no task/loss semantics moved inside ONNX +* no metadata/tracking inputs were added to the model +* no TensorRT/quantization/custom-op scope was introduced +* ONNX dependencies remain optional +* no large generated ONNX files are tracked +* no unrelated changes are included + +--- + +# Completion criteria + +Task 17 is complete when: + +1. An isolated export package exists. +2. The native ROOT-GNN model remains unchanged in its public inference behavior. +3. A tensor-only export adapter exists if DGL cannot be exported directly. +4. Adapter outputs match native `EdgeNetwork` outputs. +5. Adapter outputs match native fine-tuned model outputs. +6. Multiclass ROOT-GNN can be exported. +7. Binary fine-tuned ROOT-GNN can be exported. +8. Raw logits are the exported model output. +9. Graph batching semantics are preserved. +10. Node/edge/global message-passing semantics are preserved. +11. Variable graph size is supported to the documented extent. +12. Variable batch size is supported to the documented extent. +13. ONNX opset is explicit. +14. ONNX structural validation succeeds. +15. ONNX Runtime CPU inference works. +16. PyTorch vs ONNX multiclass parity passes. +17. PyTorch vs ONNX binary fine-tuning parity passes. +18. Checkpoint-based export works. +19. DDP-produced normalized checkpoints can be exported. +20. Export metadata/schema information is written. +21. `gnn4colliders export` works through the stable Python export API. +22. Invalid/incompatible exports fail clearly. +23. ONNX dependencies are optional. +24. Normal package imports work without ONNX dependencies. +25. Export documentation exists. +26. README contains a validated export example. +27. Full existing tests still pass. +28. Legacy source remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. direct-DGL export findings +4. export adapter design +5. tensor input contract +6. graph-membership representation +7. dynamic dimension support +8. supported ROOT-GNN model types +9. ONNX opset +10. export API +11. checkpoint export API +12. CLI export workflow +13. ONNX metadata/sidecar format +14. native-vs-adapter parity results +15. adapter-vs-ONNX parity results +16. multiclass export results +17. binary fine-tuning export results +18. single-node/zero-edge behavior +19. variable node/edge test results +20. variable batch-size test results +21. DDP checkpoint export results +22. ONNX Runtime validation results +23. numerical tolerances used +24. optional dependency changes +25. known export limitations +26. behavior intentionally deferred +27. validation commands and results + +After validation succeeds, create one Git commit containing only Task 17 changes. + +Use: + +```text id="6p7aqi" +feat: add ROOT-GNN ONNX export +``` + +Before committing, inspect the final diff and ensure no unrelated files or generated ONNX artifacts are included. diff --git a/tasks/task18.md b/tasks/task18.md new file mode 100644 index 0000000000000000000000000000000000000000..d60c728d7315eadd535a44e33e0f0312aacecca5 --- /dev/null +++ b/tasks/task18.md @@ -0,0 +1,1934 @@ +# Task 18: Legacy Compatibility Audit and Cleanup + +Perform a focused cleanup of the GNN4Colliders rewrite now that the ROOT-GNN v1 stack, documentation, performance work, and ONNX export path are in place. + +This task should determine exactly which compatibility shims and legacy references are still necessary, remove obsolete migration scaffolding, consolidate remaining compatibility code, and leave the new implementation cleanly separated from the frozen legacy tree. + +This task builds on: + +```text +Task 8: ROOT-GNN model and transfer/fine-tuning +Task 9: task/loss/metric semantics +Task 11: checkpoint compatibility +Task 12: inference/output compatibility +Task 13: Hydra configuration and CLI +Task 14: DDP/Perlmutter execution +Task 16: documentation and migration closure +Task 17: ONNX export +``` + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +docs/export.md +``` + +and any compatibility-specific documentation introduced in earlier tasks. + +Then inspect: + +```text +src/gnn4colliders/ +tests/ +configs/ +scripts/ +legacy/ +``` + +Pay particular attention to: + +```text +src/gnn4colliders/models/root_gnn/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +src/gnn4colliders/data/ +tests/parity/ +``` + +Do not delete the legacy source tree in this task unless the repository explicitly no longer needs it for parity/reference purposes. + +--- + +# Goal + +Finish the migration boundary cleanly. + +The desired architecture is: + +```text +new production code + ↓ +clean public APIs + ↓ +small explicit compatibility adapters + ↓ +legacy artifacts / checkpoints / output formats +``` + +not: + +```text +new production code + ↓ +hidden legacy assumptions everywhere +``` + +By the end of this task, it should be clear: + +```text +Which legacy behaviors remain supported? +Which compatibility adapters are required? +Which migration-only shims can be removed? +Which legacy concepts no longer exist in public APIs? +Which historical formats are intentionally unsupported? +``` + +--- + +# 1. Audit all references to `legacy/` + +Search the entire new codebase for: + +```text +legacy +root_gnn_dgl +tracking +tracking_info +module. +_orig_mod. +sys.path +dynamic import +``` + +Classify every hit. + +Use categories such as: + +```text +production compatibility +test-only parity +documentation/history +obsolete migration scaffolding +unexpected coupling +``` + +Do not remove code simply because it contains the word `legacy`. + +--- + +# 2. No production imports from legacy + +The new production package should not import executable logic from: + +```text +legacy/ +``` + +Verify: + +```text +src/gnn4colliders/ +``` + +has no dependency on legacy modules at runtime. + +If any exist, replace them with new implementations or explicit compatibility parsing. + +Parity tests may still import or execute legacy code. + +--- + +# 3. Freeze the compatibility boundary + +Identify all remaining compatibility responsibilities. + +Likely examples include: + +```text +legacy checkpoint loading +legacy state_dict key normalization +legacy classifier/backbone mapping +legacy NPZ output compatibility +legacy metadata/tracking conversion +``` + +Keep these responsibilities isolated. + +Prefer modules named clearly, such as: + +```text +compat/ +legacy_checkpoint.py +legacy_outputs.py +legacy_metadata.py +``` + +only if multiple compatibility modules genuinely justify a package. + +Do not create a broad abstraction layer without need. + +--- + +# 4. Consider a dedicated compatibility package + +If compatibility code is currently scattered across: + +```text +models/ +training/ +inference/ +data/ +``` + +consider consolidating it under: + +```text +src/gnn4colliders/compat/ +``` + +For example: + +```text +src/gnn4colliders/compat/ + __init__.py + checkpoint.py + metadata.py + outputs.py +``` + +Use this only if it improves boundaries. + +Do not move clean core APIs into `compat/`. + +--- + +# 5. Public API must remain modern + +Public APIs should use: + +```text +EventMetadata +GraphBatch +sample_id +weight +fold +EdgeNetwork +FineTunedEdgeNetwork +``` + +Do not expose public APIs requiring: + +```text +tracking[:, 0] +tracking[:, 1] +tracking_info +legacy module/class paths +``` + +Legacy structures should appear only at explicit compatibility entry points. + +--- + +# 6. Remove positional tracking from new production code + +Search for all uses of: + +```text +tracking +tracking_info +``` + +in production code. + +For each: + +* replace with named metadata if accidentally retained +* keep only if it is an explicit compatibility adapter +* document why it remains + +The normal new pipeline must not construct or propagate a positional tracking tensor. + +--- + +# 7. Legacy metadata adapter + +If support for old datasets/outputs still requires positional tracking, provide one explicit conversion boundary. + +Conceptually: + +```python +metadata = EventMetadata.from_legacy_tracking( + tracking_row, + ... +) +``` + +or equivalent. + +This helper should map known positions explicitly: + +```text +tracking[0] -> fold +tracking[1] -> weight +``` + +Do not propagate the original array beyond the adapter. + +--- + +# 8. Reject ambiguous tracking layouts + +Do not attempt to infer arbitrary historical tracking layouts. + +If a compatibility adapter only understands: + +```text +column 0 = fold +column 1 = weight +``` + +validate the input shape and fail clearly otherwise. + +Do not silently guess. + +--- + +# 9. State-dict compatibility audit + +Review all compatibility handling for: + +```text +module. +_orig_mod. +legacy module names +renamed model components +classifier mappings +``` + +Ensure there is one canonical normalization/mapping implementation. + +Do not keep duplicate key-remapping code in: + +```text +models/ +training/checkpoint.py +export/ +``` + +Reuse one compatibility layer. + +--- + +# 10. State-dict normalization should be explicit + +Prefer functions such as: + +```python +normalize_legacy_state_dict_keys(...) +map_legacy_edge_network_state_dict(...) +``` + +with deterministic behavior. + +Do not repeatedly use broad: + +```python +strict=False +``` + +without validating missing/unexpected keys. + +--- + +# 11. Remove obsolete Task 8 compatibility helpers + +If Task 11 or Task 18 has superseded temporary compatibility helpers introduced during model migration, remove the duplicate helper. + +Update imports/tests accordingly. + +There should be one supported mapping path. + +--- + +# 12. Checkpoint compatibility matrix + +Create or update a precise compatibility matrix. + +For example: + +| Artifact | Supported | Notes | +| ---------------------------------------- | --------: | ---------------------- | +| new-format checkpoint | yes | full resume | +| legacy model weights | yes | mapped to new ROOT-GNN | +| legacy pretrained backbone | yes | fine-tuning | +| legacy optimizer state | maybe | document exact support | +| legacy scheduler state | maybe | document exact support | +| arbitrary historical checkpoint variants | no | unsupported | + +Use actual implementation status. + +--- + +# 13. New checkpoints remain canonical + +Do not alter the Task 11 new checkpoint format merely to resemble historical files. + +Legacy files adapt into the new representation. + +The new checkpoint schema remains the canonical supported format. + +--- + +# 14. Legacy checkpoint conversion + +If useful, provide an explicit conversion utility: + +```python +convert_legacy_checkpoint(...) +``` + +that produces a new-format checkpoint. + +Only implement this if it simplifies long-term support. + +A conversion utility should: + +```text +load legacy artifact +normalize model keys +extract supported state +create new checkpoint schema +write explicit compatibility metadata +``` + +Do not claim full conversion of unsupported optimizer/trainer state. + +--- + +# 15. Conversion should be explicit + +Do not silently rewrite a legacy checkpoint on load. + +Conversion is a separate user action/API. + +Ordinary compatibility loading may continue to work without mutating the source file. + +--- + +# 16. Legacy output compatibility audit + +Review Task 12 support for historical output fields such as: + +```text +scores +labels +tracking_info +``` + +Determine whether a real downstream consumer still requires legacy output. + +If yes: + +* keep one explicitly named legacy writer +* document it as compatibility-only + +If no: + +* remove the writer +* remove stale config options +* remove tests that only protect unused historical output + +Do not preserve compatibility indefinitely without a use case. + +--- + +# 17. Legacy writer naming + +If retained, make compatibility obvious. + +Prefer: + +```python +write_legacy_npz(...) +``` + +rather than making legacy behavior the generic: + +```python +write_npz(...) +``` + +The default writer should use the modern named schema. + +--- + +# 18. Modern NPZ remains canonical + +The primary output format should retain named fields such as: + +```text +sample_id +logits +scores +predictions +labels +weight +fold +``` + +as implemented. + +Do not revert to positional matrices for legacy convenience. + +--- + +# 19. ROOT output compatibility + +Audit whether ROOT score writing contains legacy-specific branch naming or behavior. + +Separate: + +```text +modern output contract +``` + +from: + +```text +legacy-compatible branch layout +``` + +if both are required. + +Do not embed hidden legacy names in generic inference logic. + +--- + +# 20. Config compatibility audit + +Search for any compatibility with old YAML patterns such as: + +```yaml +module: +class: +args: +``` + +The new Hydra configuration should not support arbitrary legacy dynamic import syntax. + +If a config conversion helper exists, isolate it as an offline migration tool. + +Do not put legacy YAML interpretation in the normal CLI path. + +--- + +# 21. Legacy config conversion + +Only if useful, add a small conversion/reference tool that helps users manually map active legacy configs to new semantic configs. + +This tool may extract recognized values. + +It should not promise arbitrary automatic translation. + +Prefer documentation over complex parsing if only a few active configs exist. + +--- + +# 22. Remove stale config aliases + +Review Hydra config aliases introduced during migration. + +Remove aliases that: + +* are undocumented +* were temporary +* duplicate canonical names +* are no longer used + +Keep aliases only when compatibility value outweighs confusion. + +Document retained aliases. + +--- + +# 23. CLI compatibility audit + +Review the CLI for old naming carried forward merely for migration convenience. + +Prefer current commands: + +```text +prepare +train +evaluate +predict +export +``` + +Do not add old script-name aliases unless actively needed. + +Remove stale deprecated flags/config hooks if they have no supported use. + +--- + +# 24. Remove `sys.path` migration hacks + +Search the entire new package for: + +```python +sys.path.insert(...) +sys.path.append(...) +``` + +Production package code should not modify import paths. + +Remove any migration-era hacks. + +Tests should use normal installed-package imports. + +--- + +# 25. Import architecture audit + +Verify: + +```text +gnn4colliders.data +gnn4colliders.features +gnn4colliders.graphs +gnn4colliders.models +gnn4colliders.tasks +gnn4colliders.training +gnn4colliders.inference +gnn4colliders.distributed +gnn4colliders.export +``` + +do not depend on arbitrary working-directory layout. + +Normal package installation must be sufficient. + +--- + +# 26. Remove obsolete global mutable state + +Audit new production code for migration-era globals introduced to emulate legacy behavior. + +Examples: + +```text +global config objects +global dataset state +global device state +global current fold +global checkpoint path +``` + +Remove low-risk accidental globals. + +Do not perform a major refactor if state is intentional and properly encapsulated. + +--- + +# 27. RNG compatibility cleanup + +Ensure no compatibility code has reintroduced: + +```python +torch.manual_seed(...) +np.random.seed(...) +``` + +inside constructors or hidden helpers. + +Legacy initialization differences should be handled by: + +```text +fixed-weight parity +explicit caller seeding +state_dict compatibility +``` + +not hidden global seed mutation. + +--- + +# 28. Model architecture cleanup + +Review `models/root_gnn/` for migration-only structure. + +Remove: + +```text +unused aliases +temporary wrappers +duplicate legacy class names +dead code paths +``` + +only if no public/test/compatibility user remains. + +Keep clean modern names canonical. + +--- + +# 29. Fine-tuning compatibility cleanup + +Ensure transfer learning now uses the explicit Task 8 API. + +Remove any remaining code that: + +```text +mutates arbitrary model.children() +removes "last layer" by numeric index +depends on legacy module ordering +``` + +unless contained entirely in a compatibility state-dict converter. + +--- + +# 30. Training compatibility cleanup + +Review Trainer/Task integration for legacy terminology such as: + +```text +finish function +tracking +test-as-validation +``` + +Modern production code should use: + +```text +task +metadata +validation +test +``` + +Keep legacy terms only in compatibility docs/tests. + +--- + +# 31. Loss compatibility audit + +Task 9 intentionally preserved active legacy mathematical semantics. + +Do not "clean up" the mathematics in Task 18. + +Instead ensure compatibility comments/documentation distinguish: + +```text +legacy-compatible scientific behavior +``` + +from: + +```text +legacy implementation structure +``` + +The former remains; the latter should be removed where unnecessary. + +--- + +# 32. Padding compatibility audit + +Review any retained legacy padding modes. + +Classify each as: + +```text +active required behavior +supported compatibility +dead historical behavior +``` + +Remove code for unused modes only if confident no standard/reference config uses them. + +Do not remove behavior solely because it looks odd. + +--- + +# 33. Cache compatibility audit + +Review cache schema handling. + +Ensure new code does not silently load old incompatible cache artifacts without validation. + +If legacy caches are unsupported, state this explicitly. + +If conversion is supported, isolate it. + +Prefer rebuilding data caches over maintaining complex cache compatibility unless a real operational need exists. + +--- + +# 34. Feature schema compatibility + +Ensure: + +```text +FEATURE_SCHEMA_VERSION +GRAPH_SCHEMA_VERSION +CACHE_SCHEMA_VERSION +``` + +or equivalent are current and documented. + +Compatibility adapters should not bypass schema checks casually. + +--- + +# 35. Export compatibility audit + +Review Task 17 for duplicated compatibility logic. + +ONNX export should consume the normalized new model/checkpoint interface. + +It should not independently know historical checkpoint key prefixes. + +Remove such duplication if present. + +--- + +# 36. DDP compatibility audit + +New distributed checkpoints should already avoid `module.` prefixes. + +Confirm DDP production code does not rely on legacy prefix handling. + +Legacy prefix normalization belongs only to compatibility loaders. + +--- + +# 37. Tests: classify parity vs compatibility + +Review: + +```text +tests/parity/ +``` + +and other compatibility tests. + +Separate mentally and, where helpful, structurally: + +```text +scientific parity tests +artifact compatibility tests +``` + +Do not remove scientific parity tests merely because migration is complete. + +They remain valuable regression tests. + +--- + +# 38. Keep core parity tests + +Continue protecting: + +```text +feature construction +graph topology +edge features +ROOT-GNN forward +fine-tuning forward +loss semantics +metrics +``` + +against established legacy behavior. + +These are not temporary migration tests if they encode important scientific contracts. + +--- + +# 39. Remove obsolete migration tests + +Delete tests only when they protect implementation scaffolding that no longer exists. + +Examples: + +```text +temporary adapter created before final API +duplicate state-dict converter +deprecated config alias with no support commitment +``` + +Do not reduce useful behavioral coverage. + +--- + +# 40. Legacy executable test dependency + +Determine whether parity tests currently require running the entire legacy stack/environment. + +If so, consider whether a subset can be converted into frozen golden fixtures. + +For example: + +```text +legacy input +expected features +expected edge indices +expected edge features +fixed state_dict +expected logits +expected loss +``` + +This can reduce dependence on an increasingly old Python/DGL environment. + +Do not eliminate direct legacy execution tests if they still provide unique value. + +--- + +# 41. Golden fixture strategy + +Where appropriate, store small deterministic expected values rather than large artifacts. + +Prefer: + +```text +tiny tensors +small JSON/NPZ fixtures +``` + +Do not commit large ROOT/model binaries solely for compatibility tests. + +--- + +# 42. Legacy environment isolation + +If some parity tests require the old environment, mark/document them clearly. + +They should not make ordinary modern development impossible. + +Use the project's existing pytest marker conventions if present. + +Do not silently skip critical parity coverage in the canonical validated environment. + +--- + +# 43. `legacy/` README + +Add or update: + +```text +legacy/README.md +``` + +explaining: + +```text +why legacy exists +which tree is the actual behavioral reference +which abandoned rewrites are not targets +whether legacy code is frozen +whether production code may import it +when it may eventually be removed +``` + +Keep this concise. + +--- + +# 44. Mark abandoned rewrites clearly + +If directories such as: + +```text +legacy/physicsnemo/ +``` + +remain, mark them clearly as: + +```text +historical/abandoned rewrite attempt +not parity target +not production implementation +``` + +Avoid future developer confusion. + +--- + +# 45. Documentation cleanup + +Update: + +```text +README.md +docs/architecture.md +docs/migration.md +docs/export.md +``` + +to reflect the final compatibility boundary. + +Remove language suggesting the new rewrite is still transitional if it is now canonical. + +--- + +# 46. Migration document status + +`docs/migration.md` should clearly state: + +```text +ROOT-GNN rewrite is canonical +legacy remains reference/compatibility source +``` + +and list the remaining supported legacy artifacts. + +Do not leave "TODO migrate X" entries for completed functionality. + +--- + +# 47. Compatibility support policy + +Add a concise support policy. + +For example: + +```text +We preserve: +- active scientific behavior through parity tests +- explicitly documented historical checkpoint formats +- explicitly documented legacy output formats where needed + +We do not preserve: +- arbitrary old dynamic-import configs +- undocumented historical cache formats +- every experimental model class +``` + +Use actual project decisions. + +--- + +# 48. Deprecation language + +If any compatibility API is intended for later removal, mark it clearly. + +Use: + +```text +deprecated compatibility API +``` + +only if there is an actual replacement and removal intent. + +Do not add noisy runtime warnings everywhere unless helpful. + +Documentation may be sufficient. + +--- + +# 49. Avoid breaking supported users casually + +Do not remove a compatibility behavior that is already documented as supported without: + +```text +updating docs +updating tests +explaining the intentional break +``` + +This task is cleanup, not arbitrary breaking-change release. + +--- + +# 50. Search for stale legacy names + +Search for historical class/module names such as: + +```text +Edge_Network +Transferred_Learning_Finetuning +GCN.py +tracking +finish +``` + +in new production code. + +Modern code should generally use: + +```text +EdgeNetwork +FineTunedEdgeNetwork +task postprocessing +metadata +``` + +Compatibility/documentation references are fine. + +--- + +# 51. Search for duplicate behavior + +Look for multiple implementations of: + +```text +phi wrapping +state_dict normalization +checkpoint reconstruction +score conversion +event weight extraction +sample ID creation +``` + +Migration often creates duplicates. + +Consolidate only where clearly redundant. + +Do not create a giant generic utility module. + +--- + +# 52. No `utils.py` dumping ground + +Do not solve cleanup by moving unrelated compatibility functions into: + +```text +utils.py +``` + +Prefer domain-specific modules. + +--- + +# 53. Public export cleanup + +Review all package `__init__.py` exports. + +Remove public exports for: + +```text +migration-only helpers +internal adapters +obsolete aliases +``` + +Keep the supported public surface small. + +Compatibility helpers may be imported from: + +```text +gnn4colliders.compat +``` + +if intentionally user-facing. + +--- + +# 54. Public API import test + +Add or update a test that verifies the intended public API imports. + +For example: + +```python +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork +from gnn4colliders.tasks import BinaryClassificationTask +from gnn4colliders.training import Trainer, CheckpointManager +from gnn4colliders.inference import Predictor +``` + +Use actual public APIs. + +Do not include deprecated internal helpers. + +--- + +# 55. Optional compatibility imports + +If compatibility is intentionally exposed, make it obvious: + +```python +from gnn4colliders.compat import load_legacy_checkpoint +``` + +rather than exposing it as a normal checkpoint API with ambiguous semantics. + +--- + +# 56. Dependency cleanup + +Review `pyproject.toml`. + +Remove dependencies introduced during migration that are no longer used. + +Do not remove dependencies based only on static import search if they are required by optional workflows/config. + +Verify extras: + +```text +root-gnn +onnx/export +dev +``` + +or actual current names. + +--- + +# 57. Optional-extra boundaries + +Ensure: + +```text +core install +ROOT-GNN install +ONNX install +development install +``` + +remain logically separated. + +Legacy compatibility should not force users to install the legacy environment. + +--- + +# 58. Legacy dependencies must not leak + +The new environment should not require: + +```text +Python 3.8 +old DGL +old CUDA +legacy ROOT Python bindings +``` + +merely to import/use the new production package. + +Only dedicated parity/reference workflows may require historical dependencies. + +--- + +# 59. Remove unused files + +Identify obviously obsolete new-tree files created during earlier tasks. + +Examples: + +```text +empty compatibility modules +placeholder scripts +superseded config files +unused migration helpers +``` + +Delete only after confirming no import/config/test references. + +--- + +# 60. Generated artifacts + +Remove accidental repository artifacts such as: + +```text +__pycache__/ +*.egg-info/ +profiling traces +temporary checkpoints +temporary NPZ/ONNX outputs +``` + +if present. + +Update `.gitignore` where appropriate. + +Do not delete legitimate fixtures. + +--- + +# 61. Dead-code detection + +Use search and test coverage reasoning to identify obvious dead code. + +Do not add a new static-analysis dependency solely for this task. + +Prefer targeted manual audit. + +--- + +# 62. TODO/FIXME audit + +Search: + +```text +TODO +FIXME +HACK +XXX +legacy +temporary +migration +``` + +in production code. + +Resolve or reclassify each relevant migration-era note. + +Do not leave stale comments claiming something is temporary after it became the final implementation. + +--- + +# 63. Comments should explain compatibility + +Keep compatibility comments focused on why behavior exists. + +For example: + +```text +Legacy checkpoints may contain the `module.` prefix from DDP. +Normalize it at this boundary so the canonical model state remains prefix-free. +``` + +Avoid comments that narrate obvious code. + +--- + +# 64. Error messages + +Compatibility failures should explain what is unsupported. + +For example: + +```text +Unsupported legacy checkpoint layout: expected model_state_dict and epoch fields. +``` + +Prefer actionable errors over raw KeyError/shape errors. + +Do not catch every exception indiscriminately. + +--- + +# 65. Compatibility warnings + +When loading supported-but-legacy artifacts, consider a concise informational warning/log. + +For example: + +```text +Loading legacy ROOT-GNN checkpoint through compatibility adapter. +``` + +Do not spam warnings per tensor/key. + +Make warning behavior testable if added. + +--- + +# 66. Unsupported legacy artifacts + +Explicitly reject known unsupported formats rather than partially loading them. + +Examples may include: + +```text +experimental model families +unknown tracking layout +obsolete cache schema +unrecognized checkpoint architecture +``` + +Fail early. + +--- + +# 67. Experimental legacy models + +Confirm the new production code does not accidentally expose unsupported historical models such as: + +```text +GCN_global +GCN_global_2way +attention variants +MultiModel +Clustering +``` + +unless they were intentionally migrated in a later task. + +Document them as unsupported legacy experiments. + +--- + +# 68. Legacy transfer variants + +Likewise, retain only the active transfer/fine-tuning behavior migrated in Task 8. + +Do not keep compatibility wrappers for unused historical transfer classes without evidence they are needed. + +--- + +# 69. Legacy config examples + +Remove obsolete old YAML examples from the main documentation. + +If valuable historically, keep them under: + +```text +legacy/ +``` + +or migration docs. + +Normal users should see semantic Hydra configs. + +--- + +# 70. Re-run active reference configs + +After cleanup, ensure the canonical reference workflows still compose/run: + +```text +multiclass pretraining +binary training +pretrained fine-tuning +resume +evaluate +predict +export +``` + +Cleanup must not break the v1 stack. + +--- + +# 71. Checkpoint round-trip regression + +Run Task 11 checkpoint tests. + +Ensure compatibility cleanup does not affect: + +```text +new-format save/load +resume +best/latest selection +weight-only loading +``` + +--- + +# 72. Legacy checkpoint regression + +Run explicit legacy checkpoint compatibility tests. + +Verify all formats still documented as supported remain supported. + +Do not weaken expectations merely to simplify adapters. + +--- + +# 73. Fine-tuning regression + +Verify: + +```text +legacy pretrained checkpoint + -> +compatibility loader + -> +new backbone + -> +FineTunedEdgeNetwork +``` + +still works for the supported legacy artifact. + +This is one of the most valuable compatibility workflows. + +--- + +# 74. Inference compatibility regression + +If modern inference supports loading historical model weights, test: + +```text +legacy checkpoint + -> +new Predictor + -> +modern PredictionResult +``` + +The output should use the modern schema unless the user explicitly requests a legacy writer. + +--- + +# 75. ONNX compatibility regression + +If a supported legacy checkpoint can be normalized into a new model, verify it can still reach the Task 17 export path. + +Do not make ONNX exporter directly consume legacy layouts. + +The flow should be: + +```text +legacy checkpoint + -> +compatibility normalization + -> +new model + -> +ONNX export +``` + +--- + +# 76. No legacy logic inside ONNX adapter + +Verify Task 17's export adapter has no: + +```text +legacy key normalization +tracking conversion +old class-name handling +``` + +Those belong upstream. + +--- + +# 77. DDP regression + +Run Task 14 tests. + +Cleanup of state-dict prefix handling must not break: + +```text +new DDP checkpoint save +single-process load +distributed resume +``` + +New DDP behavior and legacy DDP compatibility are separate concepts. + +--- + +# 78. Performance regression + +Do not intentionally degrade Task 15 optimized hot paths just to centralize compatibility code. + +Compatibility adapters should generally operate at: + +```text +load +conversion +output +``` + +boundaries, not inside every training batch. + +--- + +# 79. Compatibility code should not be in hot loops + +Ensure no legacy conversion runs: + +```text +per node +per edge +per batch +``` + +during normal modern training. + +Convert once at ingestion/load boundaries. + +--- + +# 80. Architecture-neutral audit + +Use cleanup to verify shared infrastructure remains future-model friendly. + +Search for ROOT-GNN compatibility assumptions inside: + +```text +data/ +features/ +tasks/ +training/ +inference/ +distributed/ +``` + +Move graph/model-specific compatibility toward: + +```text +compat/ +graphs/ +models/root_gnn/ +``` + +where appropriate. + +Avoid major refactors. + +--- + +# 81. Future ROOT-Transformer must not inherit legacy tracking + +Ensure no shared API requires: + +```text +tracking +GraphSample +DGLGraph +``` + +for all models. + +Shared `EventSample`/metadata should remain architecture-neutral. + +Document any remaining coupling as technical debt. + +--- + +# 82. Test fixture cleanup + +Review test fixtures for duplicate or stale legacy data representations. + +Prefer reusable tiny fixtures for: + +```text +EventSample +GraphSample +GraphBatch +state_dict +``` + +Keep explicit legacy fixtures only where testing compatibility. + +--- + +# 83. Avoid huge fixture binaries + +Do not commit real production checkpoints/datasets. + +Use tiny deterministic synthetic fixtures. + +If a legacy checkpoint fixture is needed, construct a minimal compatible dictionary in tests. + +--- + +# 84. Compatibility tests naming + +Use clear test names such as: + +```text +test_legacy_checkpoint_module_prefix +test_legacy_checkpoint_orig_mod_prefix +test_legacy_tracking_to_metadata +test_legacy_npz_writer +``` + +Do not hide compatibility behavior inside generic tests. + +--- + +# 85. Deprecation tests + +If an API is formally deprecated, add a minimal test for the warning/behavior. + +Do not test exact warning prose too rigidly. + +--- + +# 86. Documentation: compatibility guide + +If compatibility behavior is substantial, add: + +```text +docs/compatibility.md +``` + +Otherwise keep it in `docs/migration.md`. + +Document: + +```text +supported historical checkpoints +supported output compatibility +unsupported artifacts +conversion examples +``` + +Keep it practical. + +--- + +# 87. Example legacy checkpoint conversion + +If a conversion API is implemented, document a small example. + +Conceptually: + +```bash +uv run python -m ... convert-legacy-checkpoint ... +``` + +or a Python API. + +Do not add a CLI command unless conversion is a real user workflow. + +A documented Python utility may be enough. + +--- + +# 88. No broad legacy CLI + +Do not add: + +```text +gnn4colliders legacy ... +``` + +as a generic compatibility subsystem. + +Keep compatibility focused. + +--- + +# 89. Migration closure statement + +Update migration docs to state that production development should now target: + +```text +src/gnn4colliders/ +``` + +and not modify the frozen legacy implementation except when maintaining parity fixtures/documentation. + +--- + +# 90. Legacy deletion criteria + +Document what must be true before `legacy/` can eventually be removed. + +For example: + +```text +all required scientific contracts have stable golden fixtures +historical checkpoint compatibility no longer requires direct source inspection +no active parity test executes legacy code +repository policy approves archival/removal +``` + +Do not delete it in this task unless explicitly instructed. + +--- + +# 91. Validate documentation links + +Ensure README/docs references to: + +```text +compatibility +migration +legacy +checkpoint conversion +``` + +point to real files/sections. + +Avoid broken links after moving docs. + +--- + +# 92. Validation + +Run searches such as: + +```bash +grep -R "tracking\\[" src/gnn4colliders +grep -R "tracking_info" src/gnn4colliders +grep -R "sys.path" src/gnn4colliders +grep -R "legacy/" src/gnn4colliders +``` + +Use a more appropriate search tool if available. + +Review every hit. + +Run focused compatibility tests: + +```bash +uv run pytest tests/parity -v +``` + +Run checkpoint tests: + +```bash +uv run pytest tests/unit/training -v +``` + +Run inference tests: + +```bash +uv run pytest tests/unit/inference -v +``` + +Run export tests: + +```bash +uv run pytest tests/unit/export -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run distributed tests: + +```bash +uv run pytest tests/unit/distributed -v +uv run pytest tests/integration -k distributed -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +Validate public CLI workflows: + +```bash +uv run gnn4colliders --help +uv run gnn4colliders prepare --help +uv run gnn4colliders train --help +uv run gnn4colliders evaluate --help +uv run gnn4colliders predict --help +uv run gnn4colliders export --help +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* production code does not import executable legacy modules +* positional tracking is absent from modern APIs +* compatibility logic is isolated +* new checkpoints remain canonical +* modern output schemas remain canonical +* no duplicate state-dict mapping remains +* no old dynamic-import YAML path exists in normal execution +* scientific parity tests remain intact +* no supported legacy artifact was broken unintentionally +* no major new functionality was added +* no unrelated refactor is included + +--- + +# Completion criteria + +Task 18 is complete when: + +1. Every production reference to legacy concepts has been audited. +2. Production code does not import executable logic from `legacy/`. +3. Positional tracking arrays are absent from modern public APIs. +4. Legacy tracking conversion, if retained, exists only at an explicit compatibility boundary. +5. Legacy checkpoint key normalization has one canonical implementation. +6. Duplicate state-dict mapping helpers have been removed. +7. New checkpoint format remains canonical. +8. Supported legacy checkpoint workflows still pass. +9. Unsupported historical checkpoint behavior is documented clearly. +10. Resume and pretrained transfer remain distinct. +11. Modern NPZ output remains canonical. +12. Legacy NPZ/output behavior is explicitly compatibility-only or removed. +13. Old dynamic-import YAML is absent from normal CLI/config paths. +14. `sys.path` migration hacks are absent from production code. +15. Obsolete migration globals/helpers are removed. +16. ROOT-GNN production model code uses modern APIs/names. +17. Task/training code contains no accidental legacy positional semantics. +18. Cache/schema compatibility behavior is explicit. +19. Task 17 export uses normalized new models rather than legacy logic. +20. DDP new-format state handling remains unaffected. +21. Scientific parity tests remain intact. +22. Obsolete migration-only tests are removed only where justified. +23. `legacy/README.md` clearly describes the frozen legacy tree. +24. Abandoned historical rewrites are clearly marked. +25. Documentation accurately describes supported compatibility. +26. Public exports are cleaned up. +27. Unused dependencies/files introduced during migration are removed where safe. +28. Full modern pretraining/fine-tuning/inference/export workflows still work. +29. Full tests pass. +30. Lint/format checks pass. +31. Legacy source remains frozen. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. files removed +4. legacy-reference audit findings +5. positional tracking audit findings +6. compatibility package/module structure +7. legacy metadata conversion behavior +8. checkpoint compatibility consolidation +9. state-dict normalization behavior +10. legacy checkpoint formats supported +11. legacy checkpoint formats unsupported +12. legacy optimizer/scheduler resume status +13. legacy output compatibility status +14. config compatibility cleanup +15. CLI compatibility cleanup +16. `sys.path`/import cleanup +17. model migration-shim cleanup +18. training/task cleanup +19. cache/schema compatibility findings +20. export compatibility cleanup +21. DDP compatibility regression results +22. parity-test cleanup +23. golden-fixture changes +24. legacy executable-test dependency status +25. public API cleanup +26. dependency cleanup +27. generated/dead-file cleanup +28. TODO/FIXME audit results +29. documentation updates +30. remaining compatibility debt +31. criteria for eventual `legacy/` removal +32. validation commands and results + +After validation succeeds, create one Git commit containing only Task 18 changes. + +Use: + +```text +refactor: consolidate legacy compatibility boundaries +``` + +Before committing, inspect the final diff and ensure no unrelated scientific or architectural changes are included. diff --git a/tasks/task19.md b/tasks/task19.md new file mode 100644 index 0000000000000000000000000000000000000000..46e831edc17abff417cfcfa5ad6b4e77a9277e90 --- /dev/null +++ b/tasks/task19.md @@ -0,0 +1,1797 @@ +# Task 19: Packaging, Release, and CI Hardening + +Harden GNN4Colliders for reproducible installation, automated validation, and a clean first release of the rewritten ROOT-GNN stack. + +This task should make the repository easier to install, test, package, and release without changing scientific behavior. + +This task builds on: + +```text +Task 13: Hydra configuration and CLI +Task 14: distributed / Perlmutter execution +Task 15: profiling and performance +Task 16: documentation and migration closure +Task 17: ONNX export +Task 18: legacy compatibility cleanup +``` + +Before making changes, read: + +```text +AGENTS.md +README.md +pyproject.toml +docs/architecture.md +docs/migration.md +docs/compatibility.md +docs/export.md +``` + +where those files exist. + +Then inspect: + +```text +.github/ +configs/ +scripts/ +src/gnn4colliders/ +tests/ +benchmarks/ +pyproject.toml +uv.lock +``` + +Also inspect any existing: + +```text +CI workflows +release scripts +version files +build configuration +package metadata +``` + +Do not change scientific algorithms, parity behavior, model semantics, task semantics, or data preprocessing in this task. + +--- + +# Goal + +Make the rewritten project release-ready. + +The desired lifecycle is: + +```text +source repository + ↓ +clean dependency metadata + ↓ +automated lint / test / package validation + ↓ +build sdist + wheel + ↓ +install built artifacts in clean environment + ↓ +smoke-test public CLI/imports + ↓ +publishable release artifact +``` + +The task should answer: + +```text +Can a fresh user install the package reproducibly? +Can CI detect regressions automatically? +Can optional dependencies be installed independently? +Can a built wheel/sdist install cleanly? +Does the installed CLI work outside the source checkout? +Is the package metadata ready for release? +``` + +--- + +# 1. Audit `pyproject.toml` + +Review the entire package metadata/configuration. + +Verify: + +```text +project name +version strategy +description +README +license +authors/maintainers if appropriate +Python requirement +dependencies +optional dependencies +CLI entry points +build backend +package discovery +``` + +Remove stale migration-era metadata. + +Do not add fields with guessed values. + +--- + +# 2. Confirm package naming + +Keep naming consistent: + +```text +repository: GNN4Colliders +distribution/project name: choose the existing canonical project name +Python import package: gnn4colliders +CLI: gnn4colliders +``` + +Do not rename the Python package in this task. + +If PyPI/distribution naming differs intentionally from the repository name, document it clearly. + +--- + +# 3. Versioning strategy + +Choose and implement one simple, explicit versioning strategy. + +Prefer either: + +```text +static version in pyproject.toml +``` + +or: + +```text +single-source version from gnn4colliders.__version__ +``` + +Do not introduce dynamic Git-derived versioning unless the repository already uses it or there is a strong reason. + +Keep release mechanics understandable. + +--- + +# 4. Add `__version__` + +If useful and not already present, expose: + +```python +import gnn4colliders + +gnn4colliders.__version__ +``` + +Keep it synchronized with the package metadata. + +Add a test. + +Do not create two manually maintained version strings. + +--- + +# 5. Python support metadata + +Verify `requires-python` matches the validated runtime. + +Use the actual supported version established by the project. + +Do not broaden support to untested Python versions for appearance. + +Do not retain stale legacy Python constraints. + +--- + +# 6. Core dependencies + +Review dependencies required for architecture-neutral core functionality. + +Distinguish carefully between: + +```text +core package +ROOT/Awkward data support +ROOT-GNN support +ONNX/export support +development/testing +``` + +Do not force large ROOT-GNN-specific dependencies into the smallest package installation unless they are genuinely core. + +--- + +# 7. Optional dependency extras + +Review and rationalize extras. + +Prefer a clear layout conceptually such as: + +```text +root-gnn +onnx +dev +``` + +or the actual naming already established. + +Potential responsibilities: + +```text +root-gnn: + PyTorch/DGL-specific ROOT-GNN runtime + +onnx: + onnx + onnxruntime + +dev: + pytest + pytest-cov + ruff + build/twine or equivalent validation tools +``` + +Do not create a large matrix of overlapping extras. + +--- + +# 8. Avoid duplicate dependency declarations + +Do not repeat the same dependency inconsistently across: + +```text +project.dependencies +optional-dependencies +tool-specific dependency groups +``` + +Use the existing uv dependency model consistently. + +--- + +# 9. DGL installation reality + +Preserve the validated DGL installation strategy from earlier tasks. + +Do not pretend DGL is available from a source/index combination that does not actually work. + +If a custom wheel index or explicit environment step remains necessary, document it. + +Do not bake private machine-specific URLs into package metadata unless they are intentionally public and supported. + +--- + +# 10. Lockfile audit + +Review: + +```text +uv.lock +``` + +for consistency with `pyproject.toml`. + +Regenerate only if dependency metadata changes require it. + +Do not perform broad dependency upgrades merely because newer versions exist. + +This task is hardening, not dependency modernization. + +--- + +# 11. Build backend + +Verify the package has a clean, standard build backend. + +Use the existing backend if it works. + +Do not switch build systems without need. + +Ensure package files under: + +```text +src/gnn4colliders/ +``` + +are included correctly. + +--- + +# 12. Source distribution + +Build an sdist. + +Conceptually: + +```bash +uv build +``` + +or the current project-standard build command. + +Verify the sdist includes required source/config/documentation files and excludes generated artifacts. + +--- + +# 13. Wheel build + +Build a wheel. + +Verify the resulting wheel includes: + +```text +Python package +CLI entry point metadata +required package data +version metadata +``` + +and excludes: + +```text +tests if not intended +benchmarks if not intended +legacy tree unless intentionally packaged +temporary outputs +cache files +``` + +Make an explicit decision about what belongs in the installed artifact. + +--- + +# 14. Do not package `legacy/` by accident + +The frozen: + +```text +legacy/ +``` + +tree is repository reference material. + +It should not be installed into the Python wheel unless there is a concrete runtime requirement. + +Prefer excluding it from distribution. + +Parity tests can use it from the source repository. + +--- + +# 15. Package data audit + +Determine whether runtime needs any package data such as: + +```text +default configs +schema files +small templates +``` + +If yes, ensure they are included intentionally. + +Do not rely on source-tree-relative filesystem access that disappears after installation. + +--- + +# 16. Hydra config packaging + +This is especially important. + +Determine whether the Hydra configuration tree is: + +```text +repository-only external config +``` + +or: + +```text +required installed package data +``` + +The installed CLI must be able to find its canonical configs when invoked outside the source repository. + +If configs are required by the CLI, package them appropriately or move/copy them into a supported package-data location. + +Do not rely on: + +```text +Path("configs") +``` + +relative to the current working directory. + +--- + +# 17. Installed CLI test + +The following must work from a clean directory unrelated to the Git checkout: + +```bash +gnn4colliders --help +``` + +and, where possible: + +```bash +gnn4colliders train --help +gnn4colliders evaluate --help +gnn4colliders predict --help +gnn4colliders export --help +``` + +This is a critical packaging test. + +--- + +# 18. Clean-environment wheel installation test + +Create a temporary clean environment. + +Install the built wheel rather than the source tree. + +Test: + +```text +import gnn4colliders +public imports +CLI help +Hydra config discovery +``` + +Do not rely on editable installation for this test. + +--- + +# 19. Sdist installation test + +Where practical, create a clean environment and install the built sdist. + +Verify the package builds/installs correctly from source distribution. + +This catches missing source files that wheel-only testing may miss. + +--- + +# 20. Optional extra installation tests + +Test important installation modes independently. + +At minimum, where feasible: + +```text +base/core +root-gnn +onnx +dev +``` + +The exact modes depend on the final extras. + +The package should fail clearly when an optional workflow is requested without its extra. + +--- + +# 21. Core import must remain lightweight + +If ROOT-GNN dependencies are optional, then: + +```python +import gnn4colliders +``` + +should not import DGL/ONNX unnecessarily. + +Avoid optional-dependency imports at package top level. + +--- + +# 22. ROOT-GNN optional dependency failure + +If a user invokes ROOT-GNN functionality without the required extra, raise a clear message. + +Do not leave them with an obscure import traceback if a clean project-specific error is straightforward. + +--- + +# 23. ONNX optional dependency failure + +Likewise, Task 17 export functionality should fail clearly when ONNX extras are absent. + +Validate this behavior in a clean installation if practical. + +--- + +# 24. Public import contract + +Define a concise supported import surface. + +Examples may include: + +```python +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork +from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask +from gnn4colliders.training import Trainer, CheckpointManager +from gnn4colliders.inference import Predictor +``` + +Use actual APIs. + +Add a smoke test. + +--- + +# 25. Avoid exporting internal APIs + +Audit package `__init__.py` files. + +Do not expose: + +```text +migration helpers +internal state-dict conversion details +private collective helpers +temporary factory functions +``` + +as top-level public APIs unless intentionally supported. + +--- + +# 26. Add CI if absent + +Use GitHub Actions if the repository is hosted on GitHub and no other CI system is established. + +Prefer a small set of understandable workflows rather than many fragmented files. + +For example: + +```text +.github/workflows/ci.yml +.github/workflows/package.yml +``` + +Use fewer files if simpler. + +--- + +# 27. CI philosophy + +CI should validate: + +```text +format/lint +unit tests +integration tests +parity tests that work in the modern environment +package build +installed-package smoke tests +``` + +Do not require Perlmutter, Slurm, or GPUs for normal CI. + +--- + +# 28. CI Python version + +Test the canonical supported Python version. + +If the package genuinely supports more than one Python version, a small matrix is acceptable. + +Do not create a broad matrix of unvalidated Python versions. + +--- + +# 29. CI dependency installation + +Use the project's canonical uv workflow. + +Conceptually: + +```text +install uv +uv sync ... +uv run ... +``` + +Pin/install tooling in a reproducible manner where practical. + +Do not maintain a second pip requirements stack solely for CI. + +--- + +# 30. CI cache + +Use dependency caching if straightforward. + +Do not make caching correctness-sensitive. + +CI should work when caches are cold. + +--- + +# 31. Lint job + +Run: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +or the final project-standard commands. + +Keep lint errors blocking. + +--- + +# 32. Unit tests in CI + +Run the ordinary unit suite. + +Keep tests CPU-safe. + +Do not silently exclude core unit areas without explanation. + +--- + +# 33. Integration tests in CI + +Run integration tests that do not require external production data or GPUs. + +Use tiny generated fixtures. + +Do not make CI depend on NERSC filesystem access. + +--- + +# 34. Parity tests in CI + +Run parity tests that are valid under the modern environment. + +If some parity tests require the old legacy environment and cannot run in canonical CI, separate them clearly. + +Do not silently pretend the whole parity suite runs if it does not. + +--- + +# 35. Legacy-dependent parity marker + +If necessary, introduce or use a marker such as: + +```text +legacy_env +``` + +for tests requiring the historical runtime. + +Normal CI may exclude these, but: + +* document this +* keep modern golden-fixture parity tests running +* provide a separate manual workflow if useful + +Do not let all scientific parity disappear from CI. + +--- + +# 36. Distributed CPU tests in CI + +If Task 14's two-rank CPU distributed tests are reliable and fast enough, include them. + +If they are flaky/slow, keep a focused subset. + +Do not require NCCL/CUDA. + +--- + +# 37. ONNX tests in CI + +If the ONNX extra is supported and installable in CI, run the Task 17 CPU ONNX export/runtime tests in either: + +```text +main CI +``` + +or: + +```text +separate optional-dependency job +``` + +Avoid forcing ONNX into the base install if it is optional. + +--- + +# 38. Package build job + +Add a CI job that builds: + +```text +sdist +wheel +``` + +and validates them. + +Use the actual project build command. + +Do not publish from normal pull-request CI. + +--- + +# 39. Wheel install smoke job + +After building the wheel: + +1. create a clean environment +2. install the wheel +3. run import smoke tests +4. run CLI help +5. verify package version + +This should not import from the source checkout accidentally. + +Be careful about current working directory/PYTHONPATH. + +--- + +# 40. Verify source checkout is not masking packaging failures + +When testing the installed wheel, run from a temporary directory outside the repository. + +This is mandatory. + +Otherwise Python may import: + +```text +src/gnn4colliders/ +``` + +and hide missing-wheel files. + +--- + +# 41. Package contents test + +Inspect wheel/sdist file lists. + +Verify intentional inclusion/exclusion. + +It is acceptable to add a small script/test that checks for critical files such as packaged configs. + +Do not overfit to every filename. + +--- + +# 42. No generated artifacts in distributions + +Ensure distributions exclude: + +```text +__pycache__ +*.pyc +*.egg-info +outputs +checkpoints +predictions +profiles +benchmark-results +ONNX artifacts +temporary ROOT fixtures +``` + +unless a specific tiny fixture is intentionally package data. + +--- + +# 43. Test fixture packaging + +Ordinary test fixtures should not automatically enter production wheels. + +Keep them under tests unless runtime examples explicitly require them. + +--- + +# 44. Release artifact validation + +Use standard tooling to validate built package metadata. + +For example, if appropriate: + +```text +twine check +``` + +or an equivalent standard package validation. + +Do not introduce an elaborate release framework. + +--- + +# 45. README rendering + +Ensure `README.md` renders acceptably as package long description. + +Avoid unsupported repository-only Markdown extensions if they break package rendering. + +Do not rewrite the whole README solely for PyPI aesthetics. + +--- + +# 46. License audit + +Verify the repository has a clear license file if the project is intended for distribution. + +Ensure `pyproject.toml` metadata points to it correctly. + +Do not invent a license. + +If no license decision exists, report it as a release blocker rather than choosing one. + +--- + +# 47. Citation / scientific attribution + +If the project has a citation requirement, consider adding or validating: + +```text +CITATION.cff +``` + +only if appropriate and supported by actual project information. + +Do not fabricate publication titles, DOIs, author lists, or affiliations. + +If this information is unavailable, leave it as a documented follow-up. + +--- + +# 48. Changelog + +Add: + +```text +CHANGELOG.md +``` + +if absent and useful. + +Start with the rewrite/release milestone. + +Use a concise structure such as: + +```text +Unreleased +Added +Changed +Fixed +Compatibility +``` + +Do not reconstruct an imaginary historical changelog. + +--- + +# 49. Release notes + +Document the major v1 rewrite characteristics: + +```text +new layered architecture +named metadata replacing tracking arrays +ROOT-GNN parity +transfer learning +checkpointing +inference +Hydra/CLI +DDP/Perlmutter +ONNX export +``` + +Only list functionality actually complete. + +--- + +# 50. Semantic version + +If this will be the first stable rewrite release, choose an appropriate version only if the repository already has a release/version policy. + +Do not arbitrarily declare `1.0.0` if maintainers have not decided that. + +Task 19 may establish mechanics without incrementing the final release version. + +--- + +# 51. Release workflow + +If appropriate, add a manual GitHub Actions release workflow. + +Prefer a safe trigger such as: + +```text +workflow_dispatch +``` + +or: + +```text +tag push +``` + +depending on repository conventions. + +Do not automatically publish packages from every main-branch push. + +--- + +# 52. Publishing credentials + +Never hardcode package registry credentials. + +Use repository/environment secrets. + +If publishing is not actually configured, build a release-ready workflow without assuming secret names beyond documented placeholders. + +--- + +# 53. Trusted publishing + +If the target registry supports trusted publishing and repository policy permits it, structure the release workflow so it can use it. + +Do not configure external project identifiers by guessing. + +Leave clearly documented placeholders if necessary. + +--- + +# 54. Dry-run release + +Provide a safe release-validation path that: + +```text +builds artifacts +checks metadata +installs artifacts +runs smoke tests +``` + +without publishing. + +This should be runnable locally and/or in CI. + +--- + +# 55. Release checklist + +Add a concise release checklist. + +For example: + +```text +[ ] version updated +[ ] changelog updated +[ ] full tests pass +[ ] parity tests pass +[ ] lint/format pass +[ ] wheel/sdist build +[ ] clean wheel install succeeds +[ ] CLI smoke succeeds +[ ] docs current +[ ] tag/release prepared +``` + +Use actual project needs. + +--- + +# 56. Make release validation scriptable + +If useful, add a thin script such as: + +```text +scripts/dev/check_release.sh +``` + +or a Python equivalent. + +It may run: + +```text +lint +tests +build +artifact validation +clean-install smoke +``` + +Do not duplicate logic excessively with CI. + +Prefer the same underlying commands locally and in CI. + +--- + +# 57. Shell portability + +If shell scripts are added, use: + +```bash +set -euo pipefail +``` + +where appropriate. + +Do not hardcode site-specific paths. + +Validate syntax. + +--- + +# 58. CI must not depend on Slurm + +Do not run: + +```text +sbatch +srun against NERSC +``` + +in normal CI. + +Perlmutter/Slurm scripts may receive static syntax checks only. + +--- + +# 59. Slurm script syntax check + +Continue validating: + +```bash +bash -n scripts/slurm/*.sh +``` + +where scripts exist. + +Do not treat this as proof of Perlmutter correctness; it is syntax validation only. + +--- + +# 60. Benchmark scripts + +Do not run full Task 15 benchmarks in CI. + +At most add a smoke import/argument test. + +Performance assertions are inappropriate for shared CI hardware. + +--- + +# 61. Docs smoke validation + +Where practical, test documented public imports and CLI commands. + +Do not add a full documentation website toolchain. + +--- + +# 62. Config composition CI + +Validate committed Hydra reference configs compose. + +This catches stale configuration after code changes. + +Test representative: + +```text +multiclass pretraining +binary fine-tuning +resume/evaluation config if applicable +Perlmutter profile composition +export config +``` + +Do not execute production workloads. + +--- + +# 63. Test markers + +Review pytest markers. + +Document them in: + +```text +pyproject.toml +``` + +or pytest configuration. + +Examples may include: + +```text +integration +distributed +legacy_env +gpu +``` + +Use only markers the suite actually needs. + +Avoid unregistered-marker warnings. + +--- + +# 64. GPU tests + +Mark tests that genuinely require CUDA. + +Do not make them part of normal CPU CI unless a GPU runner exists. + +Ensure important logic also has CPU coverage where feasible. + +--- + +# 65. Perlmutter-only tests + +Any tests requiring: + +```text +NERSC filesystem +Slurm +specific CUDA/DGL environment +``` + +must be opt-in/manual. + +Document how to run them. + +Do not let site-specific tests contaminate ordinary installation. + +--- + +# 66. Coverage + +Do not chase a coverage percentage arbitrarily. + +If pytest-cov is already used, it is acceptable to report coverage. + +Do not fail CI on a newly invented high threshold unless project policy requires it. + +Focus on meaningful behavioral tests. + +--- + +# 67. Type checking + +Do not add mypy/pyright as a required CI gate unless the project has already adopted type checking. + +If type annotations are strong and a checker is already configured, run it. + +Do not turn Task 19 into a repo-wide typing migration. + +--- + +# 68. Security scanning + +Do not add multiple dependency/security bots/scanners in this task unless already part of repository policy. + +A simple dependency update mechanism may be considered separately. + +Keep scope focused on packaging/release correctness. + +--- + +# 69. Dependency pinning policy + +Document the chosen policy. + +For example: + +```text +direct dependency ranges in pyproject +fully resolved uv.lock for development +``` + +Do not over-pin every transitive dependency manually. + +--- + +# 70. Reproducible development setup + +Validate: + +```bash +git clone ... +uv sync --extra root-gnn +uv run pytest +``` + +from a clean checkout, to the extent external DGL installation permits. + +Document any extra DGL/index step precisely. + +--- + +# 71. Editable vs built install + +Make clear that: + +```text +uv sync +``` + +is the development workflow, while wheels/sdists represent release installation. + +Test both. + +Do not assume editable/source install behavior guarantees wheel correctness. + +--- + +# 72. Installed config discovery + +This deserves a dedicated test. + +From a clean installed wheel, verify the CLI can find: + +```text +Hydra root config +model configs +task configs +trainer configs +environment configs +export configs +``` + +without the repository `configs/` directory being present. + +Fix package-resource resolution if necessary. + +--- + +# 73. Resource access + +Use: + +```text +importlib.resources +``` + +or an equivalent standard installed-package resource mechanism where appropriate. + +Do not derive package-data paths by traversing upward from `__file__` into repository layout. + +--- + +# 74. Move configs only if necessary + +If Task 13 put configs at repository root and packaging them reliably requires relocation, consider moving canonical runtime configs under: + +```text +src/gnn4colliders/configs/ +``` + +or another package-data location. + +If you move them: + +* update Hydra discovery +* update docs +* update tests +* avoid duplicate canonical copies + +Do not keep two divergent config trees. + +--- + +# 75. Config-source-of-truth rule + +There must be one canonical runtime config tree. + +Examples/documentation may reference it. + +Do not maintain: + +```text +configs/ +``` + +and: + +```text +src/gnn4colliders/configs/ +``` + +as separate manually synchronized copies. + +--- + +# 76. CLI entry point + +Verify `pyproject.toml` exposes exactly the intended command. + +For example: + +```toml +[project.scripts] +gnn4colliders = "gnn4colliders.cli:main" +``` + +Use the actual module. + +Add installed-wheel smoke coverage. + +--- + +# 77. Module execution + +If useful, also support: + +```bash +python -m gnn4colliders +``` + +but only if this is easy and intentional. + +Do not add it merely for symmetry. + +The installed console script is the primary requirement. + +--- + +# 78. Exit codes + +CLI commands should return nonzero exit status on user-facing failure. + +Packaging smoke tests can verify `--help` returns success. + +Do not refactor the entire error-handling system. + +--- + +# 79. Package metadata classifiers + +Add only accurate classifiers if used. + +Do not guess: + +```text +development status +license +OS support +GPU support +``` + +Prefer minimal correct metadata to verbose speculative metadata. + +--- + +# 80. URLs + +If official repository/documentation URLs are known from the existing project metadata, ensure they are correct. + +Do not invent a homepage or issue tracker. + +--- + +# 81. Artifact size + +Inspect wheel/sdist sizes. + +Unexpectedly large artifacts often indicate accidental inclusion of: + +```text +legacy data +checkpoints +ROOT files +ONNX models +profiling traces +``` + +Investigate unusual size. + +Do not enforce an arbitrary size limit. + +--- + +# 82. `.gitignore` + +Review `.gitignore` for generated release/development artifacts. + +Ensure it covers appropriate items such as: + +```text +dist/ +build/ +*.egg-info/ +__pycache__/ +outputs/ +checkpoints/ +profiles/ +benchmark-results/ +``` + +Use actual project needs. + +Do not ignore source/config files. + +--- + +# 83. Build artifacts + +Do not commit: + +```text +dist/*.whl +dist/*.tar.gz +``` + +to normal source control unless repository policy explicitly requires vendored releases. + +CI/release systems should create them. + +--- + +# 84. Pre-commit hooks + +Do not introduce pre-commit/prek as a required developer workflow unless already decided. + +CI lint/test gates are sufficient for this task. + +If existing hooks are present, ensure they match CI. + +--- + +# 85. Branch protection + +Do not attempt to configure GitHub branch protection from repository code. + +You may document recommended required checks if useful. + +Keep repository changes within source-controlled capabilities. + +--- + +# 86. CI job names + +Use stable, meaningful job names because they may later become required checks. + +For example: + +```text +lint +test +package +onnx +``` + +Avoid overly specific temporary names. + +--- + +# 87. Fail fast but preserve diagnostics + +Use job separation or step ordering so basic formatting failures do not obscure packaging failures unnecessarily. + +Keep CI readable. + +Do not over-optimize workflow runtime at the cost of coverage. + +--- + +# 88. CI concurrency + +If appropriate, use standard concurrency cancellation for superseded pull-request runs. + +Do not add complex concurrency behavior. + +--- + +# 89. Release workflow permissions + +Use minimum required GitHub Actions permissions. + +Do not grant broad write permissions unnecessarily. + +If publishing is not enabled, release validation should remain read-only. + +--- + +# 90. Supply-chain scope + +Do not implement signing/SBOM/provenance systems unless repository policy already requires them. + +These can be follow-up hardening tasks. + +The priority is correct package build and installation. + +--- + +# 91. Release documentation + +Add: + +```text +docs/releasing.md +``` + +if helpful. + +Cover: + +```text +version update +changelog +local validation +CI validation +build +artifact inspection +tag/release +publish step if configured +``` + +Keep it concise enough to follow. + +--- + +# 92. README installation update + +Ensure README distinguishes: + +```text +developer source install +released package install +optional ROOT-GNN dependencies +optional ONNX dependencies +``` + +Only document package-registry install commands if the package is actually published there. + +Do not claim availability on PyPI unless true. + +--- + +# 93. Unreleased-package wording + +If the package is not yet published, say so clearly. + +Use source/uv installation as the current supported path. + +Task 19 may prepare release mechanics without actually publishing. + +--- + +# 94. No release publication unless explicitly requested + +Do not publish to PyPI or create a remote GitHub release as part of this task unless explicitly instructed. + +Build and validate locally/CI only. + +Do not push tags automatically. + +--- + +# 95. Preserve legacy frozen state + +Do not modify scientific legacy source. + +Packaging configuration may exclude: + +```text +legacy/ +``` + +from distributions, but do not delete or rewrite it. + +--- + +# 96. Scientific regression guard + +Run the full parity/test suite before declaring packaging changes complete. + +This task should be behavior-neutral. + +If packaging changes require moving runtime configs, ensure all scientific config semantics remain unchanged. + +--- + +# 97. Clean-checkout validation + +Where practical, test from a fresh checkout or temporary copy. + +The goal is to catch hidden dependencies on: + +```text +untracked files +local editable metadata +stale caches +personal environment variables +``` + +Document anything that cannot be fully automated. + +--- + +# 98. Validation + +Run: + +```bash +uv sync --extra root-gnn +``` + +and optional extras as applicable. + +Run lint: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +Run tests: + +```bash +uv run pytest +``` + +Run focused distributed tests: + +```bash +uv run pytest tests/unit/distributed -v +uv run pytest tests/integration -k distributed -v +``` + +Run ONNX tests using the appropriate extra: + +```bash +uv run pytest tests/unit/export -v +uv run pytest tests/integration -k "export or onnx" -v +``` + +Build artifacts: + +```bash +uv build +``` + +or the final project-standard equivalent. + +Validate artifact metadata with the chosen standard tooling. + +Inspect wheel/sdist contents. + +Install the wheel into a fresh temporary environment. + +From a directory outside the repository, validate: + +```bash +python -c "import gnn4colliders; print(gnn4colliders.__version__)" +gnn4colliders --help +gnn4colliders train --help +gnn4colliders evaluate --help +gnn4colliders predict --help +gnn4colliders export --help +``` + +Validate packaged Hydra configuration discovery. + +If Slurm scripts exist: + +```bash +bash -n scripts/slurm/*.sh +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* no scientific behavior changed +* no legacy code was modified +* no built wheel/sdist is accidentally tracked +* canonical runtime configs work from an installed wheel +* optional dependencies remain optional +* CI does not require GPU/Slurm/NERSC +* no publishing occurred +* no unrelated changes are included + +--- + +# Completion criteria + +Task 19 is complete when: + +1. `pyproject.toml` accurately describes the package. +2. Versioning has one clear source of truth. +3. Python support metadata matches validated support. +4. Dependencies and optional extras are cleanly separated. +5. DGL installation requirements are documented accurately. +6. `uv.lock` is consistent with project metadata. +7. Wheel build succeeds. +8. Sdist build succeeds. +9. Distribution contents have been audited. +10. `legacy/` and generated artifacts are not accidentally packaged. +11. Required runtime configs/resources are included. +12. Installed CLI works outside the source repository. +13. Installed Hydra configuration discovery works. +14. Clean wheel installation succeeds. +15. Sdist installation succeeds where practical. +16. Public imports work from the installed artifact. +17. Base installation does not require optional ONNX dependencies. +18. Optional ROOT-GNN behavior is validated according to final packaging design. +19. Optional ONNX behavior is validated according to final packaging design. +20. CI runs lint/format checks. +21. CI runs core unit tests. +22. CI runs appropriate integration tests. +23. CI runs modern parity coverage. +24. CPU distributed tests run in CI where stable. +25. ONNX CPU tests run in an appropriate optional job where stable. +26. CI builds wheel and sdist. +27. CI clean-installs the built wheel and smoke-tests it. +28. Reference Hydra configs are validated in CI. +29. Release-validation commands are documented/scriptable. +30. Changelog/release documentation exists where appropriate. +31. No package publication occurs without explicit instruction. +32. Full test suite passes. +33. Legacy source remains frozen. +34. No scientific behavior changes are included. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. packaging metadata changes +4. versioning strategy +5. supported Python metadata +6. core dependency structure +7. optional extras structure +8. DGL installation handling +9. runtime package-data/config strategy +10. config source-of-truth location +11. wheel build result +12. sdist build result +13. wheel contents audit +14. sdist contents audit +15. clean wheel installation result +16. clean sdist installation result +17. installed public-import results +18. installed CLI results +19. installed Hydra config discovery results +20. optional dependency installation results +21. CI workflow structure +22. lint CI result +23. unit/integration CI coverage +24. parity CI coverage +25. distributed CI coverage +26. ONNX CI coverage +27. package-build CI coverage +28. release-validation workflow +29. changelog/release-doc changes +30. remaining release blockers +31. publishing steps intentionally not performed +32. validation commands and results + +After validation succeeds, create one Git commit containing only Task 19 changes. + +Use: + +```text +build: harden packaging and CI +``` + +Before committing, inspect the final diff and ensure no unrelated scientific, model, or migration changes are included. diff --git a/tasks/task2.md b/tasks/task2.md new file mode 100644 index 0000000000000000000000000000000000000000..615f5fb11045ed4518fb33924458b534beabf714 --- /dev/null +++ b/tasks/task2.md @@ -0,0 +1,549 @@ +# Task 2: Environment and Developer Tooling Setup + +Set up a modern, reproducible Python development environment for **GNN4Colliders**. + +This task is about: + +* dependency management +* packaging/runtime dependencies +* development tooling +* linting/formatting +* tests +* reproducibility of the developer environment + +Do **not** begin migrating legacy ML implementation in this task. + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +pyproject.toml +README.md +``` + +Also inspect the current Git status and existing project structure. + +--- + +## Goals + +After this task, a developer should be able to clone the repository and run something close to: + +```bash +uv sync + +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +The environment should be simple, reproducible, and appropriate for future PyTorch/DGL development. + +Use **uv** as the Python package/environment manager. + +Use **Ruff** for linting and formatting. + +Use **pytest** for testing. + +Do not introduce unnecessary tooling or runtime optimizations yet. + +--- + +# 1. Use uv for project dependency management + +Configure the repository to use `uv`. + +The repository should contain a generated: + +```text +uv.lock +``` + +Use the existing `pyproject.toml` rather than introducing a separate requirements file unless there is a compelling compatibility reason. + +Do not add Poetry, Pipenv, Conda metadata, or duplicate package-management systems. + +The intended common workflow should be: + +```bash +uv sync +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +Document this workflow in the README. + +--- + +# 2. Python version + +Choose a modern Python version compatible with the scientific stack we expect to use. + +Prefer Python 3.12 unless current dependency compatibility clearly requires otherwise. + +Set the supported Python version explicitly in `pyproject.toml`. + +For the ROOT-GNN environment, use `requires-python = ">=3.12,<3.13"` and add +a `.python-version` file containing `3.12`. + +Do not inherit the legacy Python 3.8 requirement. + +If PyTorch/DGL compatibility prevents using Python 3.12, report the issue rather than silently changing the target. + +--- + +# 3. Runtime dependencies + +Add the clearly required non-GPU-sensitive scientific dependencies: + +```text +numpy +awkward +uproot +scikit-learn +hydra-core +``` + +These are expected to support: + +* ROOT/Awkward data access +* numerical operations +* metrics +* experiment configuration + +Do not add speculative dependencies. + +--- + +# 4. PyTorch and DGL + +The project will require: + +```text +torch +dgl +``` + +However, GPU/CUDA compatibility must be treated carefully. + +Before selecting versions: + +1. inspect any relevant current environment/module information already present in the repository +2. inspect the legacy environment only as a reference +3. determine whether a straightforward modern PyTorch/DGL installation can be represented safely in the project dependency configuration + +Do not blindly preserve legacy versions such as PyTorch 2.0.1 or DGL 1.1.1. + +Do not hardcode CUDA-specific wheels unless there is a clear reason and the chosen approach is documented. + +If DGL/PyTorch/CUDA installation requires platform-specific handling, structure the project so that: + +* normal developer dependencies remain clean +* Perlmutter/HPC installation guidance can be documented separately +* the package metadata does not unnecessarily encode one machine's CUDA stack + +If there is material uncertainty about the correct DGL/PyTorch version combination, document the issue and defer the exact GPU installation choice rather than guessing. + +Do not add PyTorch Geometric. + +### Compatibility handoff to Task 2B + +When implementing both tasks from a clean checkout, use this verified +Linux/Perlmutter tuple rather than resolving the latest packages: + +```text +Python 3.12 +PyTorch 2.2.2+cu121 +DGL 2.4.0+cu121 +CUDA runtime 12.1 +``` + +DGL is not available as a usable Linux wheel from PyPI. The verified official +source is the flat index below; pin the DGL extra as `dgl==2.4.0+cu121` and +PyTorch as `torch==2.2.2`, sourced from the official CUDA 12.1 PyTorch index. +Constrain NumPy to `<2` for this PyTorch build, and include matplotlib in the +development dependencies because the legacy parity import path requires it. + +```toml +[[tool.uv.index]] +name = "dgl" +url = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html" +format = "flat" +explicit = true + +[tool.uv.sources] +dgl = { index = "dgl" } +``` + +--- + +# 5. Development dependency group + +Create a development dependency group using the modern `pyproject.toml` dependency-group mechanism if supported by the selected tooling. + +Include: + +```text +pytest +pytest-cov +ruff +``` + +Do not add Black, Flake8, isort, or similar tools whose role is already covered by Ruff. + +Do not add Nox yet. + +Do not add Numba yet. + +Do not add Lightning, Kedro, DeepSpeed, CuPy, or other architectural/performance frameworks. + +--- + +# 6. Ruff configuration + +Configure Ruff in `pyproject.toml`. + +Use Ruff for both: + +* linting +* formatting + +Adopt sensible defaults appropriate for a scientific Python package. + +At minimum, enable checks covering: + +* Pyflakes-style errors +* pycodestyle errors +* import sorting + +Do not turn on an extremely aggressive rule set that produces large amounts of low-value churn. + +Prefer maintainability and signal over maximum lint strictness. + +Configure an appropriate line length, preferably around 88–100 characters. + +Do not rewrite legacy code solely to satisfy Ruff. + +Exclude or appropriately scope linting so the frozen `legacy/` tree does not become part of the new-code lint burden unless explicitly intended. + +The new codebase should be linted. + +--- + +# 7. Pytest configuration + +Configure pytest in `pyproject.toml`. + +Set sensible test discovery for: + +```text +tests/ +``` + +Ensure running: + +```bash +uv run pytest +``` + +does not accidentally collect tests from: + +```text +legacy/ +``` + +unless parity tests intentionally invoke legacy behavior. + +The project currently has: + +```text +tests/unit/ +tests/integration/ +tests/parity/ +tests/fixtures/ +``` + +Preserve that structure. + +Add only the minimum smoke test necessary to verify the new package imports correctly, if no tests currently exist. + +For example, an import/package test is acceptable. + +Do not implement ML behavior in tests during this task. + +--- + +# 8. Git ignore cleanup + +Review `.gitignore`. + +Ensure generated development artifacts are ignored, including at least: + +```text +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ +.ruff_cache/ +.coverage +htmlcov/ +.venv/ +``` + +Also ensure local environment artifacts and generated Python packaging metadata are not tracked. + +Inspect whether any of these artifacts are currently tracked. + +If generated artifacts such as: + +```text +src/gnn4colliders.egg-info/ +__pycache__/ +``` + +are tracked, remove them from Git tracking without altering useful source files. + +Do not delete legacy model artifacts merely because they are large; the legacy tree is out of scope unless explicitly addressed. + +--- + +# 9. Pre-commit / prek + +Do **not** make Git hooks a hard dependency for this task. + +However, evaluate whether a small hook configuration would be useful. + +If adding one, keep it limited to fast checks such as: + +```text +ruff check +ruff format --check +``` + +Do not run the full test suite on every commit. + +Prefer `prek` only if it integrates cleanly with the current project and does not create additional complexity. + +If there is no clear benefit yet, defer hook tooling and document it as a future option. + +--- + +# 10. Do not add Nox yet + +Do not introduce Nox in this task. + +The current development commands are simple enough to be expressed directly as: + +```bash +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +Nox may be introduced later if the project needs: + +* multi-version Python testing +* CPU/GPU test environments +* packaging validation +* documentation builds +* more complex CI orchestration + +--- + +# 11. Do not add runtime optimization dependencies + +Do not add optimization packages merely because they may improve speed later. + +Specifically, do not add: + +```text +numba +cupy +deepspeed +lightning +torch-geometric +``` + +or similar dependencies in this task. + +Runtime performance will be addressed after behavioral parity exists and profiling identifies bottlenecks. + +Performance-critical areas may eventually include: + +* ROOT/Awkward ingestion +* feature construction +* graph construction +* batching +* DGL loading +* PyTorch training +* distributed training + +Do not optimize them yet. + +--- + +# 12. README update + +Update the README with a concise developer setup section. + +Include: + +```bash +git clone ... +cd GNN4Colliders + +uv sync + +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +If GPU-specific PyTorch/DGL setup requires additional steps, explain that clearly and separately. + +On Perlmutter, uv's default home-directory cache can exhaust quota during +large PyTorch/DGL installs. If extraction fails with `Disk quota exceeded`, +retry with a scratch-backed cache: + +```bash +export UV_CACHE_DIR=/pscratch/sd/$USER/uv-cache-gnn4colliders +uv sync --extra root-gnn +``` + +Do not imply that training or inference is implemented yet. + +Document the distinction between: + +* general developer environment +* future GPU/HPC runtime environment + +--- + +# 13. AGENTS.md update + +If necessary, add a short environment/tooling section to `AGENTS.md`. + +It should establish these conventions: + +```text +uv -> dependency/environment management +ruff -> linting and formatting +pytest -> testing +``` + +Future agents should use: + +```bash +uv run ... +``` + +for project commands where practical. + +Agents should not add new dependencies casually. + +Any new major dependency should have a clear architectural or runtime justification. + +Avoid duplicating extensive instructions already present in AGENTS.md. + +--- + +# 14. Optional convenience commands + +Do not introduce Makefiles, task runners, or shell wrappers just to alias commands such as: + +```bash +uv run pytest +uv run ruff check . +``` + +Keep the tooling stack minimal. + +If there is a compelling reason for a task runner, report it instead of adding one automatically. + +--- + +# Validation + +After making changes, validate the environment. + +Run the relevant commands, preferably: + +```bash +uv sync +uv run python -c "import gnn4colliders; print(gnn4colliders.__file__)" +uv run pytest +uv run ruff check . +uv run ruff format --check . +``` + +Also inspect: + +```bash +git status +git diff +``` + +Ensure: + +* generated `.egg-info` is not tracked +* `__pycache__` is not tracked +* legacy code has not been reformatted or otherwise modified +* no ML implementation has been added +* the lockfile is present if dependency resolution succeeded + +If PyTorch/DGL cannot be installed because of platform/CUDA constraints, do not conceal the failure. + +For ROOT-GNN validation, a skipped DGL test is not success. The current parity +characterization suite also contains a separate expectation issue: its +two-node no-self-loop edge test expects four edges, while the legacy graph +constructor produces two. Report that failure without changing legacy code or +weakening the test in this environment task. + +Report: + +* what succeeded +* what failed +* the relevant compatibility issue +* the recommended next step + +--- + +# Git + +After validation succeeds, create a single Git commit containing only Task 2 changes. + +Use a concise commit message such as: + +```text +chore: set up development environment and tooling +``` + +Before committing: + +* inspect `git diff` +* ensure no unrelated files are included +* ensure generated artifacts are excluded +* ensure relevant validation passes + +Do not amend or rewrite previous commits. + +--- + +# Completion report + +At the end, report: + +1. files created +2. files modified +3. runtime dependencies added +4. development dependencies added +5. Python version selected +6. PyTorch/DGL handling decision +7. validation commands run and results +8. anything intentionally deferred +9. commit hash and commit message + +Keep this task limited to environment and developer tooling setup. diff --git a/tasks/task20.md b/tasks/task20.md new file mode 100644 index 0000000000000000000000000000000000000000..1e91591697e63b70d509acb633c4f2383014c22a --- /dev/null +++ b/tasks/task20.md @@ -0,0 +1,2164 @@ +# Task 20: Expand GNN4Colliders Test Coverage and Regression Protection + +Strengthen the GNN4Colliders automated test suite so it protects not only the currently characterized happy paths, but also invariants, failure modes, serialization boundaries, configuration behavior, numerical edge cases, distributed execution, export dynamic-shape behavior, and complete end-to-end workflows. + +The repository currently has approximately: + +```text +25 test files +75 collected tests +74 passing +1 skipped when ONNX dependencies are unavailable +``` + +Existing coverage already includes: + +```text +unit: + features + graphs + data + models/root_gnn + tasks + training + inference + compat + distributed + export + package/public API + +integration: + synthetic ROOT workflows + reduced real ROOT fixture workflows + model/pretraining/transfer workflows + +parity: + graph topology + edge features + node features + dataset semantics + model forward + transfer learning + checkpoint prefix handling +``` + +Do not simply add more tests around behavior that is already well protected. + +The goal is to identify important untested contracts and add high-value regression coverage. + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +pyproject.toml +``` + +Then inspect all current tests: + +```text +tests/unit/ +tests/integration/ +tests/parity/ +``` + +and relevant production code under: + +```text +src/gnn4colliders/ +``` + +Also inspect: + +```text +configs/ +scripts/ +benchmarks/ +.github/workflows/ +``` + +where present. + +--- + +# Goal + +Develop a substantially more comprehensive automated test suite that protects the current ROOT-GNN v1 implementation against: + +```text +silent numerical regressions +shape/schema regressions +metadata misalignment +cache corruption/incompatibility +checkpoint corruption/incompatibility +configuration drift +batching inconsistencies +device-placement bugs +distributed-only regressions +export dynamic-shape bugs +unexpected input edge cases +public API/package regressions +``` + +Prioritize tests of contracts and invariants, not implementation details. + +A good test should remain valid through internal refactors as long as externally relevant behavior is preserved. + +--- + +# 1. Begin with a test-gap audit + +Before writing tests, inspect the existing suite and produce a short internal test matrix. + +For each major subsystem classify coverage as: + +```text +happy path +boundary conditions +invalid input +serialization round trip +cross-component integration +parity +distributed +optional dependency +``` + +Subsystems should include at minimum: + +```text +data / ROOT I/O +metadata +features +graphs +cache +splits +batching +models +transfer learning +tasks / losses / metrics +trainer +checkpointing +inference +output writers +compatibility +configuration / factories +CLI +distributed +ONNX/export +packaging/resources +``` + +Use this audit to avoid duplicate low-value tests. + +--- + +# 2. Test architecture-level invariants + +Add tests for invariants that should remain true across implementation changes. + +Examples: + +```text +sample identity remains aligned with labels and predictions +batching does not alter per-sample model outputs +cache round-trip does not alter scientific content +checkpoint round-trip does not alter inference output +different inference batch sizes do not alter per-event predictions +evaluation mode does not mutate model parameters +modern APIs never expose positional tracking tensors +``` + +Prefer these over tests of private helper call sequences. + +--- + +# 3. Expand feature tests + +Current feature tests already cover schema, values, empty vectors, and immutability. + +Add high-value cases such as: + +```text +single-object event +mixed object collections +all supported node/object types +extreme but finite pt/eta/phi values +phi near -pi and +pi +zero pt +negative eta +feature scaling boundary cases +input float32 vs float64 behavior +``` + +Test that: + +```text +feature column order is invariant +feature output shape matches object count +input arrays are never mutated +empty collections concatenate correctly +``` + +Do not change scientific feature definitions. + +--- + +# 4. Feature metamorphic tests + +Add metamorphic/invariant tests where useful. + +Examples: + +```text +adding one object increases node count by exactly one +reordering input objects produces the corresponding reordered feature rows +changing only phi changes only phi-dependent fields +changing scaling changes only the expected scaled columns +``` + +Use deterministic fixtures. + +--- + +# 5. Expand graph topology tests + +Current tests already cover fully connected topology, source-major ordering, zero nodes, self-loop behavior, and invalid counts. + +Add broader topology invariants. + +For multiple `N` values, verify: + +```text +edge count = N * (N - 1) +every ordered pair i != j appears exactly once +no self-loop appears +every valid source has N - 1 outgoing edges +every valid destination has N - 1 incoming edges +``` + +Test several sizes, not only 1–3 nodes. + +Keep runtime small. + +--- + +# 6. Graph edge-feature invariants + +Expand edge-feature tests beyond fixed examples. + +For every edge: + +```text +deta(i,j) = eta_i - eta_j +dphi is wrapped into the documented interval +dR >= 0 +dR(i,j) == dR(j,i) +deta(i,j) == -deta(j,i) +wrapped dphi(i,j) == -wrapped dphi(j,i) where not on branch-cut ambiguity +``` + +Include phi values near wrapping boundaries. + +--- + +# 7. Graph input validation + +Add clear failure tests for malformed graph inputs: + +```text +wrong node feature width +nonmatching feature lengths +NaN/Inf where prohibited +invalid node counts +invalid edge indices +wrong dtype where unsupported +``` + +Test public APIs rather than internal helpers when possible. + +Do not over-constrain behavior that is intentionally supported. + +--- + +# 8. EventMetadata tests + +Strengthen named-metadata contracts. + +Test: + +```text +fold +weight +sample_id +extra metadata +``` + +including: + +```text +zero weight +negative weight +very large finite weight +string sample IDs +metadata immutability if intended +extra metadata preservation +``` + +Ensure no normal modern path reconstructs positional tracking arrays. + +--- + +# 9. Stable sample identity tests + +If `sample_id` is derived from: + +```text +source file +tree +entry index +``` + +or equivalent, add tests that verify: + +```text +same source event -> same sample_id +different entry -> different sample_id +different source file -> different sample_id +dataset/cache round-trip -> same sample_id +``` + +Do not depend on absolute temporary-directory names if the design intentionally normalizes them. + +--- + +# 10. ROOT I/O boundary tests + +Expand ROOT ingestion tests using tiny generated ROOT files. + +Cover: + +```text +single event +multiple files +multiple events +empty jagged collections +optional branch presence/absence +missing required branch +wrong tree name +empty tree +``` + +Verify deterministic file/event ordering. + +Do not require production data for these tests. + +--- + +# 11. Multi-file ordering + +Add explicit tests for multiple ROOT input files. + +Verify the documented ordering contract across: + +```text +file order +entry order within file +sample_id generation +labels +metadata +``` + +This is important for inference alignment. + +--- + +# 12. ROOT error-path tests + +Add clear error tests for: + +```text +missing file +missing tree +missing required branch +incompatible branch shape +malformed configured branch name +``` + +Test for informative exception types/messages without overfitting exact prose. + +--- + +# 13. Split invariants + +Expand split tests. + +Verify: + +```text +train/validation/test are disjoint +their union matches expected selected events +split order is deterministic +fold-based splitting does not modify metadata +``` + +Test invalid configurations such as: + +```text +overlapping folds +unknown fold +empty required split +duplicate fold declarations +``` + +according to current intended behavior. + +--- + +# 14. Cache corruption and invalidation + +Current cache tests cover round-trips and schema validation. + +Add tests for: + +```text +wrong feature schema version +wrong graph schema version +wrong cache schema version +preprocessing fingerprint mismatch +truncated/corrupt cache file +missing cache metadata +metadata/sample count mismatch +``` + +The loader should fail clearly instead of silently accepting incompatible artifacts. + +--- + +# 15. Cache determinism + +Given identical input/configuration, verify cache generation produces semantically identical cached samples. + +Do not require byte-for-byte identical serialization unless the format guarantees it. + +Compare: + +```text +sample IDs +features +graphs +labels +metadata +``` + +--- + +# 16. Batch collation tests + +Add stronger `GraphBatch` tests. + +Verify: + +```text +graph count +node/edge counts +label order +global-feature order +weight order +sample-id order +``` + +for multiple heterogeneous graph sizes. + +Include: + +```text +1 graph +multiple graphs +single-node graph +mixed small/large graphs +``` + +--- + +# 17. Batch unbatch equivalence + +Where DGL supports the operation, test the invariant: + +```text +individual GraphSamples + -> batch + -> unbatch +``` + +preserves per-graph: + +```text +node features +edge features +topology +``` + +and verify the surrounding batch metadata remains aligned. + +Do not require exact object identity. + +--- + +# 18. Model shape matrix + +Current model tests cover standard output shapes. + +Expand to a compact parameter matrix over: + +```text +batch size +number of nodes +number of processing steps +hidden dimension +output size +with/without global features +``` + +Avoid combinatorial explosion. + +Select representative boundary combinations. + +--- + +# 19. Model gradient tests + +Add tests verifying gradients flow through expected trainable parameters. + +For a tiny batch: + +```text +forward +loss +backward +``` + +Verify: + +```text +trainable model parameters receive finite gradients +frozen parameters receive no gradients +classifier receives gradients during fine-tuning +unfrozen backbone receives gradients +``` + +Do not test exact gradient values unless needed for parity. + +--- + +# 20. Gradient finiteness + +For representative deterministic inputs, verify: + +```text +loss is finite +trainable gradients are finite +updated parameters remain finite +``` + +Use ordinary float32 reference execution. + +--- + +# 21. Repeated-forward invariants + +Test repeated inference on the same graph/model. + +Verify: + +```text +same input + eval mode -> same output +graph node/edge data does not accumulate temporary fields +input graph is not persistently mutated +``` + +This protects scoped temporary DGL state. + +--- + +# 22. Model serialization invariance + +Test: + +```text +model state_dict + -> fresh identical model + -> load state_dict +``` + +and verify identical logits. + +This should be a fast unit-level regression test separate from full checkpoint tests. + +--- + +# 23. Transfer-learning structural tests + +Expand transfer tests to verify: + +```text +pretrained classifier is not accidentally reused +backbone weights match source after transfer +new classifier is distinct from the source classifier +freeze/unfreeze changes only trainability state +freeze/unfreeze does not mutate parameter values +``` + +Test switching freeze state if the public API supports it. + +--- + +# 24. Loss reference-value tests + +Current loss tests cover binary/multiclass and weights. + +Add hand-computable reference cases: + +```text +equal logits +very confident correct logits +very confident incorrect logits +zero weights +negative weights +class imbalance +single represented class +``` + +Use explicit expected formulas. + +Do not calculate the expected result by calling the implementation under test. + +--- + +# 25. Loss permutation invariance + +If the loss should mathematically be invariant to event ordering, test: + +```text +jointly permuting logits +targets +weights +``` + +produces the same loss. + +This is valuable for weighted/class-balanced reduction semantics. + +--- + +# 26. Metrics permutation invariance + +Verify metrics are unchanged when: + +```text +scores +targets +weights +``` + +are jointly permuted. + +This protects against accidental order dependence. + +--- + +# 27. Metric degenerate cases + +Expand behavior for: + +```text +empty inputs +single sample +one represented class +all equal scores +zero-weight samples +negative weights where supported +``` + +Assert the documented result: + +```text +NaN +None +clear exception +``` + +rather than allowing misleading values. + +--- + +# 28. Trainer state-transition tests + +Expand trainer tests around lifecycle state. + +Verify: + +```text +global_step increments exactly once per optimizer step +epoch count is correct +evaluation does not increment global_step +validation does not alter model parameters +history length matches executed epochs +early stopping stops on the expected epoch +``` + +Pay particular attention to off-by-one boundaries. + +--- + +# 29. Optimizer update tests + +For a deterministic model/batch: + +```text +record parameters +train one step +``` + +Verify: + +```text +at least one trainable parameter changes +frozen parameters do not change +``` + +Do not rely only on successful backward execution. + +--- + +# 30. Scheduler tests + +For every scheduler currently supported, add a short expected learning-rate sequence. + +Test both: + +```text +fresh run +save/resume continuation +``` + +Do not add tests for unsupported schedulers. + +--- + +# 31. Early-stopping serialization + +Test: + +```text +early-stopping state + -> state_dict/checkpoint + -> new EarlyStopping object +``` + +and verify subsequent updates behave identically to an uninterrupted reference. + +--- + +# 32. Checkpoint corruption tests + +Expand checkpoint validation. + +Test: + +```text +missing schema version +unsupported schema version +missing model state +wrong model family +incompatible model config +wrong task metadata +malformed optimizer state +``` + +where the implementation is expected to detect these conditions. + +Do not require recovery from arbitrary binary corruption. + +--- + +# 33. Checkpoint architecture mismatch + +Verify loading fails for incompatible backbone architecture. + +For transfer loading: + +```text +classifier mismatch may be intentionally allowed +backbone mismatch must still fail +``` + +Do not allow broad `strict=False` behavior to hide invalid weights. + +--- + +# 34. Resume-vs-continuous training + +Strengthen the resume test. + +Compare: + +```text +Run A: + train continuously for N epochs + +Run B: + train K epochs + save + construct fresh model/trainer/optimizer/scheduler + restore + train remaining N-K epochs +``` + +Compare: + +```text +final model parameters +optimizer state +scheduler state +global_step +epoch +``` + +and training history where appropriate. + +Use CPU and explicit seeds. + +--- + +# 35. RNG restoration tests + +Where implemented, test restoration of: + +```text +Python random +NumPy +Torch CPU +``` + +and CUDA only in optional GPU tests. + +After restoration, verify subsequent generated sequences match the uninterrupted reference. + +--- + +# 36. Predictor batch-size invariance + +Run the same ordered dataset with: + +```text +batch_size = 1 +batch_size = 2 +batch_size = larger value +``` + +Verify identical per-sample: + +```text +sample IDs +logits +scores +predictions +``` + +within established numerical tolerances. + +--- + +# 37. Prediction ordering + +Use deliberately nontrivial sample IDs such as: + +```text +sample_10 +sample_2 +sample_7 +``` + +Verify inference does not implicitly sort them. + +--- + +# 38. Prediction metadata alignment + +For every output row verify correspondence among: + +```text +sample_id +label +weight +fold +logits +score +prediction +``` + +Use unique sentinel values so accidental permutation is easy to detect. + +--- + +# 39. Unlabeled inference + +Expand unlabeled inference coverage. + +Verify: + +```text +prediction succeeds +labels are represented as absent +label-dependent metrics fail or remain unavailable clearly +serialization handles missing labels correctly +``` + +--- + +# 40. NPZ schema tests + +Test the modern NPZ contract for: + +```text +required field names +dtype +shape +metadata alignment +optional-label behavior +``` + +Only maintain legacy NPZ coverage for formats still explicitly supported. + +--- + +# 41. NPZ failure cases + +Test: + +```text +existing output path +invalid parent directory where relevant +unsupported metadata dtype +inconsistent result lengths +``` + +according to current writer semantics. + +--- + +# 42. ROOT output tests + +If ROOT score writing exists, expand coverage for: + +```text +selected events +unselected events +selection_pass +binary scores +multiclass scores +original branch preservation +event ordering +``` + +Use generated tiny ROOT files. + +--- + +# 43. Legacy compatibility negative tests + +Current compatibility tests cover successful conversions. + +Add rejection tests for: + +```text +unknown tracking layout +too few tracking columns +unexpected legacy state-dict keys +partially matched backbone +unsupported experimental architecture +unsupported legacy checkpoint structure +``` + +Compatibility code should fail rather than guess. + +--- + +# 44. Modern paths must avoid legacy representations + +Add architectural regression coverage verifying modern workflows do not recreate positional tracking structures. + +At minimum ensure: + +```text +EventSample +GraphSample +GraphBatch +PredictionResult +``` + +use named metadata only. + +Avoid brittle mocks of internal helper functions. + +--- + +# 45. Configuration composition tests + +Add Hydra composition tests for canonical workflows. + +At minimum: + +```text +multiclass pretraining +binary training +fine-tuning +resume +evaluation +prediction +export +local environment +Perlmutter environment +DDP +``` + +These should compose without executing expensive workloads. + +--- + +# 46. Cross-config validation tests + +Test invalid combinations such as: + +```text +binary task + model out_size != 1 +multiclass task + incompatible class count +fine-tuning without pretrained checkpoint +resume + pretrained simultaneously when prohibited +negative batch size +zero epochs +overlapping folds +unsupported output format +unsupported model family +``` + +These should fail before expensive initialization. + +--- + +# 47. Factory tests + +Test semantic mappings: + +```text +config -> model +config -> task +config -> optimizer +config -> scheduler +config -> trainer +``` + +Unsupported names should fail clearly. + +Do not test Hydra's own internals. + +--- + +# 48. CLI end-to-end smoke tests + +Using tiny temporary fixtures, exercise: + +```text +prepare +train +evaluate +predict +export when ONNX is installed +``` + +Keep runs minimal. + +For training: + +```text +CPU +tiny data +1 epoch +``` + +is sufficient. + +--- + +# 49. CLI failure tests + +Verify nonzero failure for: + +```text +missing input file +missing checkpoint +invalid task/model combination +unsupported output format +invalid config override +``` + +Do not overfit tests to complete traceback text. + +--- + +# 50. CLI resume workflow + +Add an end-to-end test: + +```text +train briefly +save checkpoint +invoke train with checkpoint.resume +verify epoch/global_step continuation +``` + +Use tiny deterministic data. + +--- + +# 51. CLI fine-tuning workflow + +Add an end-to-end test: + +```text +create/save tiny pretrained multiclass checkpoint +invoke fine-tuning workflow +verify binary checkpoint is created +verify output classifier dimension +``` + +This protects one of the core project workflows. + +--- + +# 52. Complete end-to-end CPU smoke test + +Add one complete new-stack workflow: + +```text +generate tiny ROOT data + -> +prepare/cache + -> +train tiny multiclass ROOT-GNN + -> +save checkpoint + -> +load pretrained backbone + -> +fine-tune binary model + -> +evaluate + -> +predict +``` + +The goal is not scientific convergence. + +Verify: + +```text +every stage completes +artifacts from one stage can be consumed by the next +sample identities remain aligned +output schemas/shapes are correct +``` + +Keep runtime reasonable. + +--- + +# 53. Cross-layer round-trip test + +Add a high-value regression test: + +```text +ROOT event + -> +EventSample + -> +features + -> +GraphSample + -> +cache + -> +GraphBatch +``` + +Compare scientifically relevant values across boundaries. + +This should catch schema drift between layers. + +--- + +# 54. Real fixture tests + +Keep optional reduced-real-ROOT tests. + +Add only checks that provide value not already covered synthetically. + +Possible checks: + +```text +sample IDs are unique +all logits are finite +all feature widths are correct +all graphs satisfy N*(N-1) edge count +batched inference matches individual inference +``` + +Continue skipping when the real fixture is unavailable. + +--- + +# 55. Distributed metric parity + +Current distributed coverage compares a two-rank CPU update with single-process training. + +Add single-process vs 2-rank CPU evaluation comparison for: + +```text +loss +accuracy +ROC AUC +sample count +``` + +ROC AUC must use globally gathered predictions. + +--- + +# 56. Distributed uneven dataset sizes + +Test a dataset length not divisible by world size. + +Verify: + +```text +no unintended duplicate evaluation samples +global sample count is correct +metrics match single-process evaluation +``` + +This is a high-priority DDP edge case. + +--- + +# 57. Distributed checkpoint behavior + +Add a two-rank CPU test verifying: + +```text +only rank 0 writes checkpoint +checkpoint contains normalized model keys +fresh single-process model can load the checkpoint +``` + +Do not require CUDA. + +--- + +# 58. Distributed fine-tuning + +Add CPU DDP tests for: + +```text +frozen backbone +unfrozen backbone +``` + +Verify: + +```text +frozen backbone remains unchanged +classifier updates +unfrozen backbone updates +``` + +--- + +# 59. Distributed lifecycle cleanup + +Where practical, verify process-group initialization/finalization works cleanly in isolated spawned tests. + +Avoid tests likely to hang. + +Do not intentionally create deadlocked failure scenarios. + +--- + +# 60. ONNX dynamic-shape coverage + +When ONNX dependencies are installed, test exported models using: + +```text +different node count +different edge count +batch size 1 +batch size > 1 +single-node zero-edge graph if supported +multiclass model +fine-tuned binary model +``` + +Do not validate only the example shape used during export. + +--- + +# 61. ONNX checkpoint export integration + +Add: + +```text +checkpoint + -> +export + -> +ONNX Runtime +``` + +Compare ONNX logits with a fresh PyTorch model loaded from the same checkpoint. + +--- + +# 62. ONNX failure tests + +Test: + +```text +unsupported model +invalid checkpoint +invalid export input +unsupported dynamic case +missing optional dependency +``` + +according to the actual export contract. + +--- + +# 63. Optional dependency boundaries + +Verify: + +```text +base package can import without ONNX +ONNX export fails with actionable message when dependency is absent +DGL-dependent tests/workflows have the intended skip/fail behavior +``` + +Do not dynamically uninstall installed packages if a subprocess or isolated optional-dependency job is cleaner. + +--- + +# 64. Package-resource tests + +Verify installed runtime resources such as Hydra configs can be found through the same resource mechanism used by the built package. + +Do not rely only on source-tree paths. + +--- + +# 65. Wheel smoke validation + +If release validation does not already fully cover it, provide automated validation that: + +```text +builds wheel +installs it into a clean environment +runs import +runs gnn4colliders --help +composes representative config +``` + +This does not need to run inside ordinary pytest. + +--- + +# 66. Public API regression test + +Maintain one explicit test of supported public imports. + +Protect intentionally public APIs. + +Do not freeze private/internal symbols. + +--- + +# 67. Property-based testing assessment + +Evaluate whether Hypothesis would materially improve testing of: + +```text +graph topology invariants +phi wrapping +loss permutation invariance +metadata/sample ordering +batching +``` + +Do not add it automatically. + +If deterministic parameterized tests express the properties clearly, prefer no new dependency. + +If added, keep generated cases bounded and reproducible. + +--- + +# 68. Parameterization over duplication + +Use: + +```python +pytest.mark.parametrize +``` + +for repeated dimensions such as: + +```text +node count +batch size +dtype +output size +processing steps +``` + +Avoid dozens of nearly identical tests. + +--- + +# 69. Numerical tolerances + +Audit approximate comparisons. + +Use explicit: + +```text +rtol +atol +``` + +based on expected precision. + +Do not loosen tolerances merely because a test is flaky. + +Document unusually loose tolerances. + +--- + +# 70. NaN/Inf regression checks + +For valid representative inputs, verify finiteness of: + +```text +features +model logits +loss +trainable gradients +task scores +``` + +Do not define arbitrary behavior for intentionally invalid pathological input unless validation is part of the contract. + +--- + +# 71. Input immutability + +Where immutability is intended, test that operations do not unexpectedly mutate: + +```text +EventSample data +feature input arrays +GraphSample source data +GraphBatch metadata +checkpoint dictionaries passed to normalization +PredictionResult passed to writers +``` + +Do not require immutable Trainer state. + +--- + +# 72. Reentrancy tests + +Call important public operations repeatedly. + +Examples: + +```text +model inference +cache load +checkpoint load +predictor.predict +NPZ writer +``` + +Verify one invocation does not leak state into later invocations. + +--- + +# 73. User-facing error quality + +For major validation failures, assert useful error fragments such as: + +```text +schema version +checkpoint incompatible +missing branch +pretrained checkpoint required +optional dependency missing +``` + +Do not assert entire exact error strings unless stability is intentional. + +--- + +# 74. Standardize pytest markers + +Review and register useful markers. + +Potential categories: + +```text +integration +parity +distributed +onnx +gpu +real_data +legacy_env +slow +``` + +Use only markers that materially improve test selection. + +Avoid unnecessary marking of ordinary fast tests. + +--- + +# 75. Preserve the required ROOT-GNN parity gate + +Preserve: + +```bash +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 pytest +``` + +or the repository's final canonical equivalent. + +Required behavior: + +```text +if ROOT-GNN/DGL validation is required and DGL is unavailable: + fail clearly +``` + +rather than silently skipping required tests. + +Add direct coverage of the gate mechanism if practical. + +--- + +# 76. Document test-suite commands + +Provide documented commands for: + +```text +fast unit suite +full CPU suite +required ROOT-GNN parity suite +ONNX suite +distributed suite +real-data suite +GPU suite +``` + +Use the final marker/path conventions actually implemented. + +--- + +# 77. GPU test scaffolding + +GPU availability must remain optional. + +Add opt-in GPU tests for important device-specific behavior where valuable: + +```text +GraphBatch.to(cuda) +ROOT-GNN forward/backward on CUDA +fine-tuning freeze behavior on CUDA +checkpoint CPU/GPU map_location +single-GPU inference +``` + +Use an explicit GPU marker and clean skips when CUDA is unavailable. + +--- + +# 78. CPU-vs-GPU inference parity + +For an opt-in GPU test, compare CPU and CUDA float32 inference using a tiny deterministic model/batch. + +Use justified tolerance. + +Do not require bitwise equality. + +--- + +# 79. GPU checkpoint portability + +Add an opt-in test: + +```text +CUDA model/checkpoint + -> +save + -> +load with map_location=cpu + -> +CPU inference +``` + +Compare output semantics. + +--- + +# 80. Perlmutter smoke testing + +Do not put real Slurm/Perlmutter execution in ordinary pytest. + +If missing, provide an opt-in smoke script/config for: + +```text +single-GPU training +multi-GPU DDP +checkpoint +evaluation +``` + +using tiny workloads. + +Do not submit jobs automatically in CI. + +--- + +# 81. Slurm static validation + +Continue validating shell scripts with: + +```bash +bash -n scripts/slurm/*.sh +``` + +where present. + +Treat this only as syntax validation. + +--- + +# 82. Performance regression strategy + +Do not add wall-clock thresholds to ordinary tests. + +Keep Task 15 benchmarks as the performance-regression mechanism. + +If useful, add only a tiny benchmark smoke test proving benchmark scripts initialize and return metrics. + +--- + +# 83. Memory-safety regression tests + +Avoid exact memory thresholds. + +Instead test structural conditions that prevent known memory leaks: + +```text +epoch history does not retain autograd graphs +evaluation outputs are detached +PredictionResult tensors are on CPU +inference does not keep GPU graph references unnecessarily +``` + +--- + +# 84. Autograd retention + +Verify evaluation/inference outputs satisfy where appropriate: + +```text +requires_grad == False +grad_fn is None +``` + +This protects against subtle memory regressions. + +--- + +# 85. Improve deterministic fixtures + +Consolidate reusable fixtures for: + +```text +tiny EventSample +single-node GraphSample +small multi-node GraphSample +GraphBatch +binary task +multiclass task +tiny EdgeNetwork +FineTunedEdgeNetwork +temporary ROOT file +temporary checkpoint +``` + +Do not create a giant unrelated `conftest.py`. + +Keep fixtures near the tests that use them when practical. + +--- + +# 86. Avoid RNG fragility + +If random tensors are used: + +```text +seed explicitly +``` + +Prefer hand-constructed tensors for scientific contracts. + +Tests must not depend on execution order or previous RNG consumption. + +--- + +# 87. Test independence + +Tests should not depend on: + +```text +execution order +pre-existing outputs +repository working directory +developer caches +personal environment variables +``` + +Use: + +```text +tmp_path +monkeypatch +``` + +appropriately. + +--- + +# 88. Filesystem robustness + +Add tests that run relevant APIs from outside the repository root. + +Especially protect: + +```text +Hydra config/resource discovery +checkpoint path handling +prediction output paths +``` + +This is important for installed-package behavior. + +--- + +# 89. Path edge cases + +Where practical, include at least one temporary directory/file path containing spaces or nested directories. + +This catches unsafe path assumptions. + +Do not over-expand into OS-specific testing. + +--- + +# 90. Test organization + +Keep tests aligned with architecture: + +```text +tests/unit// +tests/integration/ +tests/parity/ +``` + +Do not create a giant catch-all test file. + +Name tests by behavior. + +--- + +# 91. Avoid implementation-trivia tests + +Do not add tests that primarily assert: + +```text +private helper was called exactly once +exact private method name +private class layout +internal dictionary iteration +``` + +unless those details are part of a public contract. + +Prefer outputs, state transitions, invariants, and errors. + +--- + +# 92. Keep runtime practical + +Comprehensive coverage should still be usable during development. + +Aim for: + +```text +unit: + fast and isolated + +integration: + tiny deterministic workflows + +optional: + GPU + real-data + legacy-environment + expensive export/distributed cases +``` + +Do not make every developer run a production training workload. + +--- + +# 93. CI layering + +If CI is present, arrange test execution roughly as: + +```text +lint + -> +unit + -> +integration + -> +parity + -> +distributed / ONNX optional jobs + -> +package smoke validation +``` + +Parallelize where sensible. + +Do not require Perlmutter or external production datasets. + +--- + +# 94. Testing documentation + +Add or update: + +```text +docs/testing.md +``` + +Document: + +```text +test layers +pytest markers +required ROOT-GNN parity gate +optional dependencies +real ROOT fixture +GPU tests +distributed tests +ONNX tests +package/release smoke tests +``` + +Make it clear what contributors should run before committing. + +--- + +# 95. README developer testing section + +Keep README concise. + +Include the primary commands and link to detailed testing docs. + +For example: + +```bash +uv run pytest +uv run pytest tests/unit +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest tests/parity +``` + +Use the actual final commands. + +--- + +# 96. AGENTS.md + +Add concise durable testing rules if absent: + +```text +Test public contracts rather than private implementation details. +Add parity coverage for scientific behavior changes. +Every bug fix should include a regression test. +Use deterministic fixtures. +Do not add timing thresholds to ordinary pytest. +Keep GPU/Slurm/real-data dependencies optional. +Do not silently skip required ROOT-GNN parity. +``` + +Do not duplicate the entire testing guide. + +--- + +# 97. Coverage report + +If `pytest-cov` is already available, run coverage to identify obvious untested modules. + +Use coverage only as one input to the audit. + +Do not add meaningless tests solely to increase a percentage. + +Do not establish an arbitrary new coverage threshold unless project policy requires one. + +--- + +# 98. Bug-discovery rule + +If new tests expose a genuine production bug: + +1. retain/add the regression test +2. make the smallest production fix necessary +3. document the bug in the completion report +4. run relevant parity tests +5. avoid unrelated refactoring + +Production bug fixes discovered by the new test suite are in scope. + +--- + +# 99. Preserve intentional scientific behavior + +If a new test exposes behavior that appears unusual but is already: + +```text +parity-tested +documented +scientifically intentional +``` + +do not change it merely because a conventional implementation would differ. + +Examples include: + +```text +legacy-compatible edge ordering +weighted loss reduction semantics +negative-weight handling +``` + +Protect established contracts. + +--- + +# 100. Validation + +Run the ordinary full suite: + +```bash +uv run pytest +``` + +Run focused layers: + +```bash +uv run pytest tests/unit -v +uv run pytest tests/integration -v +uv run pytest tests/parity -v +``` + +Run required ROOT-GNN parity: + +```bash +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest +``` + +or the final established command. + +Run distributed tests using the final marker/path: + +```bash +uv run pytest -m distributed -v +``` + +Run ONNX tests when dependencies are installed: + +```bash +uv run pytest -m onnx -v +``` + +Run optional GPU tests when CUDA is available: + +```bash +uv run pytest -m gpu -v +``` + +Run real-data tests when the fixture is available: + +```bash +GNN4COLLIDERS_ROOT_FIXTURE=/path/to/fixture.root \ + uv run pytest -m real_data -v +``` + +Use actual final marker names. + +Run lint/format: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* new tests protect meaningful contracts +* duplicate low-value tests were avoided +* fixtures are deterministic +* ordinary tests remain CPU-friendly +* GPU/ONNX/real-data/legacy dependencies remain isolated +* scientific semantics were not changed without justification +* every discovered production bug has a regression test +* no large production data/model artifacts were added +* no timing thresholds were added to ordinary pytest +* no unrelated refactoring is included + +--- + +# Completion criteria + +Task 20 is complete when: + +1. Existing test coverage has been audited before adding tests. +2. Important cross-layer invariants have explicit tests. +3. Feature boundary/metamorphic coverage is improved. +4. Graph topology invariants are tested across multiple sizes. +5. Edge-feature symmetry/wrapping invariants are covered. +6. ROOT I/O failure cases are covered. +7. Multi-file ordering/sample identity is covered. +8. Split disjointness/config validation is covered. +9. Cache corruption/invalidation is covered. +10. GraphBatch metadata/order invariants are covered. +11. Model gradient flow is tested. +12. Frozen/unfrozen transfer gradients are tested. +13. Model state-dict round-trip equivalence is covered. +14. Weighted loss reference cases are expanded. +15. Loss/metric permutation invariance is covered where appropriate. +16. Degenerate metric behavior is covered. +17. Trainer epoch/global-step transitions are tested. +18. Resume-vs-continuous behavior is strengthened. +19. Checkpoint incompatibility/corruption is tested. +20. Predictor batch-size invariance is tested. +21. Prediction metadata/sample alignment is tested. +22. Modern NPZ schema/error handling is better covered. +23. Legacy compatibility rejection paths are covered. +24. Canonical Hydra configs are composition-tested. +25. Invalid cross-config combinations are tested. +26. CLI prepare/train/evaluate/predict workflows have smoke coverage. +27. CLI resume is tested. +28. CLI fine-tuning is tested. +29. At least one complete tiny CPU end-to-end workflow is tested. +30. Distributed global metric parity is tested. +31. Uneven distributed dataset behavior is tested. +32. DDP checkpoint rank behavior is tested. +33. ONNX dynamic-shape coverage is improved when dependencies exist. +34. Optional dependency boundaries are tested. +35. Package-resource/config discovery is protected. +36. Public API imports remain explicitly tested. +37. GPU opt-in coverage exists for important device-specific contracts without making GPU mandatory. +38. Test markers are clear and registered. +39. Testing documentation is updated. +40. Full ordinary suite passes. +41. Required ROOT-GNN parity gate passes in the canonical ROOT-GNN environment. +42. Lint/format checks pass. + +--- + +# Completion report + +Report: + +1. original test file/case count +2. final test file/case count +3. test-gap audit summary +4. feature tests added +5. graph tests added +6. data/ROOT tests added +7. metadata/sample identity tests added +8. split/cache tests added +9. batching tests added +10. model/gradient tests added +11. transfer-learning tests added +12. task/loss/metric tests added +13. trainer tests added +14. checkpoint tests added +15. inference/output tests added +16. compatibility tests added +17. config/factory tests added +18. CLI tests added +19. end-to-end integration tests added +20. distributed tests added +21. ONNX tests added +22. GPU opt-in tests added +23. package/resource tests added +24. real-data test changes +25. new pytest markers +26. fixture/conftest improvements +27. production bugs discovered +28. production fixes made +29. intentionally untested areas and why +30. ordinary suite results +31. required parity gate results +32. distributed test results +33. ONNX test results +34. GPU test results if available +35. real-fixture results if available +36. lint/format results +37. remaining high-value test gaps + +After validation succeeds, create one Git commit containing only the testing improvements and any minimal regression fixes uncovered by them. + +Use: + +```text +test: expand comprehensive regression coverage +``` + +Before committing, inspect the final diff and ensure no unrelated implementation changes or large test artifacts are included. + diff --git a/tasks/task21.md b/tasks/task21.md new file mode 100644 index 0000000000000000000000000000000000000000..316f63e804b5be81019a7d1a94523285fc87a291 --- /dev/null +++ b/tasks/task21.md @@ -0,0 +1,2254 @@ +# Task 21: Full Legacy vs Rewrite End-to-End Validation + +Perform a complete, staged end-to-end validation of the rewritten GNN4Colliders ROOT-GNN stack against the frozen legacy implementation. + +The goal is to prove that the new implementation is scientifically and operationally equivalent where parity is required, while clearly documenting any intentional redesigns or unavoidable differences. + +This task should compare the two pipelines progressively: + +```text +ROOT input + ↓ +event selection / labels / metadata + ↓ +node features + ↓ +graph topology + ↓ +edge features + ↓ +batching + ↓ +fixed-weight model forward + ↓ +loss / metrics + ↓ +one training step + ↓ +short deterministic training + ↓ +checkpoint reload + ↓ +full inference + ↓ +final prediction artifacts +``` + +Do not start with a long production training comparison. + +If a final training result differs, the validation system must make it possible to determine exactly where the divergence first appears. + +Before making changes, read: + +```text +AGENTS.md +README.md +docs/architecture.md +docs/migration.md +docs/testing.md +docs/compatibility.md +``` + +where present. + +Then inspect: + +```text +src/gnn4colliders/ +tests/parity/ +tests/integration/ +legacy/ +configs/ +scripts/ +``` + +Pay particular attention to the active legacy ROOT-GNN implementation and the canonical new ROOT-GNN configs. + +Treat the frozen executable legacy behavior as the parity reference where scientific compatibility is intended. + +Do not modify the legacy implementation to make the comparison pass. + +--- + +# Goal + +Build a reproducible validation campaign that answers: + +```text +Does the same ROOT event produce the same scientific inputs? + +Does the same graph produce the same model output? + +Does the same batch produce the same loss? + +Does one optimizer step produce the same parameter update? + +Does short deterministic training follow the same trajectory? + +Can legacy checkpoints be consumed by the rewrite correctly? + +Do both implementations produce equivalent final event-level predictions? + +Are any remaining differences expected, intentional, and documented? +``` + +The final result should be a machine-readable and human-readable parity report. + +--- + +# 1. Add a dedicated validation area + +Create a focused area for end-to-end validation. + +Prefer something like: + +```text +validation/ + README.md + compare_preprocessing.py + compare_batches.py + compare_forward.py + compare_training_step.py + compare_short_training.py + compare_inference.py + run_full_validation.py +``` + +or a similarly compact structure. + +Alternatively place reusable implementation under: + +```text +src/gnn4colliders/validation/ +``` + +and keep executable drivers under: + +```text +validation/ +``` + +Do not mix this campaign into ordinary production code. + +--- + +# 2. Reuse existing parity tests + +Do not duplicate the existing low-level parity suite unnecessarily. + +Existing tests already protect things such as: + +```text +node feature schema +graph topology +edge ordering +edge features +model forward +transfer-learning forward +legacy checkpoint prefix handling +``` + +Task 21 should build on these and add cross-layer/end-to-end validation. + +Use current parity helpers where appropriate. + +--- + +# 3. Freeze a canonical validation dataset + +Choose one small but representative ROOT validation sample. + +Prefer: + +```text +hundreds to a few thousand events +``` + +rather than a huge production dataset. + +The sample should exercise, where present: + +```text +different object multiplicities +multiple classes +multiple folds +global features +positive event weights +zero event weights +negative event weights +phi values near wrapping boundaries +single-node events if valid +high-node-count events representative of production +``` + +Do not modify the input sample between legacy and rewrite runs. + +--- + +# 4. Record dataset provenance + +Create a validation manifest containing: + +```text +input file(s) +tree name +event count +file size +cryptographic checksum +selection +branch configuration +``` + +Prefer SHA-256 for file identity. + +For example: + +```json +{ + "files": [ + { + "path": "...", + "sha256": "..." + } + ], + "tree": "...", + "events": 1024 +} +``` + +Do not make an absolute personal filesystem path part of the scientific identity. + +--- + +# 5. Support external validation data + +Do not require committing a production ROOT sample into Git. + +Support an environment variable or explicit path such as: + +```text +GNN4COLLIDERS_E2E_FIXTURE +``` + +or a similarly clear name. + +If no external sample is available, the validation tooling should be able to use a small generated/synthetic fixture for smoke testing. + +--- + +# 6. Freeze equivalent legacy and new configs + +Select one active legacy configuration and construct the scientifically equivalent new Hydra configuration. + +Create an explicit mapping document. + +For example: + +```text +legacy field new field +--------------------------------------------------------- +legacy node feature config data/features config +tracking[:,0] metadata.fold +tracking[:,1] metadata.weight +n_proc_steps model.n_proc_steps +legacy batch size data.batch_size +legacy fold selection data.split.* +``` + +Use the actual final names. + +Do not compare vaguely similar configurations. + +--- + +# 7. Config parity manifest + +Store the exact validation configuration pair. + +For example: + +```text +validation/manifests/ + legacy_config.yaml + rewrite_config.yaml + config_mapping.md +``` + +If the legacy config contains machine-specific paths, create a portable validation version that changes only environmental paths, not scientific settings. + +--- + +# 8. Separate scientific and environmental configuration + +When normalizing configs, classify differences as: + +```text +scientific +environmental +implementation-only +``` + +Examples of environmental differences: + +```text +input path +output path +worker count +device +``` + +Examples of scientific differences: + +```text +feature scales +model dimensions +fold assignment +loss behavior +``` + +Scientific differences must not be silently ignored. + +--- + +# 9. Build normalized event identities + +The legacy and new pipelines may represent event identity differently. + +Create a normalized comparison key based on stable source identity such as: + +```text +source file +tree +entry index +``` + +and map it to the new: + +```text +sample_id +``` + +Do not rely on batch position alone. + +--- + +# 10. Preprocessing dump format + +Create a small machine-readable dump format for comparing preprocessing. + +Prefer: + +```text +NPZ +JSON + NPZ +``` + +rather than pickled Python objects. + +Each event should be comparable by stable ID. + +Include: + +```text +sample_id / normalized legacy identity +label +fold +weight +global features +node feature matrix +edge src +edge dst +edge feature matrix +``` + +Handle variable-size arrays cleanly. + +--- + +# 11. Compare event ordering + +First compare: + +```text +number of selected events +event identities +event order +``` + +Report: + +```text +missing in legacy +missing in rewrite +duplicates +order differences +``` + +Do not proceed as if feature parity is meaningful if the event sets already differ. + +--- + +# 12. Compare labels + +For each event compare: + +```text +label +``` + +Use exact equality. + +Report the exact event identities of any mismatch. + +--- + +# 13. Compare fold semantics + +Compare legacy fold assignment against: + +```text +metadata.fold +``` + +in the rewrite. + +Require exact equality. + +Report: + +```text +number of mismatches +first mismatching events +legacy value +rewrite value +``` + +--- + +# 14. Compare event weights + +Compare legacy event weight against: + +```text +metadata.weight +``` + +Use tight floating-point tolerance appropriate to source dtype. + +Pay special attention to: + +```text +negative weights +zero weights +very large/small weights +``` + +Do not silently take absolute values during validation. + +--- + +# 15. Compare global features + +Compare: + +```text +shape +ordering +dtype where relevant +values +``` + +event-by-event. + +Report maximum absolute and relative differences. + +--- + +# 16. Compare node count + +For every event require identical node count. + +This should fail before comparing node feature matrices if counts differ. + +Include event identity in the report. + +--- + +# 17. Compare node feature schema + +Require the rewrite's active node schema to correspond exactly to the legacy active representation: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +Verify: + +```text +column count +column order +``` + +No column permutation should be accepted merely because the same values are present. + +--- + +# 18. Compare node feature values + +Compare event-by-event and node-by-node. + +Report: + +```text +max absolute error +max relative error +event ID +node index +feature name/index +legacy value +rewrite value +``` + +Use tight float tolerances. + +--- + +# 19. Compare graph topology + +Require exact equality of: + +```text +edge count +src indices +dst indices +edge ordering +``` + +The graph is expected to be: + +```text +directed +fully connected +no self-loops +``` + +with: + +```text +E = N * (N - 1) +``` + +for valid nonempty graphs. + +Do not treat edge ordering differences as acceptable unless explicitly proven irrelevant and approved. + +--- + +# 20. Compare edge features + +Compare: + +```text +deta +dphi +dR +``` + +edge-by-edge in the validated ordering. + +Report maximum differences and exact mismatch location. + +Pay particular attention to phi branch-cut cases. + +--- + +# 21. Empty and single-node edge cases + +If present in the dataset or supported by the pipeline, explicitly validate: + +```text +N = 0 +N = 1 +``` + +according to the established contracts. + +For: + +```text +N = 1 +``` + +expect: + +```text +E = 0 +``` + +Do not invent semantics for unsupported zero-node events. + +--- + +# 22. Compare complete preprocessing content + +For each event define a preprocessing parity status. + +Conceptually: + +```text +event identity: exact +label: exact +fold: exact +weight: tolerance +globals: tolerance +nodes: tolerance +topology: exact +edges: tolerance +``` + +Produce both: + +```text +per-event details +aggregate summary +``` + +--- + +# 23. Preprocessing parity gate + +The fixed-weight model comparison should only proceed if preprocessing parity passes, unless an explicit diagnostic override is used. + +Do not let downstream differences obscure upstream failures. + +--- + +# 24. Compare batching behavior + +Construct equivalent ordered samples in both systems. + +Compare: + +```text +which events appear in each batch +batch size +number of graphs +total node count +total edge count +labels +weights +global features +``` + +If exact batching order is expected, require it. + +--- + +# 25. Handle intentional batching differences + +If the rewrite intentionally changes internal batching/prebatch implementation, compare by stable sample identity. + +Document: + +```text +legacy batch membership +rewrite batch membership +``` + +and distinguish: + +```text +scientifically irrelevant reordering +scientifically relevant change +``` + +Do not classify reordering as harmless without confirming no downstream stochastic behavior depends on it. + +--- + +# 26. Padding comparison + +If the selected legacy workflow uses padding, characterize and compare: + +```text +padding mode +padded node count +padded edge count +mask/selection behavior +``` + +Do not force obsolete padding into the new system if the canonical compatibility workflow does not require it. + +Document any intentional difference. + +--- + +# 27. Establish fixed model weights + +Do not compare independently initialized models. + +Create one deterministic model state that can be represented in both implementations. + +Preferred approaches: + +```text +new model state -> explicit legacy mapping +``` + +or: + +```text +legacy model state -> new compatibility mapping +``` + +Use the existing compatibility machinery rather than ad hoc key rewriting. + +--- + +# 28. Prefer a real historical checkpoint where useful + +In addition to synthetic deterministic weights, run at least one comparison using a real supported historical ROOT-GNN checkpoint if available. + +This validates actual compatibility. + +Do not commit a large production checkpoint to Git. + +Allow it to be supplied externally. + +--- + +# 29. Fixed-weight forward parity + +For identical preprocessed events and model weights, compare: + +```text +legacy logits +rewrite logits +``` + +event-by-event. + +Use: + +```text +model.eval() +no gradient +``` + +for both. + +--- + +# 30. Forward parity report + +Report: + +```text +shape equality +max absolute error +max relative error +mean absolute error +event with max error +class/output index +``` + +Also report count of predictions exceeding the configured tolerance. + +--- + +# 31. Compare intermediate representations + +If final logits disagree beyond tolerance, provide an optional diagnostic mode that compares intermediate stages for a tiny batch. + +Potential stages: + +```text +node encoder output +edge encoder output +global encoder output +edge update per processing step +node update per processing step +global representation +pooled graph representation +classifier input +``` + +Implement this only as a diagnostic helper. + +Do not permanently change public model APIs solely for validation. + +--- + +# 32. Compare multiclass model + +Run fixed-weight parity for the active multiclass pretraining architecture. + +Do not validate only the fine-tuned model. + +--- + +# 33. Compare fine-tuned binary model + +Run fixed-weight parity for the transfer-learning model. + +Use: + +```text +pretrained backbone +replacement binary classifier +``` + +Verify both: + +```text +frozen-backbone configuration +trainable-backbone configuration +``` + +where relevant to forward semantics. + +--- + +# 34. Compare task loss + +Using identical logits, targets, and event weights, compare legacy and rewrite loss. + +This should explicitly validate the characterized weighted loss reduction. + +Report: + +```text +legacy loss +rewrite loss +absolute difference +relative difference +``` + +Do not substitute a generic BCE/CE calculation for the actual legacy reference. + +--- + +# 35. Compare metrics + +On the same fixed predictions compare: + +```text +accuracy +ROC AUC +other active metrics +``` + +using identical weights and event sets. + +Make sure metrics are computed over the same full split. + +Do not average per-batch AUC values. + +--- + +# 36. One-step training parity + +Set up a deterministic one-step training comparison. + +Requirements: + +```text +same model weights +same batch +same train/eval mode +same optimizer type +same optimizer hyperparameters +same loss +same random state +``` + +Prefer CPU first. + +Disable: + +```text +DDP +AMP +torch.compile +``` + +for the baseline parity test. + +--- + +# 37. Compare pre-step values + +Before backward compare: + +```text +logits +loss +``` + +These should already satisfy the fixed-forward/loss gates. + +If they do not, do not interpret optimizer differences. + +--- + +# 38. Compare gradients + +After backward compare trainable parameter gradients. + +For each mapped parameter report: + +```text +shape +max absolute difference +max relative difference +gradient norm legacy +gradient norm rewrite +``` + +Identify the first/highest mismatch. + +--- + +# 39. Compare parameter updates + +After one optimizer step compare resulting model parameters. + +Report: + +```text +parameter name mapping +max absolute difference +max relative difference +parameter norm +update norm +``` + +This is a critical parity gate. + +--- + +# 40. Compare optimizer state + +Where equivalent optimizer implementations are used, compare meaningful optimizer state. + +Examples: + +```text +step count +momentum +Adam exp_avg +Adam exp_avg_sq +learning rate +``` + +Do not require byte-for-byte serialization identity. + +Compare semantics. + +--- + +# 41. Dropout and RNG + +If dropout is active during training parity, explicitly synchronize RNG behavior if possible. + +If legacy and rewrite cannot reproduce identical dropout masks because of implementation ordering, run an additional deterministic configuration with: + +```text +dropout = 0 +``` + +for strict mathematical parity. + +Then separately characterize stochastic training equivalence. + +Do not hide RNG-related differences. + +--- + +# 42. Short deterministic training comparison + +After one-step parity, run a short training campaign. + +Use something like: + +```text +3–5 epochs +CPU +num_workers = 0 +fixed batch order +fixed seed +no DDP +float32 +no AMP +no torch.compile +``` + +Keep the dataset small enough for practical validation. + +--- + +# 43. Compare epoch-level training trajectory + +After each epoch record: + +```text +training loss +validation loss +accuracy +ROC AUC +learning rate +global step +parameter norm +``` + +and any active early-stopping state. + +Produce one row per epoch per implementation. + +--- + +# 44. Compare short-training model states + +At the end of each epoch, optionally compare model parameters. + +Use this to identify when training trajectories diverge. + +Do not require bitwise equality if operation ordering differs. + +Use justified tolerances. + +--- + +# 45. Define strict vs scientific parity + +Distinguish two validation levels. + +## Strict deterministic parity + +Used for: + +```text +preprocessing +topology +fixed-weight forward +loss +one-step gradient/update +``` + +Expect very tight agreement. + +## Scientific training parity + +Used for: + +```text +multi-epoch stochastic training +GPU execution +DDP execution +``` + +Compare trends/distributions/metrics with justified tolerance. + +Do not mix these two standards. + +--- + +# 46. Early stopping parity + +If the short run reaches early-stopping logic, compare: + +```text +best metric +patience counter +stop decision +best epoch +``` + +If the validation run is too short to trigger it naturally, a focused deterministic check may be used. + +--- + +# 47. Scheduler parity + +If the active legacy workflow uses a scheduler, compare: + +```text +learning rate per step/epoch +scheduler stepping timing +``` + +Ensure the new trainer steps it at the same logical point. + +--- + +# 48. Checkpoint save/reload comparison + +For each implementation: + +```text +train + -> +save checkpoint + -> +fresh process/object + -> +load + -> +evaluate same split +``` + +Verify evaluation logits before and after reload match within that implementation. + +--- + +# 49. Legacy checkpoint -> rewrite comparison + +For a supported legacy checkpoint: + +```text +legacy checkpoint + -> +legacy model + -> +reference logits +``` + +and: + +```text +same legacy checkpoint + -> +rewrite compatibility loader + -> +new model + -> +rewrite logits +``` + +Compare event-by-event. + +This is a required compatibility gate. + +--- + +# 50. Pretrained backbone compatibility + +Validate the real transfer workflow: + +```text +legacy multiclass pretrained checkpoint + -> +new compatibility loader + -> +new ROOT-GNN backbone + -> +replacement binary classifier +``` + +Verify that transferred backbone parameters match the intended source tensors exactly/tightly. + +--- + +# 51. Resume comparison + +Do not require full legacy optimizer resume compatibility unless it is documented as supported. + +For the rewrite itself, verify normal Task 11 resume behavior remains correct. + +For legacy comparison, clearly classify: + +```text +weight compatibility +backbone compatibility +full training-resume compatibility +``` + +separately. + +--- + +# 52. Full inference comparison + +Run the same validation/test event set through both implementations. + +Use: + +```text +fixed model/checkpoint +eval mode +deterministic preprocessing +``` + +Capture per-event: + +```text +event ID +label +weight +logits +scores +prediction +``` + +and any supported selection/output fields. + +--- + +# 53. Compare by event identity + +Never compare final inference artifacts only by row position. + +Join using normalized event identity / `sample_id`. + +Detect: + +```text +missing events +duplicate events +different ordering +``` + +before numerical comparison. + +--- + +# 54. Compare logits + +Compare final event-level raw logits. + +Report: + +```text +max abs error +max rel error +mean abs error +quantiles of absolute error +worst event IDs +``` + +Use appropriate output dimension handling for binary vs multiclass. + +--- + +# 55. Compare scores + +Compare the task-level scores produced from logits. + +Examples: + +```text +binary sigmoid score +multiclass class probabilities/scores +``` + +Use the actual Task 9 semantics. + +Do not invent a legacy score transformation. + +--- + +# 56. Compare predicted classes + +Report: + +```text +number of prediction disagreements +fraction of disagreements +event IDs +legacy score/logit +rewrite score/logit +``` + +If disagreements occur only at threshold ties/numerical boundaries, flag them separately. + +Do not hide them. + +--- + +# 57. Score correlation + +For production-style validation compute event-by-event score correlation. + +Useful summaries include: + +```text +Pearson correlation +Spearman correlation +``` + +where appropriate. + +These are supplementary diagnostics, not substitutes for strict parity when strict parity is expected. + +--- + +# 58. Compare score distributions + +For a larger representative sample, compare score distributions by: + +```text +class +split +``` + +Possible diagnostics: + +```text +histogram overlays +difference histogram +quantile comparison +``` + +Do not require plotting libraries in core production code. + +Validation scripts may use existing lightweight plotting dependencies if already available. + +--- + +# 59. Physics-level metrics + +Compare final relevant metrics such as: + +```text +ROC AUC +accuracy +signal efficiency at selected working points if already part of workflow +background rejection if already part of workflow +``` + +Do not invent new physics metrics solely for this task. + +--- + +# 60. Weighted metrics + +Ensure the comparison uses the exact established event-weight semantics. + +Report both: + +```text +unweighted diagnostics +official weighted metric +``` + +only if both are useful. + +Do not replace the official weighted metric. + +--- + +# 61. Negative-weight diagnostics + +If negative event weights exist, explicitly report: + +```text +number of negative-weight events +sum of positive weights +sum of negative weights +``` + +and confirm both implementations receive identical values. + +This is important for interpreting metric differences. + +--- + +# 62. Compare final NPZ outputs + +If the legacy pipeline produces NPZ and the rewrite produces modern NPZ: + +```text +normalize both into a common comparison representation +``` + +Compare: + +```text +event identity +label +weight +logits/scores +``` + +Do not require identical field names where the new format intentionally uses named metadata. + +--- + +# 63. Compare ROOT outputs + +If ROOT score writing is part of the active legacy workflow, compare: + +```text +number of output entries +event alignment +score branches +selection_pass +class-score values +``` + +Normalize branch naming where the new format intentionally differs. + +Do not require byte-identical ROOT files. + +--- + +# 64. Validation report schema + +Create a structured report. + +Prefer JSON. + +For example: + +```json +{ + "dataset": {}, + "configs": {}, + "preprocessing": {}, + "batching": {}, + "forward": {}, + "loss": {}, + "training_step": {}, + "short_training": {}, + "checkpoint": {}, + "inference": {}, + "overall_status": "pass" +} +``` + +Keep the schema stable enough for future reruns. + +--- + +# 65. Report numerical mismatch locations + +Every numerical comparison should report not just a boolean. + +At minimum include: + +```text +max absolute error +max relative error +location of max error +count outside tolerance +total elements compared +``` + +For event data, include: + +```text +sample/event ID +``` + +--- + +# 66. Tolerance configuration + +Centralize tolerances. + +For example: + +```text +validation/tolerances.yaml +``` + +or a dataclass/config. + +Potential categories: + +```text +feature_atol +feature_rtol +edge_atol +edge_rtol +logit_atol +logit_rtol +gradient_atol +gradient_rtol +parameter_atol +parameter_rtol +metric_atol +``` + +Do not scatter unexplained tolerances through scripts. + +--- + +# 67. Use tight defaults + +Start with tight float32 tolerances for deterministic CPU comparisons. + +Only loosen them when: + +```text +the difference is understood +the reason is documented +the new tolerance still catches meaningful regressions +``` + +Do not make a failing comparison pass by arbitrarily increasing tolerance. + +--- + +# 68. Exact fields remain exact + +Do not use floating-point tolerance for: + +```text +event identity +labels +folds +node count +edge count +edge src/dst +predicted integer class +``` + +These should use exact equality. + +--- + +# 69. Generate a human-readable summary + +In addition to JSON, generate a concise Markdown or text report. + +For example: + +```text +validation/reports/latest.md +``` + +Include a table such as: + +```text +Stage Status Max abs error Notes +---------------------------------------------------- +Event identity PASS exact +Labels PASS exact +Weights PASS 0 +Node features PASS 3.1e-7 +Topology PASS exact +Edge features PASS 5.0e-7 +Fixed forward PASS 8.4e-7 +Loss PASS 1.1e-7 +One-step update PASS 2.0e-6 +Short training PASS ... +Inference PASS ... +``` + +Do not commit machine-specific result dumps by default. + +--- + +# 70. Validation artifact layout + +Prefer output such as: + +```text +validation_output/ + manifest.json + legacy/ + preprocessing.npz + predictions.npz + training_history.json + rewrite/ + preprocessing.npz + predictions.npz + training_history.json + reports/ + comparison.json + comparison.md +``` + +Generated results should normally be gitignored. + +--- + +# 71. Separate extraction from comparison + +Design the tooling so it can: + +```text +run legacy extraction +run rewrite extraction +compare existing outputs +``` + +independently. + +This is important because the legacy and new stacks may require different environments. + +Do not require both implementations to import in one Python process. + +--- + +# 72. Support separate environments + +Assume the legacy runtime may require: + +```text +older Python +older DGL +different dependencies +``` + +The preferred architecture is: + +```text +legacy environment: + write normalized validation artifacts + +new environment: + write normalized validation artifacts + +comparison environment: + compare artifacts +``` + +Do not force the new environment to install the entire historical stack. + +--- + +# 73. Legacy extraction script + +Provide a minimal script that runs inside the legacy environment and writes normalized artifacts. + +It may live under: + +```text +validation/legacy/ +``` + +It may import legacy code because its entire purpose is reference extraction. + +Do not modify legacy modules. + +--- + +# 74. Rewrite extraction script + +Provide a corresponding script using only: + +```text +src/gnn4colliders/ +``` + +No production import from `legacy/`. + +--- + +# 75. Normalized artifact contract + +Both extraction scripts must write the same conceptual schema. + +The comparator should not need to understand DGL graph objects or legacy dataset classes. + +Normalize to: + +```text +NumPy arrays +JSON metadata +``` + +where practical. + +--- + +# 76. Variable-sized event storage + +For variable node/edge counts, choose an unambiguous representation. + +Examples: + +```text +flattened arrays + offsets +one file per event +NPZ object arrays only if safe/practical +``` + +Prefer flattened arrays plus offsets for portability. + +For example: + +```text +node_features_flat +node_offsets +edge_src_flat +edge_dst_flat +edge_features_flat +edge_offsets +``` + +Document the schema. + +--- + +# 77. No pickle-only interchange + +Do not make Python pickle the only legacy/rewrite interchange format. + +The comparison artifacts should remain readable independently of implementation classes. + +--- + +# 78. Full validation driver + +Provide a high-level driver such as: + +```bash +python validation/run_full_validation.py ... +``` + +or equivalent. + +Because environments may differ, it is acceptable for the driver to support stages such as: + +```text +extract-legacy +extract-rewrite +compare +``` + +rather than launching everything in one process. + +--- + +# 79. CLI design + +Keep validation tooling separate from the normal: + +```text +gnn4colliders +``` + +scientific CLI unless there is a strong reason to expose it there. + +This is developer/release validation infrastructure. + +--- + +# 80. Smoke mode + +Provide a fast smoke mode using tiny generated data. + +For example: + +```text +--smoke +``` + +or a small config. + +Smoke mode should exercise: + +```text +preprocessing +fixed forward +loss +one training step +inference +``` + +without requiring production files. + +--- + +# 81. Full mode + +Provide a more representative validation configuration for manual/release use. + +This may use: + +```text +GNN4COLLIDERS_E2E_FIXTURE +``` + +or equivalent. + +Do not run full validation automatically on every developer test invocation if it is expensive. + +--- + +# 82. Add selected automated tests + +Convert the deterministic, cheap parts of Task 21 into normal automated tests where practical. + +Good candidates: + +```text +normalized artifact comparator +fixed-weight end-to-end tiny fixture +one-step training parity +legacy checkpoint -> rewrite inference +``` + +Do not force expensive dual-environment validation into ordinary CI. + +--- + +# 83. Existing parity gate integration + +The existing required parity gate is: + +```bash +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 pytest +``` + +Preserve it. + +Task 21 may add a separate explicit full-validation command. + +Do not make normal pytest depend on external legacy environments or production ROOT files. + +--- + +# 84. Production-style Perlmutter validation + +After CPU deterministic parity is established, provide a documented manual validation procedure for Perlmutter. + +Compare: + +```text +legacy single GPU +rewrite single GPU +rewrite multi-GPU DDP +``` + +Use the same scientific config and equivalent effective batch semantics. + +--- + +# 85. GPU comparison standard + +Do not require identical floating-point trajectories between legacy and rewrite GPU training. + +Instead compare: + +```text +initial fixed-weight inference +loss curves +validation metrics +final score distributions +event-level score correlation +``` + +while keeping deterministic CPU parity as the strict implementation gate. + +--- + +# 86. DDP validation + +For the rewrite, compare: + +```text +single-process evaluation +DDP evaluation +``` + +using the same checkpoint. + +Require equivalent global outputs/metrics. + +Then compare those outputs against the legacy single-process reference. + +Do not require a distributed legacy workflow if it adds no validation value. + +--- + +# 87. Batch-size sensitivity + +As part of inference validation, run the rewrite using at least two batch sizes. + +Verify per-event outputs remain unchanged. + +This ensures apparent legacy/rewrite agreement is not accidentally tied to one batching layout. + +--- + +# 88. Prebatch sensitivity + +If the legacy workflow uses prebatching, characterize whether prediction values depend on prebatch grouping. + +They should not in evaluation mode. + +If they do, document the legacy behavior as a significant finding. + +--- + +# 89. Training-order sensitivity + +For strict short-training parity, ensure both implementations consume the same batch/event order. + +Record the order in the validation artifacts. + +If they cannot be made identical, classify the training comparison as scientific rather than strict parity. + +--- + +# 90. Randomness manifest + +Record: + +```text +Python seed +NumPy seed +Torch seed +DataLoader seed +model seed policy +``` + +for both runs. + +Do not assume the historical stack is deterministic. + +Document any RNG source that cannot be synchronized. + +--- + +# 91. Environment manifest + +Record both runtime environments. + +Include: + +```text +Python +PyTorch +DGL +NumPy +Awkward +Uproot +device +CUDA version if relevant +GPU model if relevant +``` + +This is diagnostic metadata, not a parity criterion. + +--- + +# 92. Git/source identity + +Record: + +```text +rewrite git commit +legacy source identity/commit if available +``` + +Do not require Git metadata to run the comparison, but include it when available. + +--- + +# 93. Failure diagnosis + +When validation fails, stop reporting only: + +```text +FAIL +``` + +Report the earliest failing stage. + +For example: + +```text +preprocessing: + PASS + +batching: + PASS + +fixed forward: + FAIL + first mismatch: event 381 + max abs error: ... +``` + +This is the central value of staged validation. + +--- + +# 94. Do not hide known differences + +If the rewrite intentionally changed behavior, encode the difference explicitly in the report. + +Examples: + +```text +legacy tracking array replaced by named metadata +output NPZ field names changed +internal batching implementation changed +checkpoint schema changed +``` + +These may be: + +```text +expected difference +``` + +rather than: + +```text +parity failure +``` + +only if the underlying scientific meaning is equivalent. + +--- + +# 95. Define parity categories + +Use categories such as: + +```text +PASS +PASS_WITH_EXPECTED_DIFFERENCE +FAIL +NOT_APPLICABLE +NOT_RUN +``` + +Avoid ambiguous status. + +--- + +# 96. No silent data dropping + +The comparison code must fail or loudly report when: + +```text +events exist only in one output +duplicate event IDs occur +array lengths differ +metadata is missing +``` + +Do not compare only the intersection and report success. + +--- + +# 97. No tolerance masking of shape errors + +Shape mismatch should fail immediately. + +Do not flatten incompatible tensors and compare overlapping elements. + +--- + +# 98. No production-code redesign for parity tooling + +Task 21 may expose tiny diagnostic hooks if necessary, but do not refactor the entire new stack just to make comparison scripts convenient. + +Prefer adapters and extraction utilities. + +--- + +# 99. Document validation procedure + +Add: + +```text +docs/end_to_end_validation.md +``` + +or similar. + +Document: + +```text +purpose +dataset requirements +legacy environment setup +rewrite environment setup +config mapping +artifact extraction +comparison +tolerances +strict vs scientific parity +Perlmutter validation +interpretation of expected differences +``` + +--- + +# 100. Validation + +Run the existing modern suite first: + +```bash +uv run pytest +``` + +Run required parity: + +```bash +GNN4COLLIDERS_REQUIRE_ROOT_GNN=1 uv run pytest +``` + +Run lint/format: + +```bash +uv run ruff check . +uv run ruff format --check . +``` + +Run the Task 21 smoke validation. + +Then, when the canonical external validation dataset and legacy environment are available: + +```text +extract legacy artifacts +extract rewrite artifacts +compare preprocessing +compare batching +compare fixed forward +compare loss/metrics +compare one-step training +compare short training +compare checkpoint behavior +compare inference outputs +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy production code remains unchanged +* rewrite production code does not import legacy code +* validation artifacts use stable event identity +* exact fields use exact comparison +* floating-point tolerances are centralized +* no large generated validation artifacts are committed +* no production ROOT/checkpoint files are committed +* no parity failures are hidden by intersection-only comparison +* no unrelated scientific changes are included + +--- + +# Completion criteria + +Task 21 is complete when: + +1. A canonical end-to-end validation workflow exists. +2. Legacy and rewrite extraction can run in separate environments. +3. A normalized implementation-independent artifact format exists. +4. Validation dataset provenance is recorded. +5. Legacy and rewrite configs are explicitly mapped. +6. Event identity is normalized across implementations. +7. Event set/order differences are detected. +8. Labels are compared exactly. +9. Fold assignment is compared exactly. +10. Event weights are compared. +11. Global features are compared. +12. Node counts are compared exactly. +13. Node feature schema/order is validated. +14. Node feature values are compared. +15. Edge counts are compared exactly. +16. Edge topology/order is compared exactly. +17. Edge feature values are compared. +18. Batching behavior is characterized. +19. Fixed-weight multiclass forward parity is validated. +20. Fine-tuned binary forward parity is validated. +21. Legacy-vs-rewrite loss parity is validated. +22. Legacy-vs-rewrite metric parity is validated. +23. One-step gradient parity is characterized. +24. One-step parameter-update parity is characterized. +25. Optimizer state differences are characterized where applicable. +26. A short deterministic CPU training comparison exists. +27. Epoch-level training trajectories are reported. +28. Checkpoint save/reload invariance is validated. +29. Supported legacy checkpoint -> rewrite inference parity is validated. +30. Legacy pretrained backbone -> new fine-tuning compatibility is validated. +31. Full inference is compared event-by-event. +32. Raw logits are compared. +33. Task scores are compared. +34. Predicted classes are compared. +35. Final metrics are compared. +36. Output artifacts are compared to the extent scientifically relevant. +37. Strict and scientific parity standards are separated. +38. Numerical tolerances are centralized and documented. +39. Comparison reports identify maximum errors and mismatch locations. +40. Missing/duplicate events cause a clear validation failure. +41. Smoke validation is practical for development. +42. Full validation is documented for release/manual use. +43. Perlmutter single-GPU/DDP comparison procedure is documented. +44. Existing tests and parity gates still pass. +45. Legacy source remains frozen. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. validation directory structure +4. canonical validation dataset or fixture mechanism +5. dataset checksum/provenance +6. legacy configuration used +7. rewrite configuration used +8. config mapping summary +9. legacy environment +10. rewrite environment +11. normalized artifact schema +12. event identity mapping +13. event selection/order comparison +14. label comparison +15. fold comparison +16. weight comparison +17. global-feature comparison +18. node-count comparison +19. node-feature comparison +20. topology comparison +21. edge-feature comparison +22. batching comparison +23. multiclass fixed-forward parity +24. binary fine-tuning fixed-forward parity +25. loss parity +26. metric parity +27. gradient parity +28. one-step parameter-update parity +29. optimizer-state comparison +30. short-training trajectory comparison +31. checkpoint round-trip results +32. legacy checkpoint compatibility results +33. pretrained-backbone transfer results +34. full inference logit comparison +35. score comparison +36. prediction disagreement count +37. score-correlation results where measured +38. physics-level metric comparison +39. ROOT/NPZ output comparison +40. maximum observed numerical differences +41. expected intentional differences +42. unresolved differences +43. smoke validation results +44. Perlmutter validation results if run +45. existing pytest/parity results +46. lint/format results +47. overall parity status + +After validation succeeds, create one Git commit containing only Task 21 validation infrastructure, tests, documentation, and any minimal fixes required by genuine discovered regressions. + +Use: + +```text +test: add full legacy end-to-end validation +``` + +Before committing, inspect the final diff and ensure no generated validation outputs, production datasets, production checkpoints, or unrelated scientific changes are included. + diff --git a/tasks/task2b.md b/tasks/task2b.md new file mode 100644 index 0000000000000000000000000000000000000000..763816f846821e58ceec04e1c095174cacb474ec --- /dev/null +++ b/tasks/task2b.md @@ -0,0 +1,195 @@ +# Task 2B: Fix ROOT-GNN PyTorch/DGL Environment + +The current ROOT-GNN environment setup is incomplete. + +Running: + +```bash +uv sync --extra root-gnn +``` + +currently attempts to install DGL from PyPI and fails on Linux because the required Linux DGL wheel is distributed through DGL's own wheel repository. + +As a result, DGL-dependent parity tests are skipped. + +This is not acceptable for the canonical ROOT-GNN development environment. + +The verified target tuple is: + +```text +Python 3.12 +PyTorch 2.2.2+cu121 +DGL 2.4.0+cu121 +CUDA runtime 12.1 +``` + +Do not continue retrying DGL 2.5.0 artifacts: they returned HTTP 403 in the +target environment even though some official listings advertise them. The +verified Linux CPython 3.12 wheel is DGL 2.4.0+cu121 from the official +`torch-2.2/cu121` repository. + +## Goal + +Make the ROOT-GNN development environment reproducible on Linux/Perlmutter and ensure DGL-dependent parity tests actually run. + +## Requirements + +1. Inspect the current: + + * `pyproject.toml` + * `uv.lock` + * README environment instructions + * AGENTS.md + * Perlmutter/CUDA environment information available in the repository + +2. Determine an explicit compatible tuple for: + + * Python + * PyTorch + * CUDA runtime + * DGL + +3. Prefer Python 3.12 unless compatibility evidence requires otherwise. + +Set `requires-python = ">=3.12,<3.13"` and add `.python-version` containing +`3.12`. + +4. Do not use loose unconstrained dependencies such as: + +```toml +torch >= 2.0 +dgl >= 2.0 +``` + +for the ROOT-GNN stack. + +PyTorch and DGL compatibility must be represented intentionally. + +5. Configure `uv` so DGL is obtained from the appropriate official DGL wheel repository rather than PyPI. + +Use uv's supported custom/flat index/source configuration rather than ad-hoc manual `pip install` commands where practical. + +The DGL `repo.html` page is a flat wheel listing, so `format = "flat"` is +required. Use this exact source configuration: + +```toml +[[tool.uv.index]] +name = "dgl" +url = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html" +format = "flat" +explicit = true + +[tool.uv.sources] +dgl = { index = "dgl" } +``` + +Pin `dgl==2.4.0+cu121` and `torch==2.2.2`, sourcing PyTorch from its official +CUDA 12.1 index. Also constrain `numpy<2`; otherwise uv can select NumPy 2.x, +which produces an ABI warning with this PyTorch build. Add `matplotlib` to the +development dependencies because the legacy parity import path requires it. + +6. Keep DGL architecture-specific: + +```toml +[project.optional-dependencies] +root-gnn = [ + ... +] +``` + +Do not rename this back to a generic `ml` extra. + +7. PyTorch may remain shared if that makes sense for future model families such as ROOT-Transformer. + +8. Document GPU/CUDA assumptions clearly. + +Do not silently make one local machine's module setup a universal package requirement. + +9. After configuration, this command must succeed in the intended ROOT-GNN environment: + +```bash +uv sync --extra root-gnn +``` + +If Perlmutter home-directory extraction fails with `Disk quota exceeded`, use +a scratch-backed uv cache and rerun the command: + +```bash +export UV_CACHE_DIR=/pscratch/sd/$USER/uv-cache-gnn4colliders +uv sync --extra root-gnn +``` + +10. Validate imports explicitly: + +```bash +uv run python -c "import torch; print(torch.__version__)" +uv run python -c "import dgl; print(dgl.__version__)" +``` + +11. Run: + +```bash +uv run pytest tests/parity +``` + +Run CUDA availability checks on a GPU allocation rather than only on a login +node. Expected import output is Torch `2.2.2+cu121`, Torch CUDA `12.1`, DGL +`2.4.0+cu121`, and an available GPU. + +DGL-dependent tests must run rather than skip due to DGL being absent. + +12. If DGL cannot be made compatible with the selected modern PyTorch/CUDA combination, report the exact incompatibility and choose the newest well-supported compatible tuple rather than guessing. + +## Test behavior + +Do not treat skipped DGL tests as successful ROOT-GNN validation. + +The canonical ROOT-GNN validation procedure should fail early if: + +```python +import dgl +``` + +does not work. + +Generic package tests may still allow optional DGL behavior, but ROOT-GNN-specific CI/development validation must require it. + +The current parity characterization suite has one known expectation issue: its +two-node no-self-loop edge-feature test expects four edges, while the legacy +graph constructor produces two. Do not alter legacy code or weaken that test; +report it separately after confirming that DGL imports and the other parity +tests execute. + +## Out of scope + +Do not: + +* rewrite model code +* rewrite preprocessing +* optimize runtime performance +* add Lightning +* add PyG +* add Numba +* modify legacy behavior + +This task is only about establishing a correct PyTorch/DGL development runtime. + +## Completion report + +Report: + +* selected Python version +* selected PyTorch version +* selected DGL version +* selected CUDA target +* uv source/index configuration +* commands run +* import results +* parity test results +* any Perlmutter-specific notes + +Commit as: + +```text +chore: fix ROOT-GNN PyTorch DGL environment +``` diff --git a/tasks/task3.md b/tasks/task3.md new file mode 100644 index 0000000000000000000000000000000000000000..91b22d367bcc8f63fc40583a4fb34a6f079724ce --- /dev/null +++ b/tasks/task3.md @@ -0,0 +1,644 @@ +# Task 3: Characterize Legacy Data and Graph Behavior + +Build a deterministic characterization/parity test suite for the active legacy ROOT-GNN preprocessing path. + +The ROOT-GNN development environment is now expected to include working PyTorch and DGL support. + +Do not implement the new production data pipeline yet. + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect the relevant legacy implementation under: + +```text +legacy/root_gnn_dgl/ +``` + +Focus only on the active data/graph behavior required by the standard ROOT-GNN workflow. + +--- + +# Goal + +Capture the externally observable behavior of the legacy preprocessing stack so that the new implementation can later be rewritten against explicit parity tests. + +The relevant flow is: + +```text +ROOT/Awkward event data + -> +node feature construction + -> +graph topology + -> +edge feature construction + -> +labels / tracking / global features +``` + +The tests must characterize what the legacy implementation **actually does**. + +Do not encode assumptions solely because they appear in documentation. + +If documentation, previous expectations, and observed legacy behavior disagree, treat the running legacy implementation as the behavior to characterize and document the discrepancy. + +--- + +# Environment validation + +Before writing or modifying parity tests, verify that the ROOT-GNN environment is functional. + +Run: + +```bash +uv sync --extra root-gnn + +uv run python -c "import torch; print(torch.__version__)" +uv run python -c "import dgl; print(dgl.__version__)" +``` + +DGL-dependent parity tests must not silently skip because DGL is absent. + +If DGL cannot import, stop and report the environment failure rather than treating skipped tests as successful characterization. + +--- + +# 1. Node feature characterization + +Characterize the active legacy node-feature construction. + +Verify: + +* feature count +* feature ordering +* feature values +* dtype +* shape +* concatenation across object types where applicable +* feature scaling behavior + +The active schema is expected to include: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +But verify this against the legacy implementation rather than assuming it. + +Characterize derived behavior including: + +```text +CALC_E = pt * cosh(eta) +``` + +if that is what the active legacy path actually implements. + +Also verify: + +* constant-valued branches +* `NODE_TYPE` +* scale application +* object ordering when multiple object types are combined + +Use explicit expected arrays where practical. + +--- + +# 2. Graph topology characterization + +Characterize the graph produced by the active legacy graph builder. + +Verify: + +* whether the graph is directed +* whether self-loops are present +* number of nodes +* number of edges +* source indices +* destination indices +* edge ordering if deterministic +* behavior for small graph sizes + +Do not assume a two-node fully connected directed graph has four edges. + +Explicitly characterize the observed behavior. + +For example, if the legacy implementation produces: + +```text +0 -> 1 +1 -> 0 +``` + +for a two-node graph, encode that exact behavior. + +If the topology is fully connected without self-loops, verify the expected edge count: + +```text +N * (N - 1) +``` + +for multiple values of `N`. + +At minimum test: + +```text +N = 1 +N = 2 +N = 3 +``` + +where supported by the legacy implementation. + +For edge ordering: + +* assert exact ordering only if it is deterministic and observable +* otherwise test topology independent of ordering and document the ambiguity + +--- + +# 3. Edge feature characterization + +Characterize the active legacy edge-feature construction. + +Verify that edge features correspond to: + +```text +deta +dphi +dR +``` + +in the actual stored order. + +Test at least one case where phi wrapping matters. + +For example, use objects near: + +```text +phi = +pi +phi = -pi +``` + +so that naive subtraction differs from wrapped angular distance. + +Verify: + +```text +dR = sqrt(deta^2 + dphi^2) +``` + +if that is the actual legacy behavior. + +Important: + +Expected edge-feature arrays must correspond exactly to the graph edges produced by the legacy implementation. + +Do not duplicate expectations for edges that do not exist. + +Where ordering is deterministic, verify the feature associated with each `(src, dst)` pair. + +--- + +# 4. Tracking semantics + +Characterize tracking information. + +Verify at minimum the observed semantics of: + +```text +tracking[:, 0] +tracking[:, 1] +``` + +Expected meanings from the architecture documentation are: + +```text +column 0 -> fold identifier +column 1 -> event weight +``` + +Confirm those semantics against the legacy implementation. + +Also characterize: + +* dtype +* shape +* whether tracking information is mutated during dataset processing +* behavior of negative or non-unit weights where easily testable + +Do not expand into training-weight semantics yet. + +--- + +# 5. Labels + +Characterize labels produced by the active standard dataset path. + +Verify: + +* shape +* dtype +* scalar vs vector representation +* binary-task representation where applicable +* multiclass representation where practical + +Focus on the active configurations. + +Do not characterize every experimental loss or task variant. + +--- + +# 6. Global features + +Characterize global-feature behavior in the active path. + +Verify: + +* whether globals are present +* shape +* dtype +* empty/default behavior when global features are unused + +Do not invent a new global-feature abstraction. + +The purpose is only to record legacy behavior. + +--- + +# 7. Fixtures + +Use deterministic fixtures under: + +```text +tests/fixtures/ +``` + +Prefer, in order: + +1. small in-memory NumPy/Awkward objects +2. a very small ROOT file only where actual ROOT I/O must be tested +3. compact stored expected outputs only when inline arrays become unwieldy + +Avoid large external datasets. + +Fixtures should exercise meaningful cases such as: + +* one object +* two objects +* three or more objects +* multiple object types if relevant +* nontrivial eta values +* phi wrapping across `-pi/pi` +* different node types +* non-unit event weights + +Test empty object collections only if the active legacy path supports them meaningfully. + +Do not invent impossible legacy inputs just to increase coverage. + +--- + +# 8. Test organization + +Place characterization tests under: + +```text +tests/parity/ +``` + +A reasonable organization is: + +```text +tests/parity/ + conftest.py + test_legacy_node_features.py + test_legacy_graph_and_edges.py + test_legacy_tracking.py + test_legacy_labels_and_globals.py +``` + +Use fewer files if that produces clearer tests. + +Shared legacy import/setup helpers may live in: + +```text +tests/parity/conftest.py +``` + +Do not copy production functions from the legacy tree into tests. + +Tests should invoke the legacy implementation directly. + +--- + +# 9. Legacy imports + +Avoid permanently modifying Python import paths outside the test process. + +If the legacy package requires path adjustment because it is not installable, isolate that behavior in test setup. + +Do not modify: + +```text +legacy/root_gnn_dgl/ +``` + +to make imports easier. + +Do not introduce repository-wide `sys.path` hacks into production code. + +--- + +# 10. Expected outputs + +Prefer readable explicit expectations. + +For example: + +```python +expected_features = torch.tensor( + [ + [...], + [...], + ], + dtype=torch.float32, +) +``` + +For topology: + +```python +expected_src = torch.tensor([...]) +expected_dst = torch.tensor([...]) +``` + +Use exact equality for: + +* shapes +* integer IDs +* node-type codes +* edge endpoints +* fold values + +where deterministic. + +Use numerical comparison for floating-point values. + +Avoid opaque snapshot testing unless there is a strong reason. + +--- + +# 11. Floating-point tolerances + +Use explicit tolerances for numerical parity. + +For normal float32 feature calculations, prefer something approximately like: + +```python +torch.testing.assert_close( + actual, + expected, + rtol=1e-5, + atol=1e-6, +) +``` + +unless observed legacy behavior requires a different tolerance. + +Do not use unnecessarily loose tolerances. + +If a different tolerance is required, explain why. + +--- + +# 12. Characterization vs specification + +These tests are characterization tests. + +They answer: + +> What does the legacy system currently do? + +They do not automatically answer: + +> What should the redesigned system do forever? + +When unusual behavior is discovered, classify it as one of: + +```text +compatibility requirement +legacy observation +ambiguous behavior +likely bug requiring later design decision +``` + +Do not silently "fix" questionable behavior in this task. + +--- + +# 13. Known graph expectation correction + +A previous test incorrectly expected four edges for a two-node graph. + +The running legacy implementation produced two edges. + +Investigate and characterize this behavior explicitly. + +If confirmed, the likely topology is: + +```text +directed fully connected graph +without self-loops +``` + +which gives: + +```text +num_edges = N * (N - 1) +``` + +Do not modify legacy code to make it match the old test expectation. + +Update the parity test to match actual observed behavior. + +Also verify the actual `(src, dst)` pairs and corresponding edge features. + +--- + +# 14. Documentation of ambiguities + +Record meaningful discoveries in: + +```text +docs/migration.md +``` + +or a concise dedicated section if appropriate. + +Potential ambiguities include: + +* exact edge ordering +* single-node graph behavior +* empty graph behavior +* missing ROOT branches +* phi wrapping convention +* negative weights +* object concatenation ordering +* shape behavior for unused global features + +Only document ambiguities actually encountered. + +Do not expand documentation unnecessarily. + +--- + +# Out of scope + +Do not implement or modify: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/features/ +src/gnn4colliders/graphs/ +src/gnn4colliders/models/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +``` + +except for existing nonfunctional scaffolding if absolutely required for test discovery, and preferably not at all. + +Do not implement: + +* new data readers +* new feature builders +* new graph builders +* model code +* training code +* losses +* metrics +* checkpoint handling +* Hydra configuration behavior +* runtime optimizations + +Do not refactor legacy code. + +Do not modernize legacy code. + +Do not run Ruff against the frozen legacy tree. + +--- + +# Validation + +Run: + +```bash +uv run python -c "import torch, dgl; print(torch.__version__, dgl.__version__)" + +uv run pytest tests/parity -v + +uv run ruff check tests + +uv run ruff format --check tests +``` + +Then run the full current suite: + +```bash +uv run pytest +``` + +The expected result for Task 3 is: + +```text +all characterization tests pass +no DGL tests skipped because DGL is unavailable +``` + +If a test fails because its expectation disagrees with running legacy behavior: + +1. investigate the legacy implementation +2. determine the observed behavior +3. fix the test expectation if the test was wrong +4. document ambiguity if appropriate + +Do not modify legacy behavior to satisfy a characterization test. + +--- + +# Git inspection + +Before committing, inspect: + +```bash +git status +git diff +``` + +Verify that: + +* legacy files are unchanged +* no new production implementation exists +* fixtures are small +* no caches are tracked +* no model checkpoints were added +* no generated DGL files were added unintentionally +* no unrelated changes are included + +--- + +# Completion criteria + +Task 3 is complete when: + +1. PyTorch and DGL import successfully. +2. Active node-feature behavior is characterized. +3. Graph topology is characterized. +4. Edge features and phi wrapping are characterized. +5. Tracking semantics are characterized. +6. Labels and global-feature behavior are characterized to the extent exercised by the active path. +7. Important ambiguities are documented. +8. All parity tests pass. +9. No DGL-dependent parity tests are skipped because of missing dependencies. +10. Legacy production code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. fixtures created +4. behaviors characterized +5. graph topology observed +6. edge ordering observed +7. numerical tolerances used +8. ambiguities discovered +9. behavior classified as compatibility requirement vs observation +10. validation commands and results +11. anything intentionally deferred + +After validation succeeds, create one Git commit containing only Task 3 changes. + +Use: + +```text +test: characterize legacy data and graph behavior +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task4.md b/tasks/task4.md new file mode 100644 index 0000000000000000000000000000000000000000..7aca325aaa3d376744b6ea6359af83432bc94c12 --- /dev/null +++ b/tasks/task4.md @@ -0,0 +1,473 @@ +# Task 4: Implement Shared Physics Feature Construction + +Implement the first production component of the GNN4Colliders rewrite: the shared collider-object feature construction layer. + +This task should implement new code under: + +```text +src/gnn4colliders/features/ +``` + +Do not implement graph construction, datasets, training, or models yet. + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +tests/parity/ +legacy/root_gnn_dgl/ +``` + +Use the Task 3 characterization tests as the behavioral specification. + +--- + +## Goal + +Implement architecture-independent collider-object feature construction that reproduces the active legacy behavior where compatibility is required. + +The output of this layer should be plain feature arrays/tensors. + +It should not depend on: + +```text +DGL +graph topology +training code +model code +``` + +The same feature layer should be reusable by future model families such as: + +```text +root_gnn +root_transformer +``` + +--- + +# Scope + +Implement the active node/object feature behavior characterized in Task 3. + +At minimum, support the active feature schema: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +Verify all details against the parity tests and legacy implementation. + +Implement only behavior needed by the active standard workflow. + +Do not migrate experimental or unused feature paths unless required by tests. + +--- + +# Package structure + +Prefer a small structure such as: + +```text +src/gnn4colliders/features/ + __init__.py + objects.py +``` + +If selection logic is genuinely needed for this task, a separate module such as: + +```text +selections.py +``` + +is acceptable. + +Do not create many speculative modules. + +--- + +# API design + +Design a small explicit API. + +For example, something conceptually similar to: + +```python +def build_object_features(...): + ... +``` + +or: + +```python +def build_node_features(...): + ... +``` + +The exact function name and arguments should be chosen based on the actual legacy inputs and characterization tests. + +Prefer: + +* explicit inputs +* explicit outputs +* typed signatures +* no hidden global configuration +* no mutation of caller-owned data +* no dependency on repository-relative paths + +Do not expose legacy dynamic-import configuration behavior through this API. + +--- + +# Feature behavior + +Implement and test: + +## pt + +Preserve active legacy behavior. + +## eta + +Preserve active legacy behavior. + +## phi + +Preserve active legacy behavior. + +Do not normalize or wrap phi unless the legacy node-feature path actually does so. + +## energy + +Support the active derived-energy behavior. + +If Task 3 confirmed: + +```text +CALC_E = pt * cosh(eta) +``` + +implement that behavior directly and test it. + +Do not add alternative physics definitions unless required. + +## btag + +Preserve active legacy constants/branch behavior. + +## charge + +Preserve active legacy constants/branch behavior. + +## node_type + +Preserve active legacy node-type coding and object ordering. + +Do not redesign node-type identifiers in this task. + +--- + +# Feature ordering + +The output column order must match the characterized active behavior. + +Expected order is likely: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +but verify against tests. + +Add a named constant or schema representation if that makes the ordering explicit and reduces magic indices. + +For example: + +```python +NODE_FEATURE_NAMES = ( + "pt", + "eta", + "phi", + "energy", + "btag", + "charge", + "node_type", +) +``` + +Keep it simple. + +--- + +# Feature scaling + +If Task 3 characterized legacy feature scaling, implement that behavior. + +Scaling should be explicit. + +Avoid hidden mutable global scale state. + +Prefer passing scale information into the feature builder or representing it through a small typed configuration object. + +Do not introduce Hydra integration yet. + +--- + +# Multiple object types + +If the active legacy path combines multiple collider-object collections: + +* preserve concatenation order +* preserve node-type identifiers +* preserve feature ordering +* preserve dtype + +Characterize any additional behavior with unit tests. + +Do not generalize prematurely to every possible collider object type. + +Implement only what the active workflow requires. + +--- + +# Data representation + +Prefer the natural representation used by the implementation and tests. + +Awkward Arrays may be accepted at the boundary where useful. + +The feature layer should produce a predictable dense representation suitable for later graph or sequence construction. + +Do not return a DGL graph. + +Do not make DGL a dependency of this module. + +--- + +# Dtype behavior + +Preserve the characterized dtype where compatibility matters. + +Avoid relying on implicit dtype conversion. + +Use explicit dtype handling where appropriate. + +Add tests that assert dtype. + +--- + +# Tests + +Add new unit tests under: + +```text +tests/unit/features/ +``` + +or an equivalent clear location. + +Unit tests should test the new implementation directly. + +Then adapt or extend parity testing so that the new implementation is compared against the legacy behavior. + +The important distinction is: + +```text +tests/unit/ + tests the new implementation in isolation + +tests/parity/ + compares new behavior with legacy/reference behavior +``` + +Do not remove the legacy characterization tests. + +--- + +# Parity + +For deterministic feature construction, parity should be as strict as practical. + +Verify: + +* shape +* dtype +* column ordering +* numeric values +* object ordering +* node-type values +* scaling +* derived energy + +Use explicit tolerances only for floating-point calculations. + +Discrete values should use exact equality. + +--- + +# No graph implementation + +Do not implement: + +```text +fully connected graphs +DGLGraph construction +edge indices +deta +dphi +dR +graph padding +``` + +Those belong to the next task. + +Even if graph construction is simple, keep this task scoped to shared feature construction. + +--- + +# No ROOT I/O rewrite + +Do not implement the new ROOT reader yet unless absolutely necessary for a small interface boundary. + +Tests may continue to construct small in-memory arrays/events. + +The new feature implementation should not need to open ROOT files itself. + +ROOT/Awkward I/O will be implemented separately. + +--- + +# No configuration system + +Do not implement Hydra schemas or YAML loading in this task. + +The feature API should be usable directly from Python. + +Configuration will wrap it later. + +--- + +# Code quality + +Use: + +* type annotations +* docstrings for public functions +* small pure functions where practical +* explicit constants for schema/order +* no hidden global state + +Avoid: + +* giant utility modules +* class hierarchies without a demonstrated need +* dynamic imports +* mutation of input arrays +* model-specific assumptions + +--- + +# Documentation + +Update `docs/migration.md` to mark feature construction as implemented only after parity passes. + +If useful, add a short description of the feature schema to: + +```text +docs/architecture.md +``` + +Do not duplicate large amounts of legacy documentation. + +--- + +# Validation + +Run: + +```bash +uv run pytest tests/unit +uv run pytest tests/parity +uv run pytest + +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Expected outcome: + +```text +all relevant tests pass +no legacy files modified +no graph implementation added +``` + +Inspect: + +```bash +git status +git diff +``` + +before committing. + +--- + +# Completion criteria + +Task 4 is complete when: + +1. New shared feature construction exists under `src/gnn4colliders/features/`. +2. Active node-feature behavior matches the legacy characterization tests. +3. Feature order is explicit. +4. Dtypes are explicit. +5. Derived energy behavior is covered. +6. Node-type behavior is covered. +7. Multiple-object ordering is covered where required. +8. Unit tests pass. +9. Parity tests pass. +10. No DGL graph construction has been implemented. +11. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. public API added +4. feature schema implemented +5. parity results +6. tolerances used +7. intentional deviations, if any +8. validation commands and results +9. anything deferred + +After validation succeeds, create a single commit: + +```text +feat: implement shared collider feature construction +``` + +Do not include unrelated changes. diff --git a/tasks/task5.md b/tasks/task5.md new file mode 100644 index 0000000000000000000000000000000000000000..0aceb1844376ac866e2c2be130e894f46f410ba6 --- /dev/null +++ b/tasks/task5.md @@ -0,0 +1,660 @@ +# Task 5: Implement ROOT-GNN Graph Construction + +Implement the graph-construction layer for the active ROOT-GNN workflow. + +This task should build on the shared feature-construction code completed in Task 4. + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/features/ +tests/unit/ +tests/parity/ +legacy/root_gnn_dgl/ +``` + +Use the Task 3 characterization tests and Task 4 feature implementation as the behavioral specification. + +Do not implement datasets, caching, training, or models in this task. + +--- + +## Goal + +Implement the representation-specific transformation: + +```text +node features + -> +graph topology + -> +edge indices + -> +edge features + -> +DGLGraph +``` + +for the active ROOT-GNN path. + +The graph layer should live under: + +```text +src/gnn4colliders/graphs/ +``` + +This layer is specific to graph-based models and should not be used by future non-graph model families such as `root_transformer`. + +--- + +# 1. Scope + +Implement only the active graph behavior required by the current ROOT-GNN workflow. + +At minimum, characterize and reproduce: + +* directed graph construction +* fully connected topology +* self-loop behavior +* deterministic source/destination indices +* edge feature construction +* `[deta, dphi, dR]` feature ordering +* phi wrapping behavior +* single-node graph behavior +* dtype and shape behavior + +Do not migrate experimental graph variants unless required by active configs or parity tests. + +--- + +# 2. Package structure + +Prefer a small structure such as: + +```text +src/gnn4colliders/graphs/ + __init__.py + topology.py + edges.py + dgl.py +``` + +If a smaller structure is cleaner, use fewer modules. + +Do not create speculative abstractions. + +A reasonable division is: + +```text +topology.py + create edge indices + +edges.py + calculate edge features + +dgl.py + assemble a DGLGraph +``` + +The topology and edge-feature logic should ideally be testable without requiring graph mutation. + +--- + +# 3. Public API + +Design a small explicit API. + +Something conceptually similar to: + +```python +def fully_connected_edges( + num_nodes: int, + *, + self_loops: bool = False, +) -> tuple[torch.Tensor, torch.Tensor]: + ... +``` + +and: + +```python +def build_edge_features( + node_features: torch.Tensor, + src: torch.Tensor, + dst: torch.Tensor, + *, + eta_index: int, + phi_index: int, +) -> torch.Tensor: + ... +``` + +and eventually: + +```python +def build_dgl_graph( + node_features: torch.Tensor, +) -> dgl.DGLGraph: + ... +``` + +The exact API may differ if a simpler design is justified. + +Prefer: + +* typed signatures +* explicit inputs +* explicit outputs +* no hidden global state +* no repository-relative paths +* minimal mutation + +--- + +# 4. Graph topology + +Reproduce the active legacy topology exactly where parity requires it. + +Task 3 observed the active graph behavior as a directed fully connected graph without self-loops. + +For `N` nodes, verify: + +```text +num_edges = N * (N - 1) +``` + +where supported. + +Explicitly test: + +```text +N = 1 +N = 2 +N = 3 +``` + +For example, a two-node graph should contain: + +```text +0 -> 1 +1 -> 0 +``` + +and no: + +```text +0 -> 0 +1 -> 1 +``` + +Do not infer or introduce self-loops unless legacy behavior requires them. + +--- + +# 5. Edge ordering + +Preserve exact edge ordering if Task 3 demonstrated that it is deterministic and relied upon. + +If DGL's construction convention determines stable ordering, encode and test it. + +If exact order should not be part of the long-term API, still ensure edge features remain associated with the correct `(src, dst)` pairs. + +Do not accidentally compare feature arrays independently of their edges. + +--- + +# 6. Edge features + +Implement the active edge feature schema: + +```text +deta +dphi +dR +``` + +in exactly that order. + +The expected semantics should be verified against the Task 3 characterization. + +For each directed edge: + +```text +src -> dst +``` + +calculate the same signed `deta` and signed wrapped `dphi` convention used by the legacy implementation. + +Then calculate: + +```text +dR = sqrt(deta^2 + dphi^2) +``` + +if confirmed by parity. + +--- + +# 7. Phi wrapping + +Implement phi differences explicitly. + +Do not rely on naive: + +```python +phi_src - phi_dst +``` + +when crossing `-pi` / `+pi`. + +Use the exact wrapping convention characterized from legacy behavior. + +Add unit tests for cases including: + +```text +phi_src ~= +pi +phi_dst ~= -pi +``` + +and the reverse direction. + +Verify signs as well as magnitudes. + +Avoid unnecessarily introducing a heavyweight angular/physics dependency for this small operation. + +--- + +# 8. Node feature interface + +The graph builder should consume the output of Task 4's feature layer. + +Do not duplicate feature-construction logic. + +Do not recompute: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +inside graph code. + +The graph layer may use explicit feature indices or a shared schema constant from `gnn4colliders.features`. + +Prefer named/schema-based access over unexplained magic numbers. + +For example, if Task 4 exposes: + +```python +NODE_FEATURE_NAMES +``` + +use it or an associated index mapping rather than hardcoding `eta = features[:, 1]` throughout the code. + +--- + +# 9. DGL integration + +Use DGL only at the representation boundary. + +Prefer logic that conceptually separates: + +```text +topology calculation +edge-feature calculation +DGL object construction +``` + +Do not make every helper require a DGLGraph. + +The DGL-specific builder should: + +1. create graph topology +2. attach node features +3. attach edge features + +using the active legacy key names if compatibility requires them. + +For example, if legacy behavior uses: + +```python +graph.ndata["features"] +graph.edata["features"] +``` + +preserve those names for ROOT-GNN compatibility. + +--- + +# 10. Device behavior + +Do not add complicated automatic device movement. + +Prefer graph construction on CPU unless a strong existing reason requires otherwise. + +Do not silently move data to CUDA. + +Device placement should later be controlled by the training/application layer. + +Functions should preserve or clearly document expected device behavior. + +--- + +# 11. Dtypes + +Preserve Task 3/Task 4 dtype behavior. + +Typical expectations may include: + +```text +node features: float32 +edge features: float32 +edge indices: integer +``` + +but verify against actual characterization. + +Avoid implicit conversions. + +Add explicit tests for dtype. + +--- + +# 12. Single-node behavior + +Characterize and reproduce the active behavior for: + +```text +N = 1 +``` + +If the legacy graph has: + +```text +1 node +0 edges +``` + +ensure the new implementation behaves the same. + +Verify the shape of empty edge features. + +For example, expected behavior may be: + +```text +edge_features.shape == (0, 3) +``` + +Do not guess; use characterization tests. + +--- + +# 13. Empty graph behavior + +Do not expand into unsupported empty-event behavior unless Task 3 established it as part of the active path. + +If zero-node graph behavior is ambiguous or unsupported: + +* document it +* do not invent semantics in this task + +--- + +# 14. Tests + +Add focused unit tests under: + +```text +tests/unit/graphs/ +``` + +Suggested organization: + +```text +tests/unit/graphs/ + test_topology.py + test_edge_features.py + test_dgl_graph.py +``` + +Use fewer files if clearer. + +Unit tests should cover the new implementation directly. + +--- + +# 15. Parity tests + +Extend or adapt parity tests so the new graph implementation is compared with the characterized legacy implementation. + +At minimum compare: + +* node count +* edge count +* source indices +* destination indices +* node feature arrays +* edge feature arrays +* shapes +* dtypes + +Use direct legacy execution where practical. + +Do not delete the original legacy characterization assertions. + +The parity suite should now verify both: + +```text +legacy behavior +``` + +and: + +```text +new implementation matches legacy +``` + +where appropriate. + +--- + +# 16. Numerical tolerances + +Use tight explicit tolerances for floating-point calculations. + +Prefer something around: + +```python +torch.testing.assert_close( + actual, + expected, + rtol=1e-5, + atol=1e-6, +) +``` + +unless legacy calculations require otherwise. + +Use exact equality for: + +* source indices +* destination indices +* node counts +* edge counts +* shapes + +--- + +# 17. No dataset implementation + +Do not implement: + +```text +RootDataset +LazyDataset +PreBatchedDataset +cache files +DGL graph serialization +DataLoader integration +fold filtering +``` + +Those belong to subsequent tasks. + +The graph builder should operate on in-memory feature data only. + +--- + +# 18. No training/model implementation + +Do not implement or modify: + +```text +models/ +training/ +inference/ +``` + +Do not migrate `Edge_Network` yet. + +Do not add losses, metrics, optimizers, checkpointing, or distributed execution. + +--- + +# 19. No performance optimization yet + +Do not optimize this code with: + +```text +numba +custom CUDA kernels +torch.compile +C++ extensions +``` + +or other acceleration frameworks. + +Favor correctness, clarity, and parity. + +If graph construction later becomes a measurable bottleneck, optimize it after profiling. + +--- + +# 20. Documentation + +Update `docs/migration.md` once graph construction parity is established. + +If the architecture documentation does not already clearly state it, document that: + +```text +gnn4colliders.features +``` + +contains architecture-independent collider features, while: + +```text +gnn4colliders.graphs +``` + +contains graph-specific representation logic used by ROOT-GNN. + +Keep documentation changes concise. + +--- + +# Validation + +Run: + +```bash +uv run python -c "import torch, dgl; print(torch.__version__, dgl.__version__)" + +uv run pytest tests/unit/graphs -v + +uv run pytest tests/parity -v + +uv run pytest +``` + +Then: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy code is unchanged +* no dataset implementation was added +* no model/training code was added +* no generated graph caches are tracked +* no unrelated files changed + +--- + +# Completion criteria + +Task 5 is complete when: + +1. Graph topology is implemented under `src/gnn4colliders/graphs/`. +2. Directed fully connected no-self-loop behavior matches legacy. +3. Edge ordering matches legacy where required. +4. `[deta, dphi, dR]` matches legacy. +5. Phi wrapping matches legacy in both directions. +6. Node and edge feature attachment matches compatibility requirements. +7. Single-node behavior is covered. +8. Unit tests pass. +9. Legacy/new parity tests pass. +10. No DGL tests are skipped because DGL is unavailable. +11. No dataset, training, or model implementation has been added. +12. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. public API added +4. topology implemented +5. observed edge ordering +6. phi wrapping convention +7. DGL compatibility behavior +8. unit test results +9. parity results +10. numerical tolerances used +11. ambiguities or deferred behavior +12. validation commands and results + +After validation succeeds, create one Git commit containing only Task 5 changes. + +Use: + +```text +feat: implement ROOT-GNN graph construction +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task6.md b/tasks/task6.md new file mode 100644 index 0000000000000000000000000000000000000000..c365e6f728e4a5bb03e2570b40c34d2b0d7f7e1b --- /dev/null +++ b/tasks/task6.md @@ -0,0 +1,707 @@ +# Task 6: Implement ROOT/Awkward Data Ingestion and Sample Representation + +Implement the new shared data-ingestion layer for GNN4Colliders. + +This task should build on the completed feature-construction and graph-construction work from Tasks 4 and 5. + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/features/ +src/gnn4colliders/graphs/ +tests/unit/ +tests/parity/ +legacy/root_gnn_dgl/ +``` + +Use the legacy implementation and existing parity tests as behavioral references. + +Do not implement caching, pre-batching, training, or models in this task. + +--- + +## Goal + +Implement the shared data layer that turns ROOT data into architecture-neutral event/sample representations. + +The intended flow is: + +```text +ROOT file + -> +Uproot / Awkward + -> +event/sample representation + -> +gnn4colliders.features + -> +model-specific representation +``` + +For ROOT-GNN, the later model-specific representation is: + +```text +sample + -> +gnn4colliders.features + -> +gnn4colliders.graphs +``` + +For a future ROOT-Transformer, the same sample/data layer should be reusable without DGL. + +The new data layer should therefore not return DGL graphs directly. + +--- + +# 1. Scope + +Implement only the active shared ROOT/Awkward ingestion behavior required by the standard ROOT-GNN workflow. + +At minimum support: + +* opening ROOT files with Uproot +* selecting the configured tree +* reading required branches +* event indexing +* converting event data into a clean in-memory sample structure +* labels +* tracking information +* global features +* object collections needed by Task 4 feature construction + +Do not migrate every experimental dataset class. + +Focus on the active path only. + +--- + +# 2. Package structure + +Prefer a small structure such as: + +```text +src/gnn4colliders/data/ + __init__.py + root_io.py + sample.py + dataset.py +``` + +Possible responsibilities: + +```text +root_io.py + low-level ROOT/Uproot reading + +sample.py + typed event/sample representation + +dataset.py + dataset abstraction over one or more ROOT files +``` + +If fewer modules are clearer, use fewer. + +Do not create speculative abstractions or a deep class hierarchy. + +--- + +# 3. Architecture-neutral sample representation + +Introduce a small explicit representation for one event/sample. + +Prefer a typed dataclass or similarly simple structure. + +Conceptually something like: + +```python +@dataclass(frozen=True) +class EventSample: + objects: ... + label: ... + tracking: ... + global_features: ... +``` + +The exact fields should reflect the actual active legacy behavior. + +The sample representation should: + +* be independent of DGL +* be independent of ROOT-GNN model classes +* be usable by future model families +* have explicit types +* avoid hidden global state +* avoid mutation where practical + +Do not force raw ROOT branch names to leak throughout downstream code if a cleaner normalized sample structure is practical. + +At the same time, do not over-generalize beyond the active workflow. + +--- + +# 4. ROOT I/O + +Implement low-level ROOT reading with Uproot. + +Responsibilities may include: + +* opening a ROOT file +* selecting a tree by name +* validating requested branches +* reading selected branches +* returning Awkward arrays or an equivalent architecture-neutral structure + +Prefer explicit functions. + +For example, conceptually: + +```python +def read_tree( + path: Path, + tree_name: str, + branches: Sequence[str], +) -> ak.Array: + ... +``` + +The exact API may differ. + +Do not use PyROOT unless the active path requires behavior that Uproot cannot provide. + +Do not add DGL to this module. + +--- + +# 5. Branch selection + +Determine the branch set needed by the active standard configs. + +Do not read every branch in the tree by default if only a subset is required. + +The data layer should make required branch names explicit. + +If branch sets depend on object type or task, represent that cleanly without reproducing the old dynamic-import configuration system. + +Do not implement the full future Hydra configuration layer yet. + +A plain typed argument/config object is acceptable. + +--- + +# 6. Object collections + +Preserve the object-collection semantics needed by Task 4 feature construction. + +For each active object collection, characterize and preserve: + +* object ordering +* branch mapping +* jagged structure +* missing/empty collection behavior +* scalar vs vector branch behavior where relevant + +Do not duplicate feature construction here. + +This layer should expose raw or normalized object data. + +Task 4 remains responsible for constructing: + +```text +pt +eta +phi +energy +btag +charge +node_type +``` + +Do not calculate graph edges in this layer. + +--- + +# 7. Labels + +Characterize and implement active label extraction. + +Verify: + +* source branch or source rule +* scalar vs vector representation +* dtype +* shape +* binary-task behavior +* multiclass behavior where active configs require it + +Do not expand into every historical multi-label/experimental path unless required by active configs or existing parity tests. + +The sample representation should expose labels cleanly. + +--- + +# 8. Tracking information + +Implement tracking extraction matching the active legacy behavior. + +At minimum preserve the characterized semantics: + +```text +tracking column 0 -> fold identifier +tracking column 1 -> event weight +``` + +Verify actual source branches/rules in the legacy implementation. + +Preserve: + +* dtype +* shape +* ordering + +Do not implement fold filtering yet. + +Do not mutate tracking arrays in place. + +If the legacy code mutates tracking information as part of dataset processing, reproduce the resulting external behavior without preserving unnecessary mutation internally. + +--- + +# 9. Global features + +Implement active global-feature extraction. + +Characterize: + +* whether globals are present +* source branches +* shape +* dtype +* empty/default representation when not used + +Do not invent model-specific global semantics. + +The data layer should expose globals as generic sample data. + +--- + +# 10. Dataset abstraction + +Implement a simple dataset abstraction over ROOT data. + +Conceptually: + +```python +dataset = RootEventDataset(...) +sample = dataset[i] +``` + +The exact class/API may differ. + +The dataset should preferably support: + +* `__len__` +* `__getitem__` +* deterministic event indexing + +Do not add DGL-specific behavior. + +Do not make it inherit from a DGL dataset class. + +It may inherit from `torch.utils.data.Dataset` only if that provides clear value and does not force unwanted coupling. + +If plain Python indexing is sufficient, keep it simpler. + +--- + +# 11. One file vs multiple files + +Support the smallest practical file abstraction required by the active workflow. + +If the active standard path requires multiple ROOT files, support them explicitly. + +Preserve deterministic event ordering across files. + +Document the ordering rule. + +Do not implement distributed file sharding yet. + +Do not implement cache chunking yet. + +--- + +# 12. Event indexing + +Characterize and preserve event indexing semantics. + +Determine: + +* whether indexing is global across files +* how event offsets are computed +* whether ordering follows input file order +* whether any implicit sorting occurs + +Prefer a transparent deterministic rule. + +Add tests for file boundaries if multiple files are supported. + +--- + +# 13. Error behavior + +Characterize active behavior for: + +* missing file +* missing tree +* missing required branch +* malformed branch shape + +Do not implement an elaborate validation framework. + +Raise clear Python exceptions. + +Where legacy behavior is poor or accidental, preserve compatibility only if downstream behavior depends on it. + +Document intentional improvements in error clarity. + +--- + +# 14. Small ROOT fixture + +Create a tiny deterministic ROOT fixture if one does not already exist and if actual ROOT I/O needs direct testing. + +Store it under: + +```text +tests/fixtures/ +``` + +Keep it very small. + +It should include enough events to exercise: + +* at least one object collection +* multiple objects in an event +* different event sizes +* labels +* fold information +* non-unit event weights +* global features if active +* empty collection behavior if supported and useful + +Prefer generating the fixture in test setup if that is clearer and avoids committing binary data. + +If a committed ROOT file is used, keep it tiny. + +--- + +# 15. Tests + +Add unit tests under: + +```text +tests/unit/data/ +``` + +Suggested files: + +```text +test_root_io.py +test_sample.py +test_dataset.py +``` + +Use fewer files if clearer. + +Unit tests should cover the new implementation directly. + +--- + +# 16. Parity tests + +Extend parity tests to compare the new data layer against the active legacy behavior. + +At minimum compare, where practical: + +* event count +* per-event object arrays +* object ordering +* labels +* tracking +* global features +* dtypes +* shapes + +Do not compare DGL graphs in this task unless an existing integration test naturally verifies downstream compatibility. + +Graph parity belongs primarily to Task 5. + +--- + +# 17. Integration with Task 4 features + +Add a small integration test showing that one sample from the new dataset can be passed into the Task 4 feature builder. + +Conceptually: + +```text +ROOT fixture + -> +RootEventDataset + -> +EventSample + -> +build_object_features(...) +``` + +Verify the result matches the expected/legacy feature construction. + +Do not continue into DGL graph construction unless a tiny integration check is already natural. + +The main goal is to prove the shared data-to-feature boundary works. + +--- + +# 18. No DGL dependency in shared data layer + +The following should not be imported by: + +```text +gnn4colliders.data +``` + +unless absolutely unavoidable: + +```text +dgl +gnn4colliders.graphs +gnn4colliders.models.root_gnn +``` + +The data layer must remain reusable by future architectures. + +--- + +# 19. No caching yet + +Do not implement: + +* DGL `.bin` cache files +* chunked graph caches +* lazy loading +* ring buffers +* pre-batching +* padding +* cache hashing +* cache invalidation + +Those belong to Task 7. + +This task should read ROOT data and expose samples directly. + +--- + +# 20. No fold selection yet + +Although tracking should expose fold IDs, do not implement train/validation/test fold filtering in this task. + +That belongs with dataset orchestration/batching in Task 7. + +--- + +# 21. No DataLoader/pre-batching yet + +Do not implement: + +* GraphDataLoader +* PyTorch DataLoader tuning +* pre-batching +* padding +* multiprocessing workers +* worker seeding +* distributed samplers + +Those belong to later tasks. + +Keep this task focused on deterministic sample access. + +--- + +# 22. No model or training changes + +Do not modify: + +```text +src/gnn4colliders/models/ +src/gnn4colliders/training/ +src/gnn4colliders/inference/ +``` + +Do not migrate model code. + +Do not implement losses, metrics, optimizers, checkpoints, CLI workflows, or distributed execution. + +--- + +# 23. No configuration-system implementation + +Do not implement Hydra composition yet. + +If the new data APIs require structured settings, use small typed Python objects or explicit function arguments. + +The future config layer should wrap these APIs rather than define their internal behavior. + +--- + +# 24. Code quality + +Prefer: + +* explicit typed APIs +* `pathlib.Path` +* small pure functions +* immutable sample objects where practical +* deterministic ordering +* architecture-neutral data structures +* clear validation errors + +Avoid: + +* mutable global configuration +* repository-relative `sys.path` hacks +* dynamic imports +* giant dataset classes +* hidden file discovery +* implicit environment-specific paths +* model-specific behavior in shared data code + +--- + +# 25. Documentation + +Update `docs/migration.md` once the new data-ingestion layer reaches parity. + +Update `docs/architecture.md` only if needed to clarify the new boundary: + +```text +gnn4colliders.data + ROOT/Awkward I/O and architecture-neutral samples + +gnn4colliders.features + collider feature construction + +gnn4colliders.graphs + graph-specific representation +``` + +Keep documentation changes concise. + +--- + +# Validation + +Run: + +```bash +uv run pytest tests/unit/data -v +``` + +Then run relevant feature integration tests: + +```bash +uv run pytest tests/unit/features -v +uv run pytest tests/parity -v +``` + +Then run the full suite: + +```bash +uv run pytest +``` + +Run lint/format validation: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy code is unchanged +* no DGL-specific data implementation was added +* no cache implementation was added +* no batching implementation was added +* no training/model implementation was added +* fixtures are small +* no generated ROOT or cache artifacts are accidentally tracked + +--- + +# Completion criteria + +Task 6 is complete when: + +1. ROOT/Uproot reading exists under `src/gnn4colliders/data/`. +2. A clean architecture-neutral sample representation exists. +3. Deterministic event indexing works. +4. Required object collections are exposed correctly. +5. Labels match active legacy behavior. +6. Tracking info matches active legacy behavior. +7. Global-feature behavior is characterized and implemented. +8. Shared data code does not depend on DGL. +9. A sample can feed directly into the Task 4 feature layer. +10. Unit tests pass. +11. Relevant parity tests pass. +12. The full current suite passes. +13. Legacy code remains untouched. +14. No caching, batching, model, or training implementation has been added. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. public APIs added +4. sample representation introduced +5. ROOT I/O behavior implemented +6. object collections supported +7. label behavior +8. tracking behavior +9. global-feature behavior +10. event-ordering/indexing rule +11. parity results +12. integration-test results +13. ambiguities discovered +14. intentional deviations from legacy behavior +15. validation commands and results +16. anything deferred + +After validation succeeds, create one Git commit containing only Task 6 changes. + +Use: + +```text +feat: implement ROOT data ingestion and sample representation +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task7.md b/tasks/task7.md new file mode 100644 index 0000000000000000000000000000000000000000..3975b7808a66debc4cac46c5b573326a2097d42d --- /dev/null +++ b/tasks/task7.md @@ -0,0 +1,1108 @@ +# Task 7: Implement Metadata-Aware Dataset Caching, Splits, and Batching + +Implement the dataset-orchestration layer for the active ROOT-GNN workflow. + +This task builds on: + +```text +Task 4: shared physics feature construction +Task 5: ROOT-GNN graph construction +Task 6: ROOT/Awkward ingestion and architecture-neutral samples +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/features/ +src/gnn4colliders/graphs/ +tests/unit/ +tests/parity/ +tests/integration/ +legacy/root_gnn_dgl/ +``` + +Focus on active legacy behavior around: + +```text +RootDataset / LazyDataset +PreBatchedDataset +fold_selection +GraphDataLoader +padding +cache loading +``` + +Do not implement the model or training lifecycle yet. + +--- + +# Goal + +Implement the layer that turns processed event samples and graphs into deterministic train/validation/test datasets and batches ready for a future training loop. + +The intended flow is: + +```text +ROOT files + -> +EventSample + -> +shared physics features + -> +graph construction + -> +GraphSample + -> +cache + -> +metadata-based split selection + -> +batching / optional padding + -> +GraphBatch / DataLoader +``` + +The future training layer should consume this API without needing to understand legacy dataset internals. + +The design should also remain compatible with future model families such as `root_transformer`. + +--- + +# 1. Replace legacy tracking columns with named metadata + +The new implementation must not preserve the legacy positional tracking tensor as its public representation. + +Legacy behavior currently encodes: + +```text +tracking[:, 0] -> fold identifier +tracking[:, 1] -> event weight +``` + +Replace that with explicit named metadata. + +At minimum: + +```python +metadata["fold"] +metadata["weight"] +``` + +The representation may be typed rather than a raw dictionary, but downstream users must not need to know positional column numbers. + +Parity should preserve semantics, not the legacy representation. + +--- + +# 2. Use a typed EventMetadata structure + +Prefer a small typed structure rather than an unconstrained dictionary. + +For example: + +```python +from dataclasses import dataclass, field +from typing import Any, Mapping + +@dataclass(frozen=True) +class EventMetadata: + fold: int + weight: float + sample_id: str + extra: Mapping[str, Any] = field(default_factory=dict) +``` + +The exact implementation may differ if existing Task 6 code suggests a better design. + +Requirements: + +* `fold` has explicit meaning +* `weight` has explicit meaning +* every sample has a stable `sample_id` +* optional future metadata goes into `extra` +* avoid hidden positional semantics +* keep the structure lightweight and serializable + +Do not build a large metadata framework. + +--- + +# 3. Stable sample identity + +Every event/sample should have a stable identity. + +Use a deterministic identifier based on source provenance, for example conceptually: + +```text +source file + tree + entry index +``` + +The exact representation may be a string or structured fields. + +Example: + +```python +sample_id = "file.root:Events:12345" +``` + +The same source event should receive the same `sample_id` across repeated runs. + +Use this ID for: + +* debugging +* cache round-trips +* batch alignment tests +* future inference output association + +Do not use a random UUID. + +--- + +# 4. Separate event metadata from cache metadata + +Use distinct concepts and names. + +## EventMetadata + +Describes one physics event. + +Examples: + +```text +fold +weight +sample_id +event_number +run_number +source_file +``` + +## CacheMetadata + +Describes a cache artifact or cache chunk. + +Examples: + +```text +feature schema version +graph schema version +cache schema version +input source +preprocessing configuration +software version +``` + +Do not use one generic `metadata` object for both purposes. + +--- + +# 5. Explicit schema versioning + +Introduce simple schema-version constants. + +Prefer something like: + +```python +FEATURE_SCHEMA_VERSION = 1 +GRAPH_SCHEMA_VERSION = 1 +CACHE_SCHEMA_VERSION = 1 +``` + +Use clear names and keep versioning simple. + +Cache metadata should include these versions. + +The purpose is to prevent stale caches from silently loading after preprocessing or graph semantics change. + +Do not implement semantic-version machinery or a migration framework. + +--- + +# 6. EventSample contract + +If Task 6 currently exposes a sample type, update it if necessary so the new architecture is explicit. + +Conceptually: + +```python +@dataclass(frozen=True) +class EventSample: + objects: ... + label: ... + global_features: ... + metadata: EventMetadata +``` + +Avoid exposing legacy `tracking`. + +If Task 6 already introduced named metadata, preserve that design and extend it. + +--- + +# 7. GraphSample contract + +Introduce or update a graph-level sample representation. + +Prefer something conceptually like: + +```python +@dataclass +class GraphSample: + graph: dgl.DGLGraph + label: torch.Tensor + global_features: torch.Tensor | None + metadata: EventMetadata +``` + +This should represent one processed event for ROOT-GNN. + +Keep it small and explicit. + +Do not create inheritance hierarchies. + +--- + +# 8. GraphBatch contract + +Introduce a clear batch representation. + +Conceptually: + +```python +@dataclass +class GraphBatch: + graph: dgl.DGLGraph + labels: torch.Tensor + global_features: torch.Tensor | None + metadata: ... +``` + +Batched metadata must preserve one-to-one correspondence with events. + +For example: + +```python +batch.metadata.fold +batch.metadata.weight +batch.metadata.sample_id +``` + +or equivalent. + +The implementation may use tensors/lists internally, but the meaning must remain explicit. + +--- + +# 9. Metadata batching rules + +Define deterministic batching rules. + +For homogeneous numeric scalar fields: + +```text +fold +weight +event_number +run_number +``` + +prefer tensors where practical. + +For strings or heterogeneous metadata: + +```text +sample_id +source_file +sample_name +``` + +prefer ordered sequences/lists. + +For example: + +```python +BatchMetadata( + fold=torch.tensor([0, 1, 1]), + weight=torch.tensor([1.0, 0.8, 1.2]), + sample_id=[ + "a.root:Events:1", + "a.root:Events:2", + "b.root:Events:4", + ], +) +``` + +Do not collapse heterogeneous metadata into an opaque tensor. + +--- + +# 10. Named-field convention + +Add or enforce this project-wide rule: + +> Domain meaning should be represented by named fields or schemas, not only by positional column indices. + +Do not introduce new APIs that require users to know things like: + +```text +column 0 means fold +column 1 means weight +column 5 means event number +``` + +This rule should be reflected in `AGENTS.md` if not already present. + +--- + +# 11. Legacy compatibility mapping + +At the compatibility boundary, explicitly map: + +```text +legacy tracking[:, 0] -> EventMetadata.fold +legacy tracking[:, 1] -> EventMetadata.weight +``` + +Add parity tests for this mapping. + +Do not reconstruct a legacy tracking tensor in normal production APIs. + +A test-only compatibility helper is acceptable if needed. + +--- + +# 12. Fold selection + +Implement fold selection through named metadata: + +```python +sample.metadata.fold +``` + +not positional tracking. + +Provide a clean API. + +Conceptually: + +```python +def select_folds( + samples, + folds: Collection[int], +): + ... +``` + +Verify: + +* correct membership +* deterministic ordering +* no mutation +* sample IDs remain stable + +--- + +# 13. Split definitions + +Provide an explicit abstraction for train/validation/test splits. + +A small dataclass is acceptable: + +```python +@dataclass(frozen=True) +class SplitDefinition: + train_folds: frozenset[int] + validation_folds: frozenset[int] + test_folds: frozenset[int] +``` + +The exact API may differ. + +Do not hardcode specific fold numbers in generic data code. + +Keep the split definition independent of Hydra for now. + +--- + +# 14. Keep split machinery extensible + +Fold-based splitting is the current implementation, but avoid making the entire dataset layer inherently dependent on folds forever. + +Structure the code so alternative selectors could be added later without rewriting batching or caching. + +Do not over-engineer a generic query language. + +A simple separation between: + +```text +split definition +selection implementation +dataset/batching +``` + +is sufficient. + +--- + +# 15. Event weights + +Preserve legacy weight semantics as: + +```python +sample.metadata.weight +``` + +Maintain: + +* numeric value +* sign +* dtype where relevant +* event alignment + +Do not apply absolute value here. + +Loss-weight interpretation belongs to the future losses/training task. + +--- + +# 16. Cache architecture + +Implement processed-data caching with clear boundaries. + +Prefer a cache abstraction around: + +```text +GraphSample +``` + +for the active ROOT-GNN path. + +However, design the cache code so it does not prevent a future representation-neutral cache layer. + +Do not tightly couple all cache APIs to ROOT-GNN terminology if a generic sample cache abstraction is straightforward. + +--- + +# 17. Consider two cache levels + +Evaluate whether the architecture should support two conceptual cache stages: + +```text +Level 1: +ROOT + -> +architecture-neutral EventSample / normalized feature data + +Level 2: +features + -> +GraphSample +``` + +Do not necessarily implement both fully in this task. + +At minimum, structure the cache layer so a future ROOT-Transformer can reuse expensive ROOT preprocessing without requiring DGL graph caches. + +If implementing only graph caching now is simplest, document the future separation point. + +Do not create speculative unused code just to represent Level 1. + +--- + +# 18. CacheMetadata + +Introduce a typed or explicit cache metadata representation. + +Conceptually: + +```python +@dataclass(frozen=True) +class CacheMetadata: + cache_schema_version: int + feature_schema_version: int + graph_schema_version: int + source_files: tuple[str, ...] + tree_name: str +``` + +Additional stable fields may be added if necessary. + +Keep cache metadata human-readable and inspectable where practical. + +--- + +# 19. Cache provenance + +Record enough provenance to understand how cached data was produced. + +Consider including: + +```text +source files +tree name +feature schema version +graph schema version +cache schema version +preprocessing config fingerprint/hash +GNN4Colliders version or git commit if practical +``` + +Do not make Git availability a hard runtime requirement. + +If recording the Git commit is awkward or unreliable, make it optional. + +--- + +# 20. Configuration fingerprint + +If preprocessing configuration affects cache validity, create a deterministic representation or fingerprint. + +Avoid using unstable Python object hashes. + +Prefer deterministic serialization of the relevant configuration followed by a stable hash. + +Do not include unrelated training configuration. + +Cache validity should reflect preprocessing/representation choices only. + +--- + +# 21. Cache validity + +On load, verify relevant schema/provenance metadata. + +Do not silently load obviously incompatible caches. + +Raise a clear error for incompatible schema versions. + +Do not implement automatic migration of old cache formats in this task. + +--- + +# 22. Cache format + +Choose a simple format appropriate for the active DGL/PyTorch path. + +DGL-native graph serialization is acceptable. + +Isolate file-format details behind APIs such as: + +```python +save_graph_cache(...) +load_graph_cache(...) +``` + +The future trainer should not depend on file layout. + +--- + +# 23. Cache round-trip + +Test: + +```text +GraphSample + -> +save + -> +load + -> +GraphSample +``` + +Verify: + +* graph topology +* node features +* edge features +* label +* global features +* EventMetadata.fold +* EventMetadata.weight +* EventMetadata.sample_id +* optional metadata +* schema metadata + +Sample IDs must survive exactly. + +--- + +# 24. Chunking + +Implement deterministic chunking only if needed for realistic dataset size. + +Do not preserve legacy `np.array_split` boundaries merely for historical fidelity. + +Exact chunk IDs are implementation details unless external consumers require them. + +Chunk ordering must be deterministic. + +--- + +# 25. Lazy loading + +If chunked caches are implemented, allow data to be loaded without materializing the entire dataset. + +Keep it simple. + +Do not reproduce the legacy ring buffer unless profiling later demonstrates a need. + +--- + +# 26. Deterministic ordering + +Preserve deterministic ordering through: + +```text +input +processing +cache +split +batch +``` + +unless explicit shuffling is requested. + +Use `sample_id` in tests to verify ordering and alignment. + +Do not rely on filesystem iteration order. + +--- + +# 27. Deterministic shuffling + +Use explicit local RNG state. + +For example: + +```python +generator = torch.Generator() +generator.manual_seed(seed) +``` + +Avoid: + +* process-wide NumPy seed mutation +* hidden calls to `torch.manual_seed` +* unseeded shuffles + +Given identical: + +```text +dataset +configuration +seed +``` + +the sample order must be reproducible. + +Verify ordering using `sample_id`. + +--- + +# 28. Batching + +Batch graphs with DGL's normal batching primitives where appropriate. + +Preserve alignment between: + +```text +graph +label +global features +EventMetadata +``` + +Tests should verify batch correspondence using `sample_id`. + +--- + +# 29. Collation + +Implement an explicit collate function. + +Conceptually: + +```python +def collate_graph_samples( + samples: Sequence[GraphSample], +) -> GraphBatch: + ... +``` + +It should: + +1. batch graphs +2. stack labels +3. stack globals if present +4. collate named metadata +5. preserve sample ID order + +Do not hide this logic in an anonymous lambda. + +--- + +# 30. DataLoader + +Provide a small loader-construction API. + +Conceptually: + +```python +def build_graph_dataloader( + dataset, + *, + batch_size: int, + shuffle: bool, + seed: int, + num_workers: int = 0, +): + ... +``` + +Keep training-specific decisions out of this function. + +Do not implement DDP yet. + +--- + +# 31. Padding + +Inspect active legacy padding behavior. + +Legacy modes may include: + +```text +NONE +STEPS +FIXED +NODE +``` + +Implement: + +```text +NONE +``` + +plus only the modes required by active standard configs. + +Do not carry forward unused padding modes merely for completeness. + +--- + +# 32. Padding semantics + +For implemented padding modes, test: + +* padded node count +* padded edge count +* feature shape +* dummy feature values +* sample/metadata alignment + +Do not preserve arbitrary legacy constants unless they are active compatibility requirements. + +If a historical fixed padding size is not required, make it configurable or defer it. + +--- + +# 33. Keep padding separate from graph construction + +Task 5 should continue to build real event topology only. + +Padding belongs in batching/orchestration. + +Do not modify core graph-building semantics to support padding. + +--- + +# 34. Integration flow + +Add an end-to-end integration test covering: + +```text +tiny ROOT fixture + -> +EventSample + -> +feature builder + -> +graph builder + -> +GraphSample + -> +cache + -> +split + -> +DataLoader + -> +GraphBatch +``` + +Verify: + +* `sample_id` +* fold +* weight +* label +* globals +* graph topology +* deterministic ordering +* cache round-trip +* batch alignment + +Do not involve a neural network. + +--- + +# 35. Tests + +Add or update tests for: + +```text +EventMetadata +sample_id generation +legacy tracking -> EventMetadata mapping +SplitDefinition +fold filtering +CacheMetadata +cache schema mismatch +cache round-trip +deterministic shuffle +metadata batching +sample alignment +padding +DataLoader output +``` + +Prefer small deterministic fixtures. + +--- + +# 36. No model or training implementation + +Do not implement: + +```text +models/root_gnn/ +losses +metrics +optimizers +training loop +checkpointing +DDP +``` + +Those belong to later tasks. + +--- + +# 37. No Hydra wiring yet + +Do not wire these components directly to Hydra. + +Use typed Python APIs. + +Hydra will configure these objects from outside in a later task. + +--- + +# 38. No premature optimization + +Do not add: + +```text +numba +custom CUDA +torch.compile +large process-global caches +complex multiprocessing +``` + +Correctness, clarity, reproducibility, and stable interfaces come first. + +Document performance concerns for later profiling. + +--- + +# 39. Documentation + +Update `docs/architecture.md` to reflect the intended contracts: + +```text +EventSample + architecture-neutral event representation + +EventMetadata + named per-event provenance/training metadata + +GraphSample + graph-specific representation for ROOT-GNN + +GraphBatch + training-ready batch + +CacheMetadata + cache schema and provenance +``` + +Explicitly document: + +```text +Legacy: +tracking[:, 0] = fold +tracking[:, 1] = weight + +New: +EventMetadata.fold +EventMetadata.weight +``` + +Document that named fields replace positional semantics intentionally. + +Update `docs/migration.md` once Task 7 reaches parity. + +If `AGENTS.md` does not already include the rule, add: + +```text +Prefer named schemas and typed structures over positional domain semantics. +``` + +Keep documentation concise. + +--- + +# Validation + +Run focused tests: + +```bash +uv run pytest tests/unit/data -v +uv run pytest tests/unit/graphs -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy code remains unchanged +* no positional tracking API is introduced +* sample IDs are stable +* event metadata and cache metadata are separate +* schema versions are persisted +* cache round-trip preserves metadata +* no model/training implementation was added +* no large cache artifacts are tracked +* no `/global/cfs` or `/pscratch` paths are hardcoded +* no unrelated files are included + +--- + +# Completion criteria + +Task 7 is complete when: + +1. `EventMetadata` or equivalent typed named metadata exists. +2. Legacy fold semantics map to `EventMetadata.fold`. +3. Legacy weight semantics map to `EventMetadata.weight`. +4. Every event has a deterministic stable `sample_id`. +5. New APIs no longer expose positional tracking as the preferred representation. +6. `EventSample` uses named metadata. +7. `GraphSample` uses named metadata. +8. `GraphBatch` preserves named metadata and sample order. +9. `CacheMetadata` is distinct from `EventMetadata`. +10. Feature, graph, and cache schema versions are explicit. +11. Cache compatibility is validated on load. +12. Cache round-trip preserves sample IDs and event metadata. +13. Fold selection is implemented using named metadata. +14. Split definitions are explicit and not hardcoded. +15. Deterministic shuffling works with an explicit seed. +16. Graph batching preserves graph/label/metadata/global alignment. +17. Required active padding modes are implemented. +18. A DataLoader produces a clean `GraphBatch`. +19. End-to-end preprocessing-to-batch integration passes. +20. Relevant legacy parity tests pass. +21. Full tests pass. +22. No model or training implementation has been added. +23. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. EventMetadata design +2. sample ID strategy +3. legacy tracking-to-metadata mapping +4. EventSample contract +5. GraphSample contract +6. GraphBatch contract +7. metadata batching rules +8. CacheMetadata design +9. schema-version strategy +10. cache format +11. cache provenance/fingerprint strategy +12. cache round-trip results +13. split/fold API +14. deterministic shuffle behavior +15. DataLoader API +16. padding modes implemented +17. unit test results +18. integration test results +19. parity results +20. intentional deviations from legacy representation +21. performance concerns deferred +22. unresolved questions +23. validation commands and results + +After validation succeeds, create one Git commit containing only Task 7 changes. + +Use: + +```text +feat: implement metadata-aware dataset caching and batching +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task8.md b/tasks/task8.md new file mode 100644 index 0000000000000000000000000000000000000000..78075c09c187d1dfeac4d34578ec232732846952 --- /dev/null +++ b/tasks/task8.md @@ -0,0 +1,1144 @@ +# Task 8: Implement the ROOT-GNN Model and Transfer/Fine-Tuning Path + +Implement the active ROOT-GNN neural-network architecture, including the transfer-learning / fine-tuning path required by the current workflow. + +This task builds on: + +```text +Task 4: shared physics feature construction +Task 5: ROOT-GNN graph construction +Task 6: ROOT/Awkward ingestion and EventSample +Task 7: metadata-aware caching, splits, batching, and GraphBatch +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/data/ +src/gnn4colliders/features/ +src/gnn4colliders/graphs/ +tests/unit/ +tests/parity/ +tests/integration/ +legacy/root_gnn_dgl/models/GCN.py +legacy/root_gnn_dgl/models/loss.py +legacy/root_gnn_dgl/scripts/training_script.py +``` + +Focus on the active legacy model classes: + +```text +models.GCN.Edge_Network +models.GCN.Transferred_Learning_Finetuning +``` + +Do not migrate the entire legacy `GCN.py` file. + +Do not implement the training lifecycle, optimizer setup, loss functions, metrics, checkpoint orchestration, or DDP in this task. + +--- + +# Goal + +Implement a clean ROOT-GNN model family under: + +```text +src/gnn4colliders/models/root_gnn/ +``` + +that reproduces the active legacy forward behavior and supports: + +```text +1. training a ROOT-GNN model from scratch +2. loading a pretrained ROOT-GNN backbone +3. removing/replacing the pretrained classifier +4. fine-tuning a new classifier/head +5. optionally freezing or unfreezing the transferred backbone +``` + +The intended model-level flow is: + +```text +GraphBatch + -> +ROOT-GNN encoders + -> +message passing + -> +graph representation + -> +classifier + -> +logits +``` + +For transfer learning: + +```text +pretrained ROOT-GNN + -> +reuse learned backbone / representation + -> +replace classifier + -> +fine-tune target task +``` + +--- + +# 1. Package structure + +Prefer a small explicit structure such as: + +```text +src/gnn4colliders/models/root_gnn/ + __init__.py + blocks.py + edge_network.py + transfer.py +``` + +Possible responsibilities: + +```text +blocks.py + reusable MLP / encoder / update blocks + +edge_network.py + active ROOT-GNN architecture + +transfer.py + transfer-learning / fine-tuning model +``` + +Use fewer modules if clearer. + +Do not recreate the legacy giant `GCN.py`. + +--- + +# 2. Model naming + +Use modern class names in new code. + +Prefer: + +```python +EdgeNetwork +TransferredLearningFinetuning +``` + +or clearer names if appropriate, for example: + +```python +EdgeNetwork +FineTunedEdgeNetwork +``` + +Do not preserve awkward underscore-heavy legacy class names unless needed only for a compatibility adapter. + +Document any renamed public concepts. + +--- + +# 3. Active EdgeNetwork behavior + +Inspect the legacy `Edge_Network` implementation carefully. + +Reproduce the active architecture, including where applicable: + +* node encoder +* edge encoder +* global-feature encoder +* iterative message-passing blocks +* edge updates +* node aggregation +* node updates +* global pooling +* global updates +* decoder +* final classifier +* representation output if used downstream + +Do not assume behavior from the architecture document alone. + +Validate against the actual legacy implementation. + +--- + +# 4. Message-passing order + +Preserve the active update order. + +The architecture documentation suggests a repeated pattern approximately: + +```text +encode + -> +edge update + -> +aggregate edges into nodes + -> +node update + -> +pool nodes/edges + -> +global update +``` + +Verify the exact legacy behavior. + +Add tests that make the ordering observable where practical. + +Do not redesign the algorithm merely because another message-passing convention seems cleaner. + +--- + +# 5. Reusable MLP block + +Reimplement the active MLP-building behavior cleanly. + +The legacy path includes combinations of: + +```text +Linear +ReLU +Dropout +LayerNorm +``` + +Verify: + +* order +* number of layers +* activation placement +* dropout placement +* LayerNorm placement +* bias behavior +* initialization behavior where relevant + +Prefer a reusable typed helper in: + +```text +blocks.py +``` + +Do not add a general neural-network framework. + +--- + +# 6. Input contract + +The ROOT-GNN model should consume the graph representation produced by the new pipeline. + +Prefer a model API conceptually like: + +```python +logits = model( + graph, + global_features, +) +``` + +or, if cleaner: + +```python +logits = model(batch) +``` + +where `batch` is the Task 7 `GraphBatch`. + +Choose the API that keeps the model easy to test and avoids unnecessary coupling. + +If the model accepts `GraphBatch`, avoid making it depend on unrelated metadata such as fold or weight. + +Event weights are not model inputs unless the legacy model truly uses them. + +--- + +# 7. Node and edge feature keys + +If Task 5 preserved compatibility keys such as: + +```python +graph.ndata["features"] +graph.edata["features"] +``` + +the model may consume those directly. + +Do not rename graph keys in this task unless there is a strong reason. + +If named constants already exist for graph feature keys, use them. + +Avoid scattering magic strings. + +--- + +# 8. Global features + +Preserve the active global-feature behavior. + +Characterize: + +* whether globals are optional +* expected shape +* expected dtype +* how globals are encoded +* behavior when no globals are configured +* batching semantics + +Do not invent dummy global features unless the legacy path requires them. + +If the model supports both: + +```text +with globals +without globals +``` + +test both where active. + +--- + +# 9. Model outputs + +The model must return raw logits. + +Do not apply: + +```text +sigmoid +softmax +``` + +inside the model unless the legacy model itself does so. + +Expected output shape: + +```text +[batch_size, out_size] +``` + +Verify: + +```text +out_size = 1 +``` + +for binary tasks and the active multiclass output size where needed. + +Do not couple model output logic to loss functions. + +--- + +# 10. Representation interface + +The transfer-learning path depends on the learned representation before the final classifier. + +Expose this representation cleanly. + +Prefer one of: + +```python +representation = model.encode(...) +logits = model.classify(representation) +``` + +or: + +```python +representation = model.forward_features(...) +logits = model.forward(...) +``` + +The exact names may differ. + +The important goal is to avoid transfer learning by mutating arbitrary internal module lists. + +Make the backbone/classifier boundary explicit. + +--- + +# 11. Separate backbone and classifier conceptually + +Structure `EdgeNetwork` so that the learned representation and classifier are separable. + +Conceptually: + +```text +graph + -> +backbone + -> +representation + -> +classifier + -> +logits +``` + +The backbone may include: + +* encoders +* message passing +* global decoder + +The classifier should be a clearly identifiable final prediction head. + +This separation is required for clean transfer learning. + +Do not change numerical behavior merely to achieve the separation. + +--- + +# 12. Transfer-learning behavior + +Implement the active fine-tuning path represented by the legacy: + +```text +Transferred_Learning_Finetuning +``` + +Characterize exactly what it does. + +At minimum support: + +```text +pretrained EdgeNetwork + -> +load pretrained weights + -> +remove/replace final classifier + -> +attach target-task classifier + -> +forward through transferred representation +``` + +Do not blindly copy legacy state mutation. + +Provide an explicit new API. + +--- + +# 13. Transfer model API + +Prefer something conceptually like: + +```python +model = FineTunedEdgeNetwork.from_pretrained( + pretrained_model, + out_size=1, + freeze_backbone=True, +) +``` + +or: + +```python +model = FineTunedEdgeNetwork( + backbone=pretrained_model.backbone, + classifier=..., +) +``` + +The exact API may differ. + +Requirements: + +* clear ownership of pretrained backbone +* explicit target output size +* explicit freeze/unfreeze behavior +* testable parameter state + +Do not hide transfer behavior inside config side effects. + +--- + +# 14. Loading pretrained models + +Separate: + +```text +model architecture +``` + +from: + +```text +checkpoint file loading +``` + +Task 11 will implement full checkpoint orchestration. + +For this task, it is acceptable to support loading a plain compatible `state_dict` for parity tests. + +Do not implement: + +* epoch discovery +* optimizer-state restoration +* early-stopping restoration +* "best checkpoint" selection +* resume-training lifecycle + +Those belong later. + +If legacy checkpoints are needed for testing, implement the smallest compatibility helper necessary and keep it isolated. + +--- + +# 15. Legacy checkpoint compatibility boundary + +The existing legacy checkpoints store model state using historical naming. + +If needed to prove transfer-learning parity, add a narrowly scoped compatibility function such as: + +```python +load_legacy_edge_network_state_dict(...) +``` + +This helper may: + +* remove `module.` prefixes +* handle `_orig_mod.` prefixes if necessary +* map clearly known renamed module keys + +Do not build the complete checkpoint system yet. + +Keep compatibility code separate from the clean model implementation. + +--- + +# 16. Freeze / unfreeze behavior + +Support explicit control over transferred parameters. + +At minimum: + +```text +freeze_backbone = True +freeze_backbone = False +``` + +or equivalent. + +Add tests verifying `requires_grad`. + +For example: + +```python +assert not any(p.requires_grad for p in model.backbone.parameters()) +assert all(p.requires_grad for p in model.classifier.parameters()) +``` + +for a frozen-backbone configuration. + +Do not implicitly freeze parameters without documenting it. + +--- + +# 17. Legacy fine-tuning semantics + +Inspect exactly which parts of the legacy model are frozen/reset/reinitialized. + +The architecture analysis noted legacy fine-tuning behavior and hardcoded random seeding. + +Characterize: + +* which layers are reused +* which layers are removed +* which layers are reinitialized +* whether the backbone is frozen +* whether some backbone layers remain trainable +* initialization behavior of the new classifier + +Preserve active behavior when it affects compatibility. + +Where legacy behavior is accidental or unclear, document it rather than silently embedding it forever. + +--- + +# 18. Random initialization + +Do not mutate process-global random state inside model constructors. + +Do not reproduce hardcoded calls such as: + +```python +torch.manual_seed(2) +``` + +inside model classes. + +Instead, if deterministic initialization is needed, provide an explicit mechanism controlled by the caller. + +For example: + +```python +reset_parameters(generator=...) +``` + +or rely on externally seeded PyTorch initialization. + +Record this as an intentional internal improvement if it differs from legacy implementation mechanics. + +Parity tests should compare fixed parameters/weights where exact forward parity is needed rather than depending on global initialization side effects. + +--- + +# 19. Forward-pass parity strategy + +Do not attempt to compare two randomly initialized models. + +For strong parity testing: + +1. instantiate the legacy model +2. instantiate the new model with matching architecture +3. map/copy weights from legacy to new +4. pass identical graph/global inputs +5. compare: + + * intermediate representation where practical + * final logits + +This should make forward parity meaningful. + +Use small deterministic graphs from existing fixtures. + +--- + +# 20. State-dict mapping + +If the clean module structure changes parameter names, implement a test-only or compatibility mapping between legacy and new state dictionaries. + +Keep the mapping explicit and documented. + +Do not force the new package architecture to copy legacy module names merely to avoid writing a compatibility mapping. + +Test that all expected active parameters are accounted for. + +Fail clearly if unmapped active parameters remain. + +--- + +# 21. Unit tests for blocks + +Add focused tests under: + +```text +tests/unit/models/root_gnn/ +``` + +Suggested files: + +```text +test_blocks.py +test_edge_network.py +test_transfer.py +``` + +Test reusable blocks independently where worthwhile. + +Avoid testing PyTorch itself. + +Focus on project-specific structure and behavior. + +--- + +# 22. EdgeNetwork unit tests + +Cover at minimum: + +```text +single graph +batched graphs +out_size = 1 +multiclass out_size +global features if active +multiple processing steps +output dtype +output shape +``` + +Verify output is raw logits. + +--- + +# 23. Intermediate representation tests + +Test the backbone / representation API. + +Verify: + +```text +representation.shape +``` + +and consistency between: + +```python +representation = model.forward_features(...) +logits = model.classifier(representation) +``` + +and: + +```python +logits = model(...) +``` + +where appropriate. + +--- + +# 24. Transfer tests + +Test: + +```text +pretrained backbone reuse +new classifier output dimension +frozen backbone +unfrozen backbone +classifier trainability +forward output shape +``` + +Also verify that replacing the classifier does not modify pretrained backbone weights. + +--- + +# 25. Transfer forward parity + +Where practical, compare the new fine-tuning model against the active legacy transfer-learning class. + +Use: + +* equivalent pretrained weights +* equivalent new classifier weights +* identical graph inputs +* identical global inputs + +Compare target-task logits. + +Do not depend on training to establish parity. + +--- + +# 26. Binary task coverage + +Add direct coverage for the active binary fine-tuning case: + +```text +out_size = 1 +``` + +Verify: + +```text +logits.shape == [batch_size, 1] +``` + +Do not apply sigmoid in the model. + +--- + +# 27. Multiclass pretraining coverage + +Add direct coverage for the active multiclass pretrained model. + +Use the active output size from the current standard configuration. + +If the standard path uses 12 classes, test: + +```text +out_size = 12 +``` + +Verify output shape and parity. + +--- + +# 28. DGL mutation + +The legacy model may mutate graph node/edge data in place during forward. + +Avoid unnecessary persistent mutation where practical. + +If DGL message-passing operations require temporary graph state, prefer using: + +```python +with graph.local_scope(): + ... +``` + +or equivalent. + +This prevents model forward calls from leaking temporary node/edge fields into the caller's graph. + +Preserve final numerical behavior. + +Add a test verifying repeated forward calls on the same graph behave consistently. + +--- + +# 29. Device behavior + +Support normal PyTorch device behavior. + +Do not hardcode CUDA. + +The model should work on CPU for unit tests. + +It should move correctly under: + +```python +model.to(device) +graph.to(device) +``` + +Do not add special Perlmutter behavior. + +GPU/HPC execution belongs to later tasks. + +--- + +# 30. Dtype behavior + +Preserve expected float dtype behavior. + +Do not silently cast to a different precision. + +Do not introduce mixed precision yet. + +Mixed precision belongs to the training/performance stage. + +--- + +# 31. No loss implementation + +Do not implement: + +```text +BCEWithLogitsLoss +weighted per-label averaging +ROC AUC +accuracy +finish functions +``` + +Those belong to Task 9. + +The model should only produce representations/logits. + +--- + +# 32. No training loop + +Do not implement: + +```text +optimizer +scheduler +epochs +backward +gradient clipping +early stopping +training logs +``` + +Those belong to Task 10. + +--- + +# 33. No full checkpoint orchestration + +Do not implement: + +```text +model_epoch_N.pt +best epoch selection +last epoch selection +optimizer checkpoint state +early-stop state +resume training +``` + +Those belong to Task 11. + +A narrowly scoped state-dict compatibility helper is allowed only if necessary for Task 8 transfer parity. + +--- + +# 34. No Hydra integration yet + +Do not wire model construction directly to Hydra/YAML. + +Use explicit Python constructor arguments and typed config objects only if clearly useful. + +Hydra wiring comes later. + +--- + +# 35. No experimental architectures + +Do not migrate: + +```text +GCN_global +GCN_global_2way +attention variants +MultiModel +Clustering +unused transfer variants +``` + +unless an active standard config proves they are required. + +Focus on: + +```text +Edge_Network +Transferred_Learning_Finetuning +``` + +only. + +--- + +# 36. Integration with GraphBatch + +Add a small integration test: + +```text +tiny ROOT fixture + -> +EventSample + -> +features + -> +GraphSample + -> +GraphBatch + -> +EdgeNetwork + -> +logits +``` + +Verify output shape and deterministic execution. + +Then add transfer integration: + +```text +pretrained EdgeNetwork + -> +FineTunedEdgeNetwork + -> +same GraphBatch + -> +target-task logits +``` + +Do not train the model. + +--- + +# 37. Documentation + +Update: + +```text +docs/architecture.md +``` + +to describe: + +```text +EdgeNetwork + active ROOT-GNN model + +backbone / representation + learned reusable graph representation + +classifier + task-specific output head + +FineTunedEdgeNetwork + pretrained backbone + replacement classifier +``` + +Document the model boundary: + +```text +GraphBatch + -> +ROOT-GNN backbone + -> +representation + -> +task-specific classifier +``` + +Update: + +```text +docs/migration.md +``` + +when parity is established. + +Document intentional improvements such as avoiding constructor-level global RNG mutation or avoiding persistent DGL graph mutation. + +Keep documentation concise. + +--- + +# 38. Public exports + +Expose the intended public model APIs from: + +```python +gnn4colliders.models.root_gnn +``` + +For example: + +```python +from gnn4colliders.models.root_gnn import ( + EdgeNetwork, + FineTunedEdgeNetwork, +) +``` + +Do not export internal helper classes unnecessarily. + +--- + +# Validation + +Run focused model tests: + +```bash +uv run pytest tests/unit/models/root_gnn -v +``` + +Run integration tests: + +```bash +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format validation: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy source is unchanged +* no loss implementation was added +* no optimizer/training lifecycle was added +* no full checkpoint system was added +* no experimental legacy model families were migrated +* transfer learning is explicit and testable +* no hidden global seed mutation exists in new model code +* no persistent temporary DGL graph mutation leaks from forward +* no unrelated changes are included + +--- + +# Completion criteria + +Task 8 is complete when: + +1. `EdgeNetwork` is implemented under `src/gnn4colliders/models/root_gnn/`. +2. Active node/edge/global encoders match required legacy behavior. +3. Message-passing order matches the active legacy model. +4. The reusable learned representation/backbone boundary is explicit. +5. The final classifier is explicit and replaceable. +6. Raw logits are returned with the correct shape. +7. Binary `out_size=1` is covered. +8. Active multiclass pretraining output is covered. +9. A transfer/fine-tuning model is implemented. +10. A pretrained backbone can be reused with a replacement classifier. +11. Freeze/unfreeze behavior is explicit and tested. +12. Fine-tuning forward behavior matches the active legacy implementation where parity is required. +13. Fixed-weight forward-pass parity with legacy `Edge_Network` passes. +14. Legacy state-dict compatibility needed for this task is isolated from the clean model code. +15. Repeated forward calls do not leak unwanted graph mutation. +16. Unit tests pass. +17. Integration tests pass. +18. Parity tests pass. +19. Full tests pass. +20. No loss/training/checkpoint lifecycle implementation has been added. +21. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. `EdgeNetwork` public API +4. backbone/representation API +5. classifier boundary +6. message-passing architecture implemented +7. global-feature handling +8. binary/multiclass output behavior +9. transfer-learning API +10. legacy transfer behavior reproduced +11. backbone freeze/unfreeze behavior +12. state-dict mapping/compatibility strategy +13. intentional differences from legacy internals +14. RNG behavior +15. DGL mutation handling +16. unit test results +17. integration test results +18. forward parity results +19. transfer parity results +20. numerical tolerances used +21. ambiguities/deferred behavior +22. validation commands and results + +After validation succeeds, create one Git commit containing only Task 8 changes. + +Use: + +```text +feat: implement ROOT-GNN model and transfer learning +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tasks/task9.md b/tasks/task9.md new file mode 100644 index 0000000000000000000000000000000000000000..f9759758abe35f835517322bdb4acd7d87b5114a --- /dev/null +++ b/tasks/task9.md @@ -0,0 +1,1074 @@ +# Task 9: Implement Classification Tasks, Weighted Losses, Metrics, and Output Handling + +Implement the task-level logic that converts model logits and batch metadata into losses, predictions, and evaluation metrics. + +This task builds on: + +```text +Task 4: shared physics feature construction +Task 5: ROOT-GNN graph construction +Task 6: ROOT/Awkward ingestion and EventSample +Task 7: metadata-aware caching, splits, batching, and GraphBatch +Task 8: ROOT-GNN model and transfer/fine-tuning +``` + +Before making changes, read: + +```text +AGENTS.md +docs/architecture.md +docs/migration.md +``` + +Then inspect: + +```text +src/gnn4colliders/models/ +src/gnn4colliders/training/ +tests/unit/ +tests/parity/ +tests/integration/ +legacy/root_gnn_dgl/models/loss.py +legacy/root_gnn_dgl/scripts/training_script.py +``` + +Focus on the active classification behavior used by standard ROOT-GNN training and evaluation. + +Do not implement the epoch/training lifecycle, optimizer setup, scheduler, checkpoint orchestration, or DDP in this task. + +--- + +# Goal + +Implement explicit task semantics around: + +```text +GraphBatch + -> +Model + -> +logits + -> +classification task + ├── loss + ├── predictions + └── metrics +``` + +The model should remain responsible only for producing logits. + +The task layer should be responsible for interpreting those logits for a particular learning problem. + +Support the active: + +```text +binary classification +multiclass classification +``` + +workflows, including the transfer/fine-tuning use case. + +--- + +# 1. Introduce a task layer + +Add a small task abstraction. + +Prefer a structure such as: + +```text +src/gnn4colliders/tasks/ + __init__.py + binary_classification.py + multiclass_classification.py +``` + +or a similarly small organization. + +Do not bury classification semantics inside the future trainer. + +The trainer should eventually be able to do something conceptually like: + +```python +logits = model(batch) +loss = task.loss(logits, batch) +metrics = task.metrics(logits, batch) +``` + +The exact API may differ. + +Keep it simple. + +--- + +# 2. Keep training mechanics separate from task semantics + +The task layer should own things like: + +```text +how logits are interpreted +how targets are shaped +how event weights are applied +how predictions are produced +which metrics are computed +``` + +The future trainer should own things like: + +```text +forward pass +backward pass +optimizer step +epochs +logging +checkpointing +``` + +Do not mix those responsibilities. + +--- + +# 3. Binary classification task + +Implement the active binary-classification behavior. + +Expected model output: + +```text +[batch_size, 1] +``` + +The model returns raw logits. + +The task should handle any reshaping needed for loss/metrics in an explicit way. + +Do not apply sigmoid before `BCEWithLogitsLoss`. + +For prediction/metrics, apply sigmoid where appropriate. + +--- + +# 4. Binary targets + +Characterize the active binary target representation from the legacy workflow. + +Verify: + +* dtype +* shape +* expected values +* whether labels arrive as `[B]` or `[B, 1]` + +Do not silently rely on PyTorch broadcasting. + +Normalize target shape explicitly inside the task layer where appropriate. + +Add tests for target-shape handling. + +--- + +# 5. Legacy binary weighted loss + +Reproduce the active legacy binary-loss behavior. + +The architecture analysis indicates that the legacy code approximately does: + +```text +elementwise BCEWithLogitsLoss + -> +multiply by event weights + -> +average separately per unique label + -> +average across labels +``` + +Verify the exact implementation against: + +```text +legacy/root_gnn_dgl/scripts/training_script.py +legacy/root_gnn_dgl/models/loss.py +``` + +Do not assume ordinary weighted BCE is equivalent. + +Implement the actual active semantics. + +--- + +# 6. Event weights use named metadata + +The new implementation must obtain event weights through named metadata. + +For example: + +```python +batch.metadata.weight +``` + +or the equivalent Task 7 API. + +Do not reintroduce positional legacy access such as: + +```python +tracking[:, 1] +``` + +Add parity tests proving: + +```text +legacy tracking weight + == +new EventMetadata.weight +``` + +at the loss boundary. + +--- + +# 7. Negative weights + +Characterize the active legacy behavior for negative weights. + +The legacy CLI historically has an `--abs` path that makes weights positive. + +Do not make weights absolute by default unless that is the active task configuration. + +The task/loss API should make this choice explicit. + +For example, conceptually: + +```python +BinaryClassificationTask( + use_absolute_weights=False, +) +``` + +or an equivalent option. + +Avoid hiding this behavior in dataset processing. + +Add tests covering: + +```text +positive weights +negative weights +absolute-weight option +``` + +if the legacy active path supports them. + +--- + +# 8. Per-class balancing semantics + +If the legacy loss averages weighted loss separately for each unique label and then averages across labels, implement that explicitly. + +For binary classification, this likely means something conceptually like: + +```text +loss(label == 0) +loss(label == 1) +mean of class losses +``` + +Verify the precise denominator/weight handling. + +Do not replace it with: + +```python +torch.mean(weight * bce) +``` + +unless parity proves they are equivalent for the active implementation. + +--- + +# 9. Multiclass classification task + +Implement the active multiclass behavior. + +The active pretraining model uses a multiclass output, historically with 12 outputs in the standard configuration. + +Do not hardcode 12 into generic task logic. + +Use: + +```python +num_classes +``` + +or infer from logits/configuration where appropriate. + +Expected logits: + +```text +[batch_size, num_classes] +``` + +--- + +# 10. Multiclass target representation + +Characterize the active legacy multiclass target representation carefully. + +Determine whether the active path uses: + +```text +class indices +one-hot / multilabel-like target tensors +``` + +and which loss is actually used. + +Do not assume CrossEntropyLoss merely because the task is multiclass. + +The architecture analysis indicates the legacy default objective uses `BCEWithLogitsLoss`, including for the multiclass path. + +Verify this directly. + +Reproduce active behavior rather than replacing it with a more conventional formulation in this task. + +If the current active multiclass labels are one-hot and use BCE-with-logits, preserve that for parity. + +Document it as a compatibility behavior if it is unusual. + +--- + +# 11. Multiclass weighted loss + +Implement the active multiclass weighting/reduction behavior. + +Verify: + +* elementwise loss shape +* event-weight broadcast behavior +* per-label or per-class reduction +* final reduction + +Use explicit operations rather than relying on ambiguous broadcasting. + +Add small hand-computable tests. + +--- + +# 12. Hand-computable loss tests + +For both binary and multiclass behavior, add tests with tiny fixed tensors where expected loss can be calculated explicitly. + +For example: + +```python +logits = ... +targets = ... +weights = ... +``` + +and compare against a manually constructed expected result. + +These tests should not merely call the same helper function twice. + +The purpose is to lock down reduction semantics. + +--- + +# 13. Binary prediction semantics + +Implement binary predictions according to the active legacy behavior. + +The architecture analysis indicates: + +```text +sigmoid(logit) >= 0.5 +``` + +or equivalent threshold behavior. + +Verify exact comparison semantics: + +```text +> +>= +``` + +if it matters. + +Expose prediction utilities clearly. + +For example: + +```python +probabilities = task.probabilities(logits) +predictions = task.predict(logits) +``` + +The exact API may differ. + +--- + +# 14. Binary probabilities + +Provide access to sigmoid probabilities separately from hard predictions. + +Do not mutate logits. + +Keep: + +```text +logits +probabilities +predictions +``` + +conceptually distinct. + +This will be useful later for inference and ROC calculations. + +--- + +# 15. Multiclass prediction semantics + +Implement active multiclass prediction behavior. + +The architecture indicates predictions use: + +```text +argmax +``` + +Verify the exact dimension. + +If probability-like outputs are needed, determine whether legacy evaluation uses: + +```text +sigmoid per output +softmax +raw logits +``` + +for ROC AUC. + +Do not assume. + +Characterize the actual active implementation. + +--- + +# 16. Accuracy + +Implement active accuracy metrics. + +For binary classification: + +* use the characterized threshold +* preserve weighting semantics if legacy accuracy is weighted + +For multiclass classification: + +* use argmax +* preserve weighting semantics if legacy accuracy is weighted + +Verify whether the legacy metric is: + +```text +ordinary accuracy +weighted accuracy +``` + +and reproduce the active path. + +--- + +# 17. ROC AUC + +Implement active ROC AUC behavior using scikit-learn or an equivalent already-approved dependency. + +For binary classification, verify: + +* probability/score input +* sample weights +* handling of one-class batches/datasets + +For multiclass classification, verify: + +* one-vs-rest behavior +* score representation +* sample weights +* averaging mode + +The architecture analysis indicates weighted ROC AUC and one-vs-rest multiclass behavior, but verify the exact implementation. + +--- + +# 18. Metrics are evaluation-level, not necessarily batch-safe + +Some metrics such as ROC AUC are not meaningful or robust on every individual mini-batch. + +Design metric APIs so they can work on accumulated epoch-level predictions/targets/weights later. + +Do not force the future trainer to compute ROC AUC independently on every batch. + +Prefer metric functions that operate on: + +```text +all logits/scores for a split +all targets +all weights +``` + +This will integrate cleanly with Task 10. + +--- + +# 19. Metric inputs + +Prefer explicit metric input structures. + +A lightweight structure is acceptable, for example: + +```python +@dataclass(frozen=True) +class ClassificationOutputs: + logits: torch.Tensor + targets: torch.Tensor + weights: torch.Tensor +``` + +Do not create abstractions unless they improve clarity. + +Avoid coupling metrics directly to DGL graphs. + +--- + +# 20. CPU conversion for sklearn metrics + +If scikit-learn is used, isolate conversion from PyTorch tensors. + +For example: + +```text +tensor.detach() + -> +cpu() + -> +numpy() +``` + +Do not accidentally keep autograd graphs alive during metric computation. + +Keep conversion logic in one clear place. + +--- + +# 21. Finish/postprocessing functions + +Inspect active `finish` functions in: + +```text +legacy/root_gnn_dgl/models/loss.py +``` + +Determine which ones are actually used by the active standard configurations. + +Only implement the output transformations required by active: + +```text +binary +multiclass +fine-tuning +``` + +workflows. + +Do not migrate every historical finish function. + +If output transformation is just: + +```text +sigmoid +argmax +``` + +prefer clear task methods over carrying forward arbitrary legacy "finish function" terminology. + +--- + +# 22. Task API + +Prefer a small interface. + +Conceptually: + +```python +class BinaryClassificationTask: + def loss(self, logits, targets, weights): ... + def probabilities(self, logits): ... + def predict(self, logits): ... + def metrics(self, logits, targets, weights): ... +``` + +and: + +```python +class MulticlassClassificationTask: + ... +``` + +The exact API can use functions or dataclasses instead of classes if simpler. + +Do not introduce a deep abstract base-class hierarchy. + +Only add a shared protocol/base abstraction if it reduces real duplication. + +--- + +# 23. Batch-facing convenience API + +It is acceptable to provide convenience methods that consume `GraphBatch`: + +```python +loss = task.loss_from_batch(logits, batch) +``` + +if this keeps the future trainer clean. + +However, keep core loss functions testable using plain tensors. + +Do not make the mathematical implementation depend on DGL. + +--- + +# 24. Weight dtype/device + +Ensure event weights are compatible with logits. + +Handle: + +* device +* dtype +* shape + +explicitly. + +Do not silently move the entire batch between CPU/GPU inside loss functions. + +Small metadata tensors may be moved/cast as needed in a controlled way, but the behavior should be clear and tested. + +--- + +# 25. Empty classes / missing labels + +Characterize behavior when a loss-reduction batch contains no examples for one class. + +This matters if the legacy implementation loops over unique labels. + +Do not invent a divide-by-zero behavior. + +Implement the observed semantics and add a test if this can occur. + +--- + +# 26. Degenerate ROC AUC inputs + +Handle cases where ROC AUC cannot be computed because only one class is present. + +Choose behavior consistent with the active legacy workflow where possible. + +Prefer an explicit: + +```text +NaN +None +clear exception +``` + +over a misleading numeric value. + +Document the chosen behavior. + +Do not hide metric failures silently. + +--- + +# 27. No model changes unless required for interface cleanup + +Do not change Task 8 model mathematics. + +Model output remains raw logits. + +Only make minimal interface changes if necessary to integrate with the task layer. + +Any model change must preserve Task 8 parity. + +--- + +# 28. No training lifecycle + +Do not implement: + +```text +optimizer +backward +zero_grad +optimizer.step +epochs +schedulers +gradient clipping +early stopping +training loops +``` + +Those belong to Task 10. + +--- + +# 29. No checkpointing + +Do not implement: + +```text +checkpoint save +checkpoint resume +best epoch +last epoch +optimizer-state restoration +``` + +Those belong to Task 11. + +--- + +# 30. No Hydra wiring yet + +Do not implement full Hydra task configuration yet. + +Use explicit Python task constructors/options. + +Hydra/YAML composition will be wired in a later task. + +--- + +# 31. No distributed metrics yet + +Do not implement distributed all-gather/reduction. + +Task 10/14 will address distributed execution. + +Metric APIs should be designed so accumulated outputs can later be gathered across ranks. + +--- + +# 32. Unit tests + +Add focused tests under: + +```text +tests/unit/tasks/ +``` + +Suggested files: + +```text +test_binary_classification.py +test_multiclass_classification.py +test_metrics.py +``` + +or a smaller clear organization. + +Cover: + +```text +binary loss +binary weighting +negative weights +absolute-weight option +binary thresholding +binary probabilities +multiclass loss +multiclass predictions +accuracy +ROC AUC +shape validation +dtype/device handling +``` + +--- + +# 33. Legacy parity tests + +Add parity tests comparing new task logic to the active legacy implementations. + +Use deterministic fixed: + +```text +logits +targets +weights +``` + +rather than running a full training loop. + +Compare: + +* loss +* probabilities/scores +* predictions +* accuracy +* ROC AUC + +where deterministic and practical. + +--- + +# 34. Fine-tuning task parity + +Explicitly cover the active binary fine-tuning case produced by: + +```text +FineTunedEdgeNetwork +``` + +Use deterministic logits or a deterministic model fixture. + +Verify: + +```text +FineTunedEdgeNetwork + -> +raw binary logits + -> +BinaryClassificationTask + -> +weighted loss / metrics +``` + +No optimizer or backward pass is needed. + +--- + +# 35. Multiclass pretraining task parity + +Explicitly cover the active pretraining output. + +Verify: + +```text +EdgeNetwork(out_size=active_num_classes) + -> +raw logits + -> +MulticlassClassificationTask +``` + +matches legacy loss/prediction/metric semantics. + +--- + +# 36. Integration test + +Add an integration test covering: + +```text +tiny ROOT fixture + -> +GraphBatch + -> +EdgeNetwork or FineTunedEdgeNetwork + -> +logits + -> +task loss + -> +task predictions +``` + +Verify: + +* finite loss +* expected shape +* metadata weight use +* no positional tracking dependency + +Do not perform backward/optimizer steps yet. + +--- + +# 37. Documentation + +Update: + +```text +docs/architecture.md +``` + +to make the new boundary explicit: + +```text +model + raw logits only + +task + target semantics + weighting + loss + probabilities + predictions + metrics + +trainer + future execution lifecycle +``` + +Document that event weights come from: + +```text +EventMetadata.weight / BatchMetadata.weight +``` + +rather than legacy tracking columns. + +If the active multiclass implementation intentionally retains unusual BCE-with-logits semantics, document that as a compatibility behavior. + +Update: + +```text +docs/migration.md +``` + +when Task 9 parity is established. + +--- + +# 38. Public exports + +Expose the intended task APIs. + +For example: + +```python +from gnn4colliders.tasks import ( + BinaryClassificationTask, + MulticlassClassificationTask, +) +``` + +Do not expose internal helper functions unnecessarily. + +--- + +# Validation + +Run focused tests: + +```bash +uv run pytest tests/unit/tasks -v +``` + +Run relevant model/integration tests: + +```bash +uv run pytest tests/unit/models/root_gnn -v +uv run pytest tests/integration -v +``` + +Run parity tests: + +```bash +uv run pytest tests/parity -v +``` + +Run the full suite: + +```bash +uv run pytest +``` + +Run lint/format checks: + +```bash +uv run ruff check src tests +uv run ruff format --check src tests +``` + +Inspect: + +```bash +git status +git diff +``` + +Verify: + +* legacy code is unchanged +* weights are accessed through named metadata +* no positional tracking API has returned +* model outputs remain raw logits +* no optimizer/training loop was added +* no checkpoint system was added +* no unrelated changes are included + +--- + +# Completion criteria + +Task 9 is complete when: + +1. A clear task layer exists. +2. Binary classification behavior is implemented. +3. Multiclass classification behavior is implemented. +4. Active legacy loss semantics are reproduced. +5. Event weights use named metadata. +6. Negative/absolute-weight semantics are explicit. +7. Binary probability and threshold behavior match legacy. +8. Multiclass prediction behavior matches legacy. +9. Accuracy matches active legacy semantics. +10. Binary weighted ROC AUC matches legacy. +11. Multiclass one-vs-rest ROC AUC matches legacy. +12. Task APIs can operate on plain tensors independently of DGL. +13. Fine-tuning task parity passes. +14. Multiclass pretraining task parity passes. +15. Unit tests pass. +16. Integration tests pass. +17. Parity tests pass. +18. Full tests pass. +19. No training lifecycle or checkpoint implementation has been added. +20. Legacy code remains untouched. + +--- + +# Completion report + +Report: + +1. files created +2. files modified +3. task API design +4. binary target semantics +5. binary weighted-loss formula +6. negative-weight handling +7. multiclass target semantics +8. multiclass loss formula +9. prediction/probability behavior +10. accuracy behavior +11. ROC AUC behavior +12. degenerate metric handling +13. legacy finish-function behavior retained or replaced +14. event-weight metadata integration +15. unit test results +16. integration test results +17. parity results +18. intentional deviations from legacy internals +19. numerical tolerances used +20. unresolved ambiguities +21. validation commands and results + +After validation succeeds, create one Git commit containing only Task 9 changes. + +Use: + +```text +feat: implement classification losses and metrics +``` + +Before committing, inspect the final diff and ensure no unrelated files are included. diff --git a/tests/fixtures/.gitkeep b/tests/fixtures/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/tests/fixtures/.gitkeep @@ -0,0 +1 @@ + diff --git a/tests/integration/.gitkeep b/tests/integration/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/tests/integration/.gitkeep @@ -0,0 +1 @@ + diff --git a/tests/integration/__init__.py b/tests/integration/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/integration/test_real_root_model.py b/tests/integration/test_real_root_model.py new file mode 100644 index 0000000000000000000000000000000000000000..0377a29579d7edbb95372a3074ba5ca040993482 --- /dev/null +++ b/tests/integration/test_real_root_model.py @@ -0,0 +1,111 @@ +"""ROOT-GNN model coverage over every entry in the 64-event ROOT fixture.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import awkward as ak +import pytest +import torch + +from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples +from gnn4colliders.features import build_node_features +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + +pytest.importorskip("dgl") + +pytestmark = [pytest.mark.integration, pytest.mark.real_data] +uproot = pytest.importorskip("uproot") + +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"] +FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1] + + +def _fixture_path() -> Path: + configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE") + return Path(configured) if configured else Path("data/processed/ttH_NLO_64.root") + + +@pytest.fixture(scope="module") +def root_graph_samples(): + path = _fixture_path() + if not path.exists(): + pytest.skip(f"ROOT sample fixture is absent: {path}") + with uproot.open(path) as root_file: + arrays = root_file["output"].arrays(entry_start=0, entry_stop=64, library="ak") + samples = [] + for index, event in enumerate(ak.Array(arrays)): + features, _ = build_node_features( + event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES + ) + samples.append( + GraphSample( + graph=build_dgl_graph(features), + label=torch.tensor(index % 2, dtype=torch.long), + global_features=torch.tensor( + [float(event["Number"]), float(event["weight"])], + dtype=torch.float32, + ), + metadata=EventMetadata( + fold=index % 4, + weight=float(event["weight"]), + sample_id=f"fixture:{index}", + ), + ) + ) + return samples + + +def test_all_64_real_root_graphs_run_as_one_graph_batch(root_graph_samples): + batch = batch_graph_samples(root_graph_samples) + model = EdgeNetwork( + root_graph_samples[0].graph, + torch.zeros(1, 2), + hid_size=16, + out_size=12, + n_layers=2, + n_proc_steps=2, + ).eval() + transferred = FineTunedEdgeNetwork.from_pretrained( + model, out_size=1, freeze_backbone=True + ).eval() + with torch.no_grad(): + logits = model(batch) + transfer_logits = transferred(batch) + assert batch.graph.batch_num_nodes().shape == (64,) + assert logits.shape == (64, 12) + assert transfer_logits.shape == (64, 1) + assert torch.isfinite(logits).all() + assert torch.isfinite(transfer_logits).all() + + +def test_real_root_graph_batch_matches_individual_forward(root_graph_samples): + model = EdgeNetwork( + root_graph_samples[0].graph, + torch.zeros(1, 2), + hid_size=8, + out_size=1, + n_layers=2, + n_proc_steps=1, + ).eval() + batch = batch_graph_samples(root_graph_samples) + with torch.no_grad(): + batched = model(batch) + individual = torch.cat( + [ + model(sample.graph, sample.global_features) + for sample in root_graph_samples + ] + ) + assert torch.allclose(batched, individual, atol=1e-5, rtol=1e-5) diff --git a/tests/integration/test_real_root_sample.py b/tests/integration/test_real_root_sample.py new file mode 100644 index 0000000000000000000000000000000000000000..409d70c114b98e5855d6ede63572f10b1fd79eaa --- /dev/null +++ b/tests/integration/test_real_root_sample.py @@ -0,0 +1,76 @@ +"""Smoke tests against a small fixture derived from the real ttH ROOT sample.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import awkward as ak +import pytest +import torch +import uproot + +from gnn4colliders.features import build_node_features +from gnn4colliders.graphs import build_dgl_graph + +pytest.importorskip("dgl") + +pytestmark = [pytest.mark.integration, pytest.mark.real_data] + +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"] +FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1] + + +def _fixture_path() -> Path: + configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE") + if configured: + return Path(configured) + return Path("data/processed/ttH_NLO_64.root") + + +@pytest.fixture(scope="module") +def root_tree(): + path = _fixture_path() + if not path.exists(): + pytest.skip( + f"ROOT sample fixture is absent: {path}; download ttH_NLO.root and " + "create the reduced fixture first" + ) + root_file = uproot.open(path) + tree = root_file["output"] + try: + yield tree + finally: + root_file.close() + + +def test_real_root_sample_has_expected_fixture_shape(root_tree): + assert root_tree.num_entries == 64 + assert {"jet_pt", "MET_met", "weight", "Number"}.issubset(root_tree.keys()) + + +def test_real_root_events_build_features_and_graphs(root_tree): + arrays = root_tree.arrays(entry_start=0, entry_stop=8, library="ak") + for event in ak.Array(arrays): + event_features, lengths = build_node_features( + event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES + ) + graph = build_dgl_graph(event_features) + + assert event_features.dtype == torch.float32 + assert event_features.shape == (sum(lengths), 7) + assert graph.number_of_nodes() == event_features.shape[0] + expected_edges = event_features.shape[0] * max(event_features.shape[0] - 1, 1) + assert graph.number_of_edges() == expected_edges + assert graph.edata["features"].shape == (expected_edges, 3) + assert torch.isfinite(event_features).all() + assert torch.isfinite(graph.edata["features"]).all() diff --git a/tests/integration/test_root_gnn_model.py b/tests/integration/test_root_gnn_model.py new file mode 100644 index 0000000000000000000000000000000000000000..043285c88cc8fbb8a1523f9b4fe12b01cd748ef8 --- /dev/null +++ b/tests/integration/test_root_gnn_model.py @@ -0,0 +1,23 @@ +import pytest +import torch + +from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + +pytest.importorskip("dgl") + + +def test_graph_batch_runs_through_pretraining_and_transfer_models(): + graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]])) + sample = GraphSample( + graph=graph, + label=torch.tensor(1), + global_features=torch.tensor([1.0, 2.0]), + metadata=EventMetadata(fold=0, weight=1.0, sample_id="fixture:0"), + ) + batch = batch_graph_samples([sample]) + model = EdgeNetwork(graph, torch.zeros(1, 2), 8, 12, 2, 1) + transferred = FineTunedEdgeNetwork.from_pretrained(model, 1) + assert model(batch).shape == (1, 12) + assert transferred(batch).shape == (1, 1) diff --git a/tests/parity/.gitkeep b/tests/parity/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/tests/parity/.gitkeep @@ -0,0 +1 @@ + diff --git a/tests/parity/conftest.py b/tests/parity/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..007ce7a9c4ec97cb57e0f86272378d93c115e43c --- /dev/null +++ b/tests/parity/conftest.py @@ -0,0 +1,92 @@ +"""Helpers for importing the legacy ROOT-GNN implementation in parity tests.""" + +import os +import sys +import types +from pathlib import Path + +import pytest + +pytestmark = pytest.mark.parity + +pytestmark = pytest.mark.parity + + +def _dgl_is_importable(): + try: + import dgl # noqa: F401 + except ImportError: + return False + return True + + +@pytest.fixture(scope="session", autouse=True) +def require_root_gnn_dependencies(): + if ( + os.environ.get("GNN4COLLIDERS_REQUIRE_ROOT_GNN") == "1" + and not _dgl_is_importable() + ): + pytest.fail( + "ROOT-GNN parity requires an importable DGL installation", + pytrace=False, + ) + + +@pytest.fixture(scope="session") +def legacy_dataset_module(): + """Return the active legacy dataset module when DGL is available.""" + pytest.importorskip("dgl") + legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl" + legacy_root_string = str(legacy_root) + if legacy_root_string not in sys.path: + sys.path.insert(0, legacy_root_string) + from root_gnn_base import dataset + + return dataset + + +@pytest.fixture(scope="session") +def legacy_model_module(): + """Return the legacy active model module for fixed-weight parity tests.""" + pytest.importorskip("dgl") + legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl" + legacy_root_string = str(legacy_root) + if legacy_root_string not in sys.path: + sys.path.insert(0, legacy_root_string) + from models import GCN + + return GCN + + +@pytest.fixture(scope="session") +def legacy_dataset_module_without_dgl(): + """Import pure legacy preprocessing with a minimal DGL import shim.""" + legacy_root = Path(__file__).parents[2] / "legacy" / "root_gnn_dgl" + legacy_root_string = str(legacy_root) + if legacy_root_string not in sys.path: + sys.path.insert(0, legacy_root_string) + try: + import dgl # noqa: F401 + except ImportError: + dgl_module = types.ModuleType("dgl") + dgl_data_module = types.ModuleType("dgl.data") + dgl_data_module.DGLDataset = type("DGLDataset", (), {}) + dgl_module.data = dgl_data_module + sys.modules["dgl"] = dgl_module + sys.modules["dgl.data"] = dgl_data_module + matplotlib_module = types.ModuleType("matplotlib") + matplotlib_pyplot_module = types.ModuleType("matplotlib.pyplot") + matplotlib_module.pyplot = matplotlib_pyplot_module + sys.modules["matplotlib"] = matplotlib_module + sys.modules["matplotlib.pyplot"] = matplotlib_pyplot_module + from root_gnn_base import dataset + + return dataset + + +@pytest.fixture(scope="session") +def legacy_utils_module(legacy_dataset_module_without_dgl): + """Return legacy fold-selection helpers after the package path is set up.""" + from root_gnn_base import utils + + return utils diff --git a/tests/parity/test_legacy_graph_and_edges.py b/tests/parity/test_legacy_graph_and_edges.py new file mode 100644 index 0000000000000000000000000000000000000000..ee50ed80b4eb0f6db9ba51cdd3a82e63ec2b0c90 --- /dev/null +++ b/tests/parity/test_legacy_graph_and_edges.py @@ -0,0 +1,112 @@ +"""Characterization tests for legacy directed graph and edge construction.""" + +import numpy as np +import torch + +from gnn4colliders.graphs import build_edge_features, fully_connected_edges + + +def test_full_connected_graph_is_directed_without_self_loops(legacy_dataset_module): + graph = legacy_dataset_module.full_connected_graph(3, self_loops=False) + assert graph.is_homogeneous + assert graph.number_of_nodes() == 3 + assert graph.number_of_edges() == 6 + sources, destinations = graph.edges() + np.testing.assert_array_equal(sources.numpy(), [0, 0, 1, 1, 2, 2]) + np.testing.assert_array_equal(destinations.numpy(), [1, 2, 0, 2, 0, 1]) + + +def test_full_connected_graph_keeps_self_loop_for_single_node( + legacy_dataset_module, +): + graph = legacy_dataset_module.full_connected_graph(1, self_loops=False) + assert graph.number_of_nodes() == 1 + assert graph.number_of_edges() == 1 + + +def test_full_connected_graph_edge_counts_and_order_are_source_major( + legacy_dataset_module, +): + expected_edges = { + 1: [(0, 0)], + 2: [(0, 1), (1, 0)], + 3: [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)], + } + for n_nodes, expected in expected_edges.items(): + graph = legacy_dataset_module.full_connected_graph(n_nodes, self_loops=False) + assert graph.number_of_edges() == len(expected) + sources, destinations = graph.edges() + assert list(zip(sources.tolist(), destinations.tolist())) == expected + + +def test_new_topology_matches_legacy(legacy_dataset_module): + for n_nodes in (1, 2, 3): + legacy_graph = legacy_dataset_module.full_connected_graph( + n_nodes, self_loops=False + ) + source, destination = fully_connected_edges(n_nodes) + legacy_source, legacy_destination = legacy_graph.edges() + torch.testing.assert_close(source, legacy_source) + torch.testing.assert_close(destination, legacy_destination) + + +def test_edge_features_are_deta_dphi_dr_with_phi_wrapping(legacy_dataset_module): + dataset = object.__new__(legacy_dataset_module.EdgeDataset) + dataset.times = [0, 0] + dataset.node_branch_names = [ + ["jet_pt"], + ["jet_eta"], + ["jet_phi"], + "CALC_E", + [0], + [0], + "NODE_TYPE", + ] + dataset.node_branch_types = ["vector"] + dataset.node_feature_scales = torch.ones(7) + event = { + "jet_pt": np.array([10.0, 20.0], dtype=np.float32), + "jet_eta": np.array([1.0, -0.5], dtype=np.float32), + "jet_phi": np.array([3.1, -3.1], dtype=np.float32), + } + graph = dataset.make_graph(event) + expected_deta = torch.tensor([1.5, -1.5]) + expected_dphi = torch.tensor([-0.0831853, 0.0831853]) + expected = torch.stack( + [expected_deta, expected_dphi, torch.sqrt(expected_deta**2 + expected_dphi**2)], + dim=1, + ) + assert graph.edata["features"].shape == (2, 3) + np.testing.assert_allclose( + graph.edata["features"].numpy(), expected.numpy(), rtol=0, atol=1e-5 + ) + + +def test_new_edge_features_match_legacy(legacy_dataset_module): + dataset = object.__new__(legacy_dataset_module.EdgeDataset) + dataset.times = [0, 0] + dataset.node_branch_names = [ + ["jet_pt"], + ["jet_eta"], + ["jet_phi"], + "CALC_E", + [0], + [0], + "NODE_TYPE", + ] + dataset.node_branch_types = ["vector"] + dataset.node_feature_scales = torch.ones(7) + event = { + "jet_pt": np.array([10.0, 20.0], dtype=np.float32), + "jet_eta": np.array([1.0, -0.5], dtype=np.float32), + "jet_phi": np.array([3.1, -3.1], dtype=np.float32), + } + legacy_graph = dataset.make_graph(event) + node_features = legacy_graph.ndata["features"] + source, destination = fully_connected_edges(node_features.shape[0]) + actual = build_edge_features( + node_features, source, destination, eta_index=1, phi_index=2 + ) + np.testing.assert_allclose( + actual.numpy(), legacy_graph.edata["features"].numpy(), rtol=0, atol=1e-6 + ) diff --git a/tests/parity/test_legacy_node_features.py b/tests/parity/test_legacy_node_features.py new file mode 100644 index 0000000000000000000000000000000000000000..9230fe2d7205947aa253a97d08f75560baa76385 --- /dev/null +++ b/tests/parity/test_legacy_node_features.py @@ -0,0 +1,100 @@ +"""Characterization tests for the active legacy node feature builder.""" + +import numpy as np +import pytest +import torch + + +@pytest.fixture +def event(): + return { + "jet_pt": np.array([100.0, 50.0], dtype=np.float32), + "ele_pt": np.array([20.0], dtype=np.float32), + "mu_pt": np.array([30.0], dtype=np.float32), + "ph_pt": np.array([40.0], dtype=np.float32), + "MET_met": np.float32(25.0), + "jet_eta": np.array([1.0, -0.5], dtype=np.float32), + "ele_eta": np.array([0.25], dtype=np.float32), + "mu_eta": np.array([-0.75], dtype=np.float32), + "ph_eta": np.array([0.5], dtype=np.float32), + "jet_phi": np.array([3.0, -3.0], dtype=np.float32), + "ele_phi": np.array([0.2], dtype=np.float32), + "mu_phi": np.array([-0.4], dtype=np.float32), + "ph_phi": np.array([1.0], dtype=np.float32), + "MET_phi": np.float32(-1.2), + "jet_btag": np.array([0.8, 0.1], dtype=np.float32), + "ele_charge": np.array([-1.0], dtype=np.float32), + "mu_charge": np.array([1.0], dtype=np.float32), + } + + +@pytest.fixture +def node_schema(): + return ( + [ + ["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", + ], + ["vector", "vector", "vector", "vector", "single"], + [0.1, 1, 1, 0.1, 1, 1, 1], + ) + + +def test_node_features_have_active_schema_and_order( + legacy_dataset_module_without_dgl, event, node_schema +): + names, types, scales = node_schema + features, lengths = legacy_dataset_module_without_dgl.node_features_from_tree( + event, names, types, torch.tensor(scales, dtype=torch.float32) + ) + expected = np.array( + [ + [10.0, 1.0, 3.0, 15.431, 0.8, 0.0, 0.0], + [5.0, -0.5, -3.0, 5.638, 0.1, 0.0, 0.0], + [2.0, 0.25, 0.2, 2.063, 0.0, -1.0, 1.0], + [3.0, -0.75, -0.4, 3.884, 0.0, 1.0, 2.0], + [4.0, 0.5, 1.0, 4.511, 0.0, 0.0, 3.0], + [2.5, 0.0, -1.2, 2.500, 0.0, 0.0, 4.0], + ], + dtype=np.float32, + ) + assert lengths == [2, 1, 1, 1, 1] + assert features.shape == (6, 7) + assert features.dtype == torch.float32 + np.testing.assert_allclose(features.numpy(), expected, rtol=0, atol=2e-3) + + +def test_node_feature_builder_preserves_zero_length_vector( + legacy_dataset_module_without_dgl, node_schema +): + names, types, scales = node_schema + event = { + "jet_pt": np.array([], dtype=np.float32), + "ele_pt": np.array([20.0], dtype=np.float32), + "mu_pt": np.array([], dtype=np.float32), + "ph_pt": np.array([], dtype=np.float32), + "MET_met": np.float32(25.0), + "jet_btag": np.array([], dtype=np.float32), + "ele_charge": np.array([-1.0], dtype=np.float32), + "mu_charge": np.array([], dtype=np.float32), + "jet_eta": np.array([], dtype=np.float32), + "ele_eta": np.array([0.25], dtype=np.float32), + "mu_eta": np.array([], dtype=np.float32), + "ph_eta": np.array([], dtype=np.float32), + "jet_phi": np.array([], dtype=np.float32), + "ele_phi": np.array([0.2], dtype=np.float32), + "mu_phi": np.array([], dtype=np.float32), + "ph_phi": np.array([], dtype=np.float32), + "MET_phi": np.float32(-1.2), + } + features, lengths = legacy_dataset_module_without_dgl.node_features_from_tree( + event, names, types, torch.tensor(scales) + ) + assert lengths == [0, 1, 0, 0, 1] + assert features.shape == (2, 7) + assert features[:, 6].tolist() == [1.0, 4.0] diff --git a/tests/parity/test_legacy_tracking.py b/tests/parity/test_legacy_tracking.py new file mode 100644 index 0000000000000000000000000000000000000000..06dff934841519bf42e140cc1ccb471053dca37d --- /dev/null +++ b/tests/parity/test_legacy_tracking.py @@ -0,0 +1,77 @@ +"""Characterization of the legacy dataset item contract.""" + +import torch + + +def _dataset_item(legacy_dataset_module_without_dgl, labels, tracking, globals_): + dataset = object.__new__(legacy_dataset_module_without_dgl.RootDataset) + dataset.graphs = [object() for _ in labels] + dataset.labels = torch.as_tensor(labels) + dataset.tracking = torch.as_tensor(tracking, dtype=torch.float32) + dataset.global_features = torch.as_tensor(globals_, dtype=torch.float32) + return dataset + + +def test_dataset_item_contract_preserves_label_tracking_and_global_shapes( + legacy_dataset_module_without_dgl, +): + dataset = object.__new__(legacy_dataset_module_without_dgl.RootDataset) + graph = object() + dataset.graphs = [graph] + dataset.labels = torch.tensor([7]) + dataset.tracking = torch.tensor([[3.0, 2.5]], dtype=torch.float32) + dataset.global_features = torch.empty((1, 0), dtype=torch.float32) + item = dataset[0] + assert item[0] is graph + assert item[1].shape == torch.Size([]) + assert item[1].dtype == torch.int64 + assert item[2].shape == (2,) + assert item[2].dtype == torch.float32 + assert item[2].tolist() == [3.0, 2.5] + assert item[3].shape == (0,) + assert item[3].dtype == torch.float32 + + +def test_dataset_item_contract_preserves_binary_and_multiclass_labels( + legacy_dataset_module_without_dgl, +): + dataset = _dataset_item( + legacy_dataset_module_without_dgl, + labels=[0, 1, 11], + tracking=[[0.0, 1.0], [3.0, -2.5], [7.0, 0.0]], + globals_=[[], [], []], + ) + items = [dataset[index] for index in range(3)] + assert [item[1].item() for item in items] == [0, 1, 11] + assert all(item[1].shape == torch.Size([]) for item in items) + assert all(item[1].dtype == torch.int64 for item in items) + assert items[1][2].tolist() == [3.0, -2.5] + assert all(item[3].shape == (0,) for item in items) + + +def test_dataset_item_contract_preserves_nonempty_global_features( + legacy_dataset_module_without_dgl, +): + dataset = _dataset_item( + legacy_dataset_module_without_dgl, + labels=[1], + tracking=[[4.0, 2.0]], + globals_=[[12.5, -3.0]], + ) + item = dataset[0] + assert item[3].shape == (2,) + assert item[3].dtype == torch.float32 + assert item[3].tolist() == [12.5, -3.0] + + +def test_fold_selection_uses_tracking_column_zero_and_keeps_weights( + legacy_utils_module, +): + sample = type( + "Sample", + (), + {"tracking": torch.tensor([[0.0, 1.0], [1.0, -2.5], [2.0, 0.0], [3.0, 4.0]])}, + )() + selection = legacy_utils_module.fold_selection({"n_folds": 2, "test": [1]}, "test") + assert selection(sample).tolist() == [False, True, False, True] + assert sample.tracking[:, 1].tolist() == [1.0, -2.5, 0.0, 4.0] diff --git a/tests/parity/test_model_parity.py b/tests/parity/test_model_parity.py new file mode 100644 index 0000000000000000000000000000000000000000..aa501f2cf0304be61acfc8dd6a7854b9516e524a --- /dev/null +++ b/tests/parity/test_model_parity.py @@ -0,0 +1,110 @@ +"""Fixed-weight parity tests for the active legacy and rewritten models.""" + +from __future__ import annotations + +import pytest +import torch + +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import ( + EdgeNetwork, + FineTunedEdgeNetwork, + load_legacy_edge_network_state_dict, +) + +pytest.importorskip("dgl") + + +def _graph_and_globals(): + graph = build_dgl_graph( + torch.tensor( + [[10.0, 0.2, 0.1], [8.0, -0.3, -0.2], [4.0, 0.1, 2.8]], + dtype=torch.float32, + ) + ) + return graph, torch.tensor([[1.0, 2.0]], dtype=torch.float32) + + +def _new_model(graph, globals_): + return EdgeNetwork(graph, globals_, 8, 3, 2, 2, dropout=0.0).eval() + + +def test_edge_network_matches_legacy_forward_and_representation( + legacy_model_module, +): + graph, globals_ = _graph_and_globals() + new_model = _new_model(graph, globals_) + legacy_model = legacy_model_module.Edge_Network( + graph, globals_, 8, 3, 2, 2, dropout=0.0 + ).eval() + legacy_state = { + key.replace("classifier.", "classify."): value + for key, value in new_model.state_dict().items() + } + legacy_model.load_state_dict(legacy_state) + + with torch.no_grad(): + new_representation = new_model.forward_features(graph, globals_) + legacy_representation = legacy_model.representation(graph.clone(), globals_)[1] + new_logits = new_model(graph, globals_) + legacy_logits = legacy_model(graph.clone(), globals_) + assert torch.equal(new_representation, legacy_representation) + assert torch.equal(new_logits, legacy_logits) + + +def test_transfer_forward_matches_legacy_transfer(legacy_model_module, tmp_path): + graph, _ = _graph_and_globals() + globals_ = torch.empty((0, 0), dtype=torch.float32) + new_model = EdgeNetwork(graph, globals_, 8, 3, 2, 2, dropout=0.0).eval() + legacy_pretrained = legacy_model_module.Edge_Network( + graph, globals_, 8, 3, 2, 2, dropout=0.0 + ).eval() + legacy_state = { + key.replace("classifier.", "classify."): value + for key, value in new_model.state_dict().items() + } + legacy_pretrained.load_state_dict(legacy_state) + checkpoint_path = tmp_path / "pretrained.pt" + torch.save({"model_state_dict": legacy_pretrained.state_dict()}, checkpoint_path) + + config = { + "module": "models.GCN", + "class": "Edge_Network", + "args": { + "hid_size": 8, + "out_size": 3, + "n_layers": 2, + "n_proc_steps": 2, + "dropout": 0.0, + }, + } + legacy_transfer = legacy_model_module.Transferred_Learning_Finetuning( + str(checkpoint_path), config, graph, globals_, 8, 1, 2, 2, dropout=0.0 + ).eval() + transfer = FineTunedEdgeNetwork.from_pretrained(new_model, 1) + transfer.classifier.load_state_dict(legacy_transfer.classify.state_dict()) + + with torch.no_grad(): + expected = legacy_transfer(graph.clone(), None) + actual = transfer(graph, None) + assert torch.equal(actual, expected) + assert torch.equal( + transfer.representation(graph, None), transfer.forward_features(graph, None) + ) + + +def test_legacy_prefix_loader_handles_checkpoint_prefixes(): + graph, globals_ = _graph_and_globals() + source = _new_model(graph, globals_) + target = _new_model(graph, globals_) + prefixed = { + "module._orig_mod." + key.replace("classifier.", "classify."): value + for key, value in source.state_dict().items() + } + load_legacy_edge_network_state_dict(target, {"model_state_dict": prefixed}) + for actual, expected in zip(target.parameters(), source.parameters()): + assert torch.equal(actual, expected) + transfer = FineTunedEdgeNetwork.from_pretrained( + source, 1, state_dict={"model_state_dict": prefixed} + ) + assert isinstance(transfer.backbone.classifier, torch.nn.Identity) diff --git a/tests/parity/test_new_node_features.py b/tests/parity/test_new_node_features.py new file mode 100644 index 0000000000000000000000000000000000000000..9726594b72d74f2d2236db2a489a95fc42e43ee2 --- /dev/null +++ b/tests/parity/test_new_node_features.py @@ -0,0 +1,45 @@ +import numpy as np +import torch + +from gnn4colliders.features import build_node_features + + +def test_new_builder_matches_legacy_builder(legacy_dataset_module_without_dgl): + event = { + "jet_pt": np.array([100.0, 50.0], dtype=np.float32), + "ele_pt": np.array([20.0], dtype=np.float32), + "mu_pt": np.array([30.0], dtype=np.float32), + "ph_pt": np.array([40.0], dtype=np.float32), + "MET_met": np.float32(25.0), + "jet_eta": np.array([1.0, -0.5], dtype=np.float32), + "ele_eta": np.array([0.25], dtype=np.float32), + "mu_eta": np.array([-0.75], dtype=np.float32), + "ph_eta": np.array([0.5], dtype=np.float32), + "jet_phi": np.array([3.0, -3.0], dtype=np.float32), + "ele_phi": np.array([0.2], dtype=np.float32), + "mu_phi": np.array([-0.4], dtype=np.float32), + "ph_phi": np.array([1.0], dtype=np.float32), + "MET_phi": np.float32(-1.2), + "jet_btag": np.array([0.8, 0.1], dtype=np.float32), + "ele_charge": np.array([-1.0], dtype=np.float32), + "mu_charge": np.array([1.0], dtype=np.float32), + } + names = [ + ["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 = torch.tensor([0.1, 1, 1, 0.1, 1, 1, 1]) + expected, expected_lengths = ( + legacy_dataset_module_without_dgl.node_features_from_tree( + event, names, object_types, scales + ) + ) + actual, actual_lengths = build_node_features(event, names, object_types, scales) + assert actual_lengths == expected_lengths + torch.testing.assert_close(actual, expected, rtol=0, atol=0) diff --git a/tests/unit/.gitkeep b/tests/unit/.gitkeep new file mode 100644 index 0000000000000000000000000000000000000000..8b137891791fe96927ad78e64b0aad7bded08bdc --- /dev/null +++ b/tests/unit/.gitkeep @@ -0,0 +1 @@ + diff --git a/tests/unit/compat/__init__.py b/tests/unit/compat/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/unit/compat/test_adapters.py b/tests/unit/compat/test_adapters.py new file mode 100644 index 0000000000000000000000000000000000000000..e359dc8493491f191b1e246cc9a9dee5ff34fcb6 --- /dev/null +++ b/tests/unit/compat/test_adapters.py @@ -0,0 +1,29 @@ +import pytest +import torch + +from gnn4colliders.compat import ( + event_metadata_from_legacy_tracking, + map_legacy_edge_network_state_dict, +) +from gnn4colliders.data import EventMetadata + + +def test_legacy_tracking_is_converted_to_named_metadata(): + metadata = event_metadata_from_legacy_tracking( + [3.0, -2.5, 99.0], sample_id="event:7" + ) + assert metadata == EventMetadata(fold=3, weight=-2.5, sample_id="event:7") + + +def test_legacy_tracking_rejects_ambiguous_short_rows(): + with pytest.raises(ValueError, match="fold and weight"): + EventMetadata.from_legacy_tracking([3.0], sample_id="event:7") + + +def test_legacy_state_dict_mapping_is_canonical_and_deterministic(): + state = { + "module._orig_mod.encoder.weight": torch.ones(2, 2), + "module._orig_mod.classify.bias": torch.zeros(1), + } + mapped = map_legacy_edge_network_state_dict(state) + assert set(mapped) == {"encoder.weight", "classifier.bias"} diff --git a/tests/unit/config/test_factories.py b/tests/unit/config/test_factories.py new file mode 100644 index 0000000000000000000000000000000000000000..7a7a7bee796af7ea559ccb31ed1adef28c16fb7a --- /dev/null +++ b/tests/unit/config/test_factories.py @@ -0,0 +1,18 @@ +import pytest + +from gnn4colliders.config import build_task + + +def test_factory_builds_semantic_task_names(): + from gnn4colliders.tasks import ( + BinaryClassificationTask, + MulticlassClassificationTask, + ) + + assert isinstance(build_task({"type": "binary"}), BinaryClassificationTask) + assert isinstance(build_task({"type": "multiclass"}), MulticlassClassificationTask) + + +def test_factory_rejects_unknown_task_name(): + with pytest.raises(ValueError, match="unsupported task"): + build_task({"type": "not-a-task"}) diff --git a/tests/unit/config/test_validation.py b/tests/unit/config/test_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..a31601bc442f01e4e6b17b366fa6763ea9b359c5 --- /dev/null +++ b/tests/unit/config/test_validation.py @@ -0,0 +1,38 @@ +import pytest + +from gnn4colliders.config import validate_config + + +@pytest.mark.parametrize( + ("section", "values", "message"), + [ + ("data", {"batch_size": 0}, "batch_size"), + ("trainer", {"max_epochs": 0}, "max_epochs"), + ("model", {"out_size": 2}, "out_size"), + ], +) +def test_invalid_cross_config_values_fail_before_initialization( + section, values, message +): + config = {"data": {}, "trainer": {}, "model": {}, "task": {"type": "binary"}} + config[section].update(values) + with pytest.raises(ValueError, match=message): + validate_config(config) + + +def test_validation_rejects_conflicting_resume_and_transfer(): + config = {"checkpoint": {"resume": "resume.pt", "pretrained": "base.pt"}} + with pytest.raises(ValueError, match="distinct workflows"): + validate_config(config) + + +def test_validation_rejects_overlapping_split_folds(): + config = {"data": {"splits": {"train_folds": [0], "test_folds": [0]}}} + with pytest.raises(ValueError, match="disjoint"): + validate_config(config) + + +def test_fine_tuning_requires_pretrained_checkpoint(): + config = {"model": {"name": "fine_tuned_edge_network"}} + with pytest.raises(ValueError, match="pretrained"): + validate_config(config) diff --git a/tests/unit/data/test_batching.py b/tests/unit/data/test_batching.py new file mode 100644 index 0000000000000000000000000000000000000000..cd47c8ab04a60b17fb8849975a211f957f85688f --- /dev/null +++ b/tests/unit/data/test_batching.py @@ -0,0 +1,33 @@ +import pytest +import torch + +from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples +from gnn4colliders.graphs import build_dgl_graph + +dgl = pytest.importorskip("dgl") + + +def _sample(index: int, node_count: int) -> GraphSample: + nodes = torch.arange(node_count * 3, dtype=torch.float32).reshape(node_count, 3) + return GraphSample( + graph=build_dgl_graph(nodes), + label=torch.tensor(index), + global_features=torch.tensor([float(index), -float(index)]), + metadata=EventMetadata(index, index + 0.5, f"sample_{index}"), + ) + + +@pytest.mark.parametrize("counts", [(1,), (1, 2), (1, 3, 2)]) +def test_graph_batch_preserves_heterogeneous_order_and_metadata(counts): + samples = [_sample(index, count) for index, count in enumerate(counts)] + batch = batch_graph_samples(samples) + assert batch.labels.tolist() == list(range(len(counts))) + assert batch.metadata.sample_id == tuple(f"sample_{i}" for i in range(len(counts))) + assert batch.metadata.weight.tolist() == pytest.approx( + [i + 0.5 for i in range(len(counts))] + ) + assert batch.global_features[:, 0].tolist() == list(map(float, range(len(counts)))) + assert batch.graph.batch_num_nodes().tolist() == list(counts) + assert batch.graph.batch_num_edges().tolist() == [ + count * (count - 1) if count > 1 else 1 for count in counts + ] diff --git a/tests/unit/data/test_orchestration.py b/tests/unit/data/test_orchestration.py new file mode 100644 index 0000000000000000000000000000000000000000..eff193f4f59610050614816a618558fc96ad4d76 --- /dev/null +++ b/tests/unit/data/test_orchestration.py @@ -0,0 +1,134 @@ +from pathlib import Path + +import numpy as np +import pytest + +from gnn4colliders.data import ( + FEATURE_SCHEMA_VERSION, + BatchMetadata, + EventMetadata, + SplitDefinition, + select_folds, +) +from gnn4colliders.data.cache import CacheMetadata, GraphSampleCache +from gnn4colliders.data.graph_dataset import GraphDataLoader, GraphDataset, GraphSample + + +def _sample(sample_id: str, fold: int) -> GraphSample: + return GraphSample( + graph=object(), + label=np.asarray(fold), + global_features=None, + metadata=EventMetadata( + fold=fold, + weight=-0.5 if fold else 1.0, + sample_id=sample_id, + extra={"source_file": "events.root"}, + ), + ) + + +def test_select_folds_is_named_deterministic_and_non_mutating(): + samples = [_sample("a", 1), _sample("b", 0), _sample("c", 1)] + selected = select_folds(samples, {1}) + assert [sample.metadata.sample_id for sample in selected] == ["a", "c"] + assert [sample.metadata.sample_id for sample in samples] == ["a", "b", "c"] + + +def test_batch_metadata_preserves_order_and_extra_fields(): + events = [ + EventMetadata(2, 0.0, "sample_10", {"run": 10}), + EventMetadata(0, -3.5, "sample_2", {"run": 2, "tag": "b"}), + ] + batch = BatchMetadata.from_events(events) + assert batch.sample_id == ("sample_10", "sample_2") + assert batch.fold.tolist() == [2, 0] + assert batch.weight.tolist() == [0.0, -3.5] + assert batch.extra == {"run": (10, 2), "tag": (None, "b")} + + +def test_graph_loader_shuffle_is_seeded_by_epoch_without_global_rng_mutation(): + pytest.importorskip("dgl") + import torch + + from gnn4colliders.graphs import build_dgl_graph + + samples = [ + GraphSample( + build_dgl_graph(torch.ones(1, 3)), + torch.tensor(i), + None, + EventMetadata(0, 1.0, str(i)), + ) + for i in range(5) + ] + loader = GraphDataLoader(GraphDataset(samples), 2, shuffle=True, seed=11) + first = [item for batch in loader for item in batch.metadata.sample_id] + second = [item for batch in loader for item in batch.metadata.sample_id] + loader.set_epoch(1) + third = [item for batch in loader for item in batch.metadata.sample_id] + assert first == second + assert third != first + + +def test_split_definition_rejects_overlapping_folds(): + with pytest.raises(ValueError, match="disjoint"): + SplitDefinition(train_folds=frozenset({0}), test_folds=frozenset({0})) + + +def test_graph_cache_round_trip_and_schema_check(tmp_path: Path): + path = tmp_path / "graphs.pt" + cache = GraphSampleCache(path) + samples = (_sample("event.root:Events:3", 2),) + cache.save(samples) + loaded = cache.load() + assert loaded[0].metadata.sample_id == samples[0].metadata.sample_id + assert loaded[0].metadata.weight == -0.5 + + stale = GraphSampleCache( + path, CacheMetadata(feature_schema_version=FEATURE_SCHEMA_VERSION + 1) + ) + with pytest.raises(ValueError, match="schema mismatch"): + stale.load() + + +@pytest.mark.parametrize( + "payload, message", + [ + ({"samples": ()}, "schema mismatch"), + ({"cache_metadata": {}, "samples": ()}, "schema mismatch"), + ( + { + "cache_metadata": { + "feature_schema_version": 1, + "graph_schema_version": 1, + "cache_schema_version": 1, + } + }, + "samples", + ), + ], +) +def test_cache_rejects_missing_or_incomplete_metadata(tmp_path, payload, message): + path = tmp_path / "broken.pt" + import torch + + torch.save(payload, path) + with pytest.raises(ValueError, match=message): + GraphSampleCache(path).load() + + +def test_cache_rejects_preprocessing_fingerprint_mismatch(tmp_path): + path = tmp_path / "graphs.pt" + GraphSampleCache(path, CacheMetadata(preprocessing={"features": "a"})).save( + [_sample("event:0", 0)] + ) + with pytest.raises(ValueError, match="fingerprint"): + GraphSampleCache(path, CacheMetadata(preprocessing={"features": "b"})).load() + + +def test_cache_rejects_truncated_artifact(tmp_path): + path = tmp_path / "truncated.pt" + path.write_bytes(b"not a torch archive") + with pytest.raises(Exception): + GraphSampleCache(path).load() diff --git a/tests/unit/data/test_root_dataset.py b/tests/unit/data/test_root_dataset.py new file mode 100644 index 0000000000000000000000000000000000000000..f32e2a4e621449a522ce1a800dc19e7822dd2ce9 --- /dev/null +++ b/tests/unit/data/test_root_dataset.py @@ -0,0 +1,125 @@ +from pathlib import Path + +import awkward as ak +import numpy as np +import pytest +import uproot + +from gnn4colliders.data import EventSample, RootEventDataset +from gnn4colliders.data.root_io import branch_names_from_specs, read_tree +from gnn4colliders.features import build_object_features + + +def _write_fixture(path: Path) -> None: + with uproot.recreate(path) as root_file: + root_file["events"] = { + "eventNumber": np.array([10, 11, 12], dtype=np.int64), + "weight": np.array([1.5, -2.0, 0.25], dtype=np.float32), + "label_branch": np.array([4, 5, 6], dtype=np.int32), + "jet_pt": ak.Array([[10.0, 20.0], [], [30.0]]), + "global_x": np.array([0.5, 0.6, 0.7], dtype=np.float32), + } + + +def test_branch_selection_and_root_reading(tmp_path: Path): + path = tmp_path / "events.root" + _write_fixture(path) + assert branch_names_from_specs( + [["jet_pt"], "CALC_E", "NODE_TYPE"], + tracking_info=["eventNumber"], + global_features=["global_x"], + ) == ("jet_pt", "eventNumber", "global_x") + arrays = read_tree(path, "events", ["weight"]) + assert arrays.fields == ["weight"] + assert arrays["weight"].to_list() == [1.5, -2.0, 0.25] + + +def test_dataset_returns_architecture_neutral_samples(tmp_path: Path): + path = tmp_path / "events.root" + _write_fixture(path) + dataset = RootEventDataset( + path, + tree_name="events", + label="label_branch", + feature_branches=[["jet_pt"]], + global_features=["global_x"], + fold_var="eventNumber", + weight_var="weight", + ) + + assert len(dataset) == 3 + assert dataset.branches == ( + "jet_pt", + "eventNumber", + "weight", + "global_x", + "label_branch", + ) + sample = dataset[1] + assert isinstance(sample, EventSample) + assert sample.event_index == 1 + assert sample.objects["jet_pt"].to_list() == [] + assert sample.label == 5 + assert sample.tracking.tolist() == [11.0, -2.0] + assert sample.global_features.tolist() == pytest.approx([0.6]) + assert dataset[-1].event_index == 2 + + +def test_multiple_files_keep_input_order_and_feed_feature_builder(tmp_path: Path): + first = tmp_path / "first.root" + second = tmp_path / "second.root" + _write_fixture(first) + _write_fixture(second) + dataset = RootEventDataset( + [first, second], + tree_name="events", + label=[0, 1], + feature_branches=[["jet_pt"], [0], [0], ["CALC_E"]], + fold_var="eventNumber", + ) + + assert len(dataset) == 6 + assert dataset[3].event_index == 3 + assert dataset[3].label == 1 + features, lengths = build_object_features( + dataset[0].objects, + [["jet_pt"], [0], [0], ["CALC_E"]], + ["vector"], + [1, 1, 1, 1], + ) + assert features.shape == (2, 4) + assert lengths == [2] + + +def test_root_boundary_reports_missing_tree_and_branch(tmp_path: Path): + path = tmp_path / "events.root" + _write_fixture(path) + with pytest.raises(KeyError, match="tree"): + read_tree(path, "missing", ["weight"]) + with pytest.raises(KeyError, match="missing requested branches"): + read_tree(path, "events", ["missing_branch"]) + + +def test_root_boundary_reports_missing_file(tmp_path: Path): + with pytest.raises((OSError, FileNotFoundError)): + read_tree(tmp_path / "does-not-exist.root", "events", ["weight"]) + + +def test_sample_identity_is_stable_and_distinguishes_files(tmp_path: Path): + first = tmp_path / "first.root" + second = tmp_path / "second.root" + _write_fixture(first) + _write_fixture(second) + kwargs = dict( + tree_name="events", + label="label_branch", + feature_branches=[["jet_pt"]], + fold_var="eventNumber", + weight_var="weight", + ) + left = RootEventDataset(first, **kwargs) + repeated = RootEventDataset(first, **kwargs) + both = RootEventDataset([first, second], **kwargs) + assert left[0].metadata.sample_id == repeated[0].metadata.sample_id + assert left[0].metadata.sample_id != left[1].metadata.sample_id + assert both[0].metadata.sample_id != both[3].metadata.sample_id diff --git a/tests/unit/distributed/__init__.py b/tests/unit/distributed/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/unit/distributed/test_ddp_cpu.py b/tests/unit/distributed/test_ddp_cpu.py new file mode 100644 index 0000000000000000000000000000000000000000..40087dc882dac84ef344092f35ed0d1d280914c7 --- /dev/null +++ b/tests/unit/distributed/test_ddp_cpu.py @@ -0,0 +1,78 @@ +import multiprocessing as mp +import os +import socket + +import pytest +import torch +from torch import nn + +from gnn4colliders.distributed import finalize, initialize, prepare_model + +pytestmark = pytest.mark.distributed + + +def _port() -> int: + with socket.socket() as sock: + sock.bind(("127.0.0.1", 0)) + return sock.getsockname()[1] + + +def _worker(rank: int, port: int, queue) -> None: + os.environ.update( + RANK=str(rank), + LOCAL_RANK=str(rank), + WORLD_SIZE="2", + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + ) + context = initialize(enabled=True, backend="gloo", device="cpu") + try: + torch.manual_seed(5) + model = prepare_model(nn.Linear(1, 1), context) + optimizer = torch.optim.SGD(model.parameters(), lr=0.1) + x = torch.tensor([[1.0], [2.0]]) if rank == 0 else torch.tensor([[3.0], [4.0]]) + y = 2 * x + optimizer.zero_grad() + loss = nn.functional.mse_loss(model(x), y) + loss.backward() + optimizer.step() + queue.put( + ( + rank, + model.module.weight.detach().cpu().item(), + model.module.bias.detach().cpu().item(), + ) + ) + finally: + finalize(context) + + +@pytest.mark.skipif( + not torch.distributed.is_available(), reason="torch.distributed unavailable" +) +def test_two_rank_cpu_ddp_matches_single_process_update(): + queue = mp.get_context("spawn").Queue() + # Both workers must share one rendezvous port. + port = _port() + processes = [ + mp.get_context("spawn").Process(target=_worker, args=(rank, port, queue)) + for rank in range(2) + ] + for process in processes: + process.start() + values = [queue.get(timeout=30) for _ in processes] + for process in processes: + process.join(timeout=30) + assert process.exitcode == 0 + assert values[0][1] == values[1][1] + assert values[0][2] == values[1][2] + + torch.manual_seed(5) + reference = nn.Linear(1, 1) + optimizer = torch.optim.SGD(reference.parameters(), lr=0.1) + x = torch.arange(1.0, 5.0).reshape(-1, 1) + optimizer.zero_grad() + nn.functional.mse_loss(reference(x), 2 * x).backward() + optimizer.step() + assert torch.allclose(torch.tensor(values[0][1]), reference.weight.squeeze()) + assert torch.allclose(torch.tensor(values[0][2]), reference.bias.squeeze()) diff --git a/tests/unit/export/test_onnx_export.py b/tests/unit/export/test_onnx_export.py new file mode 100644 index 0000000000000000000000000000000000000000..7ca99483060d9c9a9d0136f132fb79fbfd1c6e6b --- /dev/null +++ b/tests/unit/export/test_onnx_export.py @@ -0,0 +1,63 @@ +import pytest +import torch + +from gnn4colliders.data import EventMetadata, GraphSample, batch_graph_samples +from gnn4colliders.export import ( + RootGNNExportAdapter, + export_root_gnn_onnx, + inputs_from_graph_batch, +) +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + +pytestmark = pytest.mark.onnx +dgl = pytest.importorskip("dgl") + + +def _batch(): + samples = [] + for index, count in enumerate((2, 3)): + nodes = torch.arange(count * 3, dtype=torch.float32).reshape(count, 3) + samples.append( + GraphSample( + build_dgl_graph(nodes), + torch.tensor(index), + torch.tensor([1.0, 2.0]), + EventMetadata(index, 1.0, f"fixture:{index}"), + ) + ) + return batch_graph_samples(samples) + + +def test_export_inputs_preserve_graph_membership_and_edges(): + batch = _batch() + inputs = inputs_from_graph_batch(batch) + src, _ = batch.graph.edges(order="eid") + assert torch.equal(inputs.node_features, batch.graph.ndata["features"]) + assert torch.equal(inputs.edge_features, batch.graph.edata["features"]) + assert torch.equal(inputs.edge_src, src) + assert inputs.node_batch.tolist() == [0, 0, 1, 1, 1] + assert inputs.global_features.shape == (2, 2) + + +@pytest.mark.parametrize("fine_tuned", [False, True]) +def test_tensor_adapter_matches_native_model(fine_tuned): + batch = _batch() + model = EdgeNetwork(batch.graph, batch.global_features, 8, 3, 2, 2).eval() + if fine_tuned: + model = FineTunedEdgeNetwork.from_pretrained(model, 1).eval() + inputs = inputs_from_graph_batch(batch) + with torch.inference_mode(): + native = model(batch.graph, batch.global_features) + adapted = RootGNNExportAdapter(model)(*inputs.as_tuple()) + torch.testing.assert_close(adapted, native, rtol=1e-5, atol=1e-5) + + +def test_onnx_export_is_optional_and_writes_validated_model(tmp_path): + pytest.importorskip("onnx") + pytest.importorskip("onnxruntime") + batch = _batch() + model = EdgeNetwork(batch.graph, batch.global_features, 8, 3, 2, 2).eval() + output = export_root_gnn_onnx(model, batch, tmp_path / "model.onnx") + assert output.exists() + assert output.with_suffix(".onnx.json").exists() diff --git a/tests/unit/features/__init__.py b/tests/unit/features/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..0259f722e9f8058fd4b6602def1e6a5a6e379205 --- /dev/null +++ b/tests/unit/features/__init__.py @@ -0,0 +1 @@ +"""Unit tests for shared collider feature construction.""" diff --git a/tests/unit/features/test_objects.py b/tests/unit/features/test_objects.py new file mode 100644 index 0000000000000000000000000000000000000000..53895bc59a2102ccd915c30595d16d87331840c6 --- /dev/null +++ b/tests/unit/features/test_objects.py @@ -0,0 +1,110 @@ +import numpy as np +import pytest +import torch + +from gnn4colliders.features import NODE_FEATURE_NAMES, build_node_features + + +@pytest.fixture +def schema(): + return ( + [ + ["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", + ], + ["vector", "vector", "vector", "vector", "single"], + [0.1, 1, 1, 0.1, 1, 1, 1], + ) + + +@pytest.fixture +def event(): + return { + "jet_pt": np.array([100.0, 50.0], dtype=np.float32), + "ele_pt": np.array([20.0], dtype=np.float32), + "mu_pt": np.array([30.0], dtype=np.float32), + "ph_pt": np.array([40.0], dtype=np.float32), + "MET_met": np.float32(25.0), + "jet_eta": np.array([1.0, -0.5], dtype=np.float32), + "ele_eta": np.array([0.25], dtype=np.float32), + "mu_eta": np.array([-0.75], dtype=np.float32), + "ph_eta": np.array([0.5], dtype=np.float32), + "jet_phi": np.array([3.0, -3.0], dtype=np.float32), + "ele_phi": np.array([0.2], dtype=np.float32), + "mu_phi": np.array([-0.4], dtype=np.float32), + "ph_phi": np.array([1.0], dtype=np.float32), + "MET_phi": np.float32(-1.2), + "jet_btag": np.array([0.8, 0.1], dtype=np.float32), + "ele_charge": np.array([-1.0], dtype=np.float32), + "mu_charge": np.array([1.0], dtype=np.float32), + } + + +def test_schema_and_values(event, schema): + names, object_types, scales = schema + features, lengths = build_node_features(event, names, object_types, scales) + assert NODE_FEATURE_NAMES == ( + "pt", + "eta", + "phi", + "energy", + "btag", + "charge", + "node_type", + ) + assert lengths == [2, 1, 1, 1, 1] + assert features.shape == (6, 7) + assert features.dtype == torch.float32 + np.testing.assert_allclose( + features.numpy(), + [ + [10.0, 1.0, 3.0, 15.431, 0.8, 0.0, 0.0], + [5.0, -0.5, -3.0, 5.638, 0.1, 0.0, 0.0], + [2.0, 0.25, 0.2, 2.063, 0.0, -1.0, 1.0], + [3.0, -0.75, -0.4, 3.884, 0.0, 1.0, 2.0], + [4.0, 0.5, 1.0, 4.511, 0.0, 0.0, 3.0], + [2.5, 0.0, -1.2, 2.500, 0.0, 0.0, 4.0], + ], + rtol=0, + atol=2e-3, + ) + + +def test_empty_vectors_and_input_immutability(event, schema): + names, object_types, scales = schema + event = dict(event) + event["jet_pt"] = np.array([], dtype=np.float32) + event["jet_eta"] = np.array([], dtype=np.float32) + event["jet_phi"] = np.array([], dtype=np.float32) + event["jet_btag"] = np.array([], dtype=np.float32) + before = event["ele_pt"].copy() + features, lengths = build_node_features(event, names, object_types, scales) + assert lengths == [0, 1, 1, 1, 1] + assert features.shape == (4, 7) + np.testing.assert_array_equal(event["ele_pt"], before) + + +def test_reordering_vector_objects_reorders_feature_rows(event, schema): + names, object_types, scales = schema + reordered = dict(event) + for name in ("jet_pt", "jet_eta", "jet_phi", "jet_btag"): + reordered[name] = event[name][::-1].copy() + first, _ = build_node_features(event, names, object_types, scales) + second, _ = build_node_features(reordered, names, object_types, scales) + torch.testing.assert_close(second[:2], first[[1, 0]]) + torch.testing.assert_close(second[2:], first[2:]) + + +def test_feature_scaling_changes_only_the_requested_columns(event, schema): + names, object_types, scales = schema + base, _ = build_node_features(event, names, object_types, scales) + changed_scales = list(scales) + changed_scales[0] *= 2 + scaled, _ = build_node_features(event, names, object_types, changed_scales) + torch.testing.assert_close(scaled[:, 0], base[:, 0] * 2) + torch.testing.assert_close(scaled[:, 1:], base[:, 1:]) diff --git a/tests/unit/graphs/__init__.py b/tests/unit/graphs/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/tests/unit/graphs/test_dgl_graph.py b/tests/unit/graphs/test_dgl_graph.py new file mode 100644 index 0000000000000000000000000000000000000000..8b089405b173aa537244093890a3c2daf199d892 --- /dev/null +++ b/tests/unit/graphs/test_dgl_graph.py @@ -0,0 +1,16 @@ +import torch + +from gnn4colliders.graphs import build_dgl_graph + + +def test_build_dgl_graph_attaches_compatible_features(): + node_features = torch.tensor( + [[10.0, 1.0, 3.1], [20.0, -0.5, -3.1]], dtype=torch.float32 + ) + graph = build_dgl_graph(node_features, eta_index=1, phi_index=2) + + assert graph.number_of_nodes() == 2 + assert graph.number_of_edges() == 2 + assert graph.ndata["features"] is node_features + assert graph.edata["features"].shape == (2, 3) + assert graph.edata["features"].dtype == torch.float32 diff --git a/tests/unit/graphs/test_edge_features.py b/tests/unit/graphs/test_edge_features.py new file mode 100644 index 0000000000000000000000000000000000000000..20fdd5cdff6dd5ff1d2a567a0998e6e35d957cfc --- /dev/null +++ b/tests/unit/graphs/test_edge_features.py @@ -0,0 +1,51 @@ +import pytest +import torch + +from gnn4colliders.graphs import build_edge_features, fully_connected_edges + + +def test_edge_features_have_legacy_order_signs_and_phi_wrapping(): + nodes = torch.tensor([[0.0, 1.0, 3.1], [0.0, -0.5, -3.1]]) + src, dst = fully_connected_edges(2) + actual = build_edge_features(nodes, src, dst, eta_index=1, phi_index=2) + expected = torch.tensor([[1.5, -0.0831853, 1.502304], [-1.5, 0.0831853, 1.502304]]) + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-6) + assert actual.dtype == torch.float32 + + +def test_empty_edges_keep_three_feature_columns(): + nodes = torch.empty((0, 3), dtype=torch.float32) + src, dst = fully_connected_edges(0) + features = build_edge_features(nodes, src, dst, eta_index=1, phi_index=2) + assert features.shape == (0, 3) + assert features.dtype == torch.float32 + + +def test_edge_features_have_symmetry_and_nonnegative_distance(): + nodes = torch.tensor([[1.0, 2.0, 3.13], [2.0, -1.0, -3.13], [3.0, 0.5, 0.2]]) + src, dst = fully_connected_edges(3) + actual = build_edge_features(nodes, src, dst, eta_index=1, phi_index=2) + reverse = build_edge_features(nodes, dst, src, eta_index=1, phi_index=2) + torch.testing.assert_close(actual[:, 0], -reverse[:, 0]) + torch.testing.assert_close(actual[:, 2], reverse[:, 2]) + assert torch.all(actual[:, 2] >= 0) + assert torch.all(actual[:, 1].abs() <= torch.pi) + + +@pytest.mark.parametrize("bad_nodes", [torch.zeros(3), torch.zeros(2, 2)]) +def test_edge_features_reject_malformed_node_features(bad_nodes): + with pytest.raises((ValueError, IndexError)): + build_edge_features( + bad_nodes, torch.tensor([0]), torch.tensor([0]), eta_index=1, phi_index=2 + ) + + +def test_edge_features_reject_mismatched_edge_indices(): + with pytest.raises(ValueError, match="equal shape"): + build_edge_features( + torch.zeros(2, 3), + torch.tensor([0]), + torch.tensor([1, 0]), + eta_index=1, + phi_index=2, + ) diff --git a/tests/unit/graphs/test_topology.py b/tests/unit/graphs/test_topology.py new file mode 100644 index 0000000000000000000000000000000000000000..ee8bd52748cf6ae0c29b8ae1c6e53240ae17858c --- /dev/null +++ b/tests/unit/graphs/test_topology.py @@ -0,0 +1,78 @@ +import pytest +import torch + +from gnn4colliders.graphs import ( + clear_topology_cache, + fully_connected_edges, + topology_cache_info, +) + + +def test_topology_cache_reuses_bounded_source_major_indices(): + clear_topology_cache() + first = fully_connected_edges(3) + second = fully_connected_edges(3) + assert torch.equal(first[0], second[0]) + assert first[0] is not second[0] + assert topology_cache_info() == {"size": 1, "max_size": 32} + for count in range(40): + fully_connected_edges(count) + assert topology_cache_info()["size"] <= 32 + clear_topology_cache() + + +@pytest.mark.parametrize( + ("num_nodes", "expected"), + [ + (1, [(0, 0)]), + (2, [(0, 1), (1, 0)]), + (3, [(0, 1), (0, 2), (1, 0), (1, 2), (2, 0), (2, 1)]), + ], +) +def test_fully_connected_edges_are_source_major(num_nodes, expected): + src, dst = fully_connected_edges(num_nodes) + assert list(zip(src.tolist(), dst.tolist())) == expected + assert src.dtype == dst.dtype == torch.long + + +def test_self_loops_can_be_enabled_and_zero_nodes_is_empty(): + src, dst = fully_connected_edges(2, self_loops=True) + assert list(zip(src.tolist(), dst.tolist())) == [ + (0, 0), + (0, 1), + (1, 0), + (1, 1), + ] + src, dst = fully_connected_edges(0) + assert src.shape == dst.shape == (0,) + + +def test_negative_node_count_is_rejected(): + with pytest.raises(ValueError, match="non-negative"): + fully_connected_edges(-1) + + +@pytest.mark.parametrize("num_nodes", [0, 1, 2, 4, 7]) +def test_topology_invariants_hold_for_multiple_graph_sizes(num_nodes): + src, dst = fully_connected_edges(num_nodes) + expected = ( + 0 if num_nodes == 0 else 1 if num_nodes == 1 else num_nodes * (num_nodes - 1) + ) + assert src.numel() == dst.numel() == expected + pairs = list(zip(src.tolist(), dst.tolist())) + assert len(pairs) == len(set(pairs)) + if num_nodes > 1: + assert all(source != destination for source, destination in pairs) + assert [src.tolist().count(index) for index in range(num_nodes)] == [ + num_nodes - 1 + ] * num_nodes + assert [dst.tolist().count(index) for index in range(num_nodes)] == [ + num_nodes - 1 + ] * num_nodes + + +def test_topology_rejects_bool_and_non_integer_counts(): + with pytest.raises(TypeError, match="integer"): + fully_connected_edges(True) + with pytest.raises(TypeError, match="integer"): + fully_connected_edges(2.0) diff --git a/tests/unit/inference/test_predictor.py b/tests/unit/inference/test_predictor.py new file mode 100644 index 0000000000000000000000000000000000000000..79892f440a32901fcdc2dae3cc3228f5a3533284 --- /dev/null +++ b/tests/unit/inference/test_predictor.py @@ -0,0 +1,107 @@ +from types import SimpleNamespace + +import numpy as np +import torch +from torch import nn + +from gnn4colliders.data import BatchMetadata +from gnn4colliders.inference import Predictor, write_npz +from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask + + +class _Model(nn.Module): + def __init__(self, outputs): + super().__init__() + self.outputs = torch.as_tensor(outputs, dtype=torch.float32) + self.dropout = nn.Dropout(0.9) + + def forward(self, batch): + count = ( + batch.labels.shape[0] + if hasattr(batch, "labels") + else len(batch.metadata.sample_id) + ) + return self.outputs[batch.offset : batch.offset + count] + + +def _loader(outputs, labels, *, offset=0, ids=("a", "b")): + metadata = BatchMetadata( + fold=torch.tensor([0, 1]), + weight=torch.tensor([1.0, 2.0]), + sample_id=ids, + extra={"source_file": ("events.root", "events.root")}, + ) + return [ + SimpleNamespace( + labels=torch.tensor(labels), + metadata=metadata, + offset=offset, + to=lambda device: SimpleNamespace( + labels=torch.tensor(labels), metadata=metadata, offset=offset + ), + ) + ] + + +def test_predictor_preserves_named_metadata_and_uses_task_scores(): + loader = _loader([[-2.0], [2.0]], [0, 1]) + result = Predictor(_Model([[-2.0], [2.0]]), BinaryClassificationTask()).predict( + loader + ) + assert result.sample_ids == ("a", "b") + assert result.scores.tolist() == [ + torch.sigmoid(torch.tensor(-2.0)).item(), + torch.sigmoid(torch.tensor(2.0)).item(), + ] + assert result.predictions.tolist() == [False, True] + assert result.weights.tolist() == [1.0, 2.0] + assert result.extra["source_file"] == ("events.root", "events.root") + + +def test_predictor_supports_unlabeled_multiclass_batches(): + metadata = BatchMetadata( + fold=torch.tensor([0]), weight=torch.tensor([1.0]), sample_id=("x",) + ) + batch = SimpleNamespace(metadata=metadata, offset=0, to=lambda device: batch) + result = Predictor( + _Model([[1.0, 3.0, 2.0]]), MulticlassClassificationTask() + ).predict([batch]) + assert result.labels is None + assert result.scores.shape == (1, 3) + assert result.predictions.tolist() == [1] + + +def test_npz_round_trip(tmp_path): + result = Predictor(_Model([[-1.0], [1.0]]), BinaryClassificationTask()).predict( + _loader([[-1.0], [1.0]], [0, 1]) + ) + path = write_npz(result, tmp_path / "predictions.npz") + with np.load(path) as values: + assert values["sample_id"].tolist() == ["a", "b"] + np.testing.assert_array_equal(values["logits"], result.logits.numpy()) + np.testing.assert_array_equal(values["labels"], result.labels.numpy()) + np.testing.assert_array_equal(values["fold"], result.fold.numpy()) + + +def test_predictor_is_reentrant_and_returns_detached_cpu_results(): + predictor = Predictor(_Model([[-1.0], [1.0]]), BinaryClassificationTask()) + first = predictor.predict(_loader([[-1.0], [1.0]], [0, 1])) + second = predictor.predict(_loader([[-1.0], [1.0]], [0, 1])) + torch.testing.assert_close(first.logits, second.logits) + assert first.logits.device.type == "cpu" + assert not first.logits.requires_grad + assert first.logits.grad_fn is None + + +def test_npz_omits_labels_for_unlabeled_results(tmp_path): + metadata = BatchMetadata( + fold=torch.tensor([0]), weight=torch.tensor([1.0]), sample_id=("sample_7",) + ) + batch = SimpleNamespace(metadata=metadata, offset=0, to=lambda device: batch) + result = Predictor(_Model([[2.0, 1.0]]), MulticlassClassificationTask()).predict( + [batch] + ) + path = write_npz(result, tmp_path / "unlabeled.npz") + with np.load(path) as values: + assert "labels" not in values + assert values["sample_id"].tolist() == ["sample_7"] diff --git a/tests/unit/models/root_gnn/test_blocks.py b/tests/unit/models/root_gnn/test_blocks.py new file mode 100644 index 0000000000000000000000000000000000000000..f916e5a1b1ba488dcc57c581e3498c4023a6cbc2 --- /dev/null +++ b/tests/unit/models/root_gnn/test_blocks.py @@ -0,0 +1,22 @@ +import pytest +from torch import nn + +from gnn4colliders.models.root_gnn.blocks import make_mlp + + +def test_make_mlp_matches_legacy_layer_order(): + block = make_mlp(3, 4, 5, 2, dropout=0.1) + assert [type(layer) for layer in block] == [ + nn.Linear, + nn.ReLU, + nn.Dropout, + nn.Linear, + nn.ReLU, + nn.Dropout, + nn.LayerNorm, + ] + + +def test_make_mlp_rejects_empty_network(): + with pytest.raises(ValueError, match="positive"): + make_mlp(3, 4, 5, 0) diff --git a/tests/unit/models/root_gnn/test_edge_cases.py b/tests/unit/models/root_gnn/test_edge_cases.py new file mode 100644 index 0000000000000000000000000000000000000000..6ff3c4a195a60c8130a02e3da40779de5942e202 --- /dev/null +++ b/tests/unit/models/root_gnn/test_edge_cases.py @@ -0,0 +1,89 @@ +import pytest +import torch + +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + +pytest.importorskip("dgl") + + +def _graph(node_features): + graph = build_dgl_graph(torch.tensor(node_features, dtype=torch.float32)) + return graph, torch.tensor([[1.0]], dtype=torch.float32) + + +def test_no_globals_uses_graph_node_count_as_legacy_fallback(): + graph, _ = _graph([[10.0, 0.2, 0.1]]) + model = EdgeNetwork(graph, torch.empty((1, 0)), 4, 1, 1, 0).eval() + with torch.no_grad(): + logits = model(graph) + assert logits.shape == (1, 1) + assert torch.isfinite(logits).all() + + +def test_missing_or_mismatched_global_shapes_fail_clearly(): + graph, globals_ = _graph([[10.0, 0.2, 0.1]]) + model = EdgeNetwork(graph, globals_, 4, 1, 1, 0) + with pytest.raises(ValueError, match="global_features are required"): + model(graph) + batched_graph = build_dgl_graph( + torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]], dtype=torch.float32) + ) + batched_graph = __import__("dgl").batch([graph, batched_graph]) + with pytest.raises(ValueError, match="one-dimensional"): + model(batched_graph, torch.tensor([1.0])) + + +def test_model_preserves_float64_dtype(): + graph = build_dgl_graph( + torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]], dtype=torch.float64) + ) + globals_ = torch.tensor([[1.0]], dtype=torch.float64) + model = EdgeNetwork(graph, globals_, 4, 1, 1, 1).double().eval() + assert model(graph, globals_).dtype == torch.float64 + + +def test_negative_processing_steps_are_rejected(): + graph, globals_ = _graph([[10.0, 0.2, 0.1]]) + with pytest.raises(ValueError, match="non-negative"): + EdgeNetwork(graph, globals_, 4, 1, 1, -1) + + +def test_single_node_graph_and_zero_processing_step_are_supported(): + graph, globals_ = _graph([[10.0, 0.2, 0.1]]) + model = EdgeNetwork(graph, globals_, 4, 1, 1, 0).eval() + assert graph.number_of_edges() == 1 + assert model(graph, globals_).shape == (1, 1) + + +def test_weighted_padded_nodes_produce_finite_logits(): + graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2], [0.0, 0.0, 0.0]]) + graph.ndata["w"] = torch.tensor([[1.0], [1.0], [0.0]]) + model = EdgeNetwork(graph, globals_, 4, 1, 1, 1).eval() + logits = model(graph, globals_) + assert logits.shape == (1, 1) + assert torch.isfinite(logits).all() + + +def test_unfrozen_transfer_backbone_receives_gradients(): + graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]) + source = EdgeNetwork(graph, globals_, 4, 3, 1, 1) + transfer = FineTunedEdgeNetwork.from_pretrained(source, 1, freeze_backbone=False) + transfer(graph, globals_).sum().backward() + assert any( + parameter.grad is not None for parameter in transfer.backbone.parameters() + ) + assert all( + parameter.grad is not None for parameter in transfer.classifier.parameters() + ) + + +def test_frozen_transfer_backbone_has_no_gradients_and_source_is_unchanged(): + graph, globals_ = _graph([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]]) + source = EdgeNetwork(graph, globals_, 4, 3, 1, 1) + source_classifier_type = type(source.classifier) + transfer = FineTunedEdgeNetwork.from_pretrained(source, 1, freeze_backbone=True) + transfer(graph, globals_).sum().backward() + assert isinstance(source.classifier, source_classifier_type) + assert isinstance(transfer.backbone.classifier, torch.nn.Identity) + assert all(parameter.grad is None for parameter in transfer.backbone.parameters()) diff --git a/tests/unit/models/root_gnn/test_edge_network.py b/tests/unit/models/root_gnn/test_edge_network.py new file mode 100644 index 0000000000000000000000000000000000000000..154a42d243556932b37b48482d46018d633e2c7c --- /dev/null +++ b/tests/unit/models/root_gnn/test_edge_network.py @@ -0,0 +1,78 @@ +import pytest +import torch + +from gnn4colliders.graphs import build_dgl_graph +from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + +dgl = pytest.importorskip("dgl") + + +def _model(out_size=1): + graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]])) + return EdgeNetwork(graph, torch.zeros(1, 2), 8, out_size, 2, 2), graph + + +def test_edge_network_returns_raw_binary_logits_without_graph_leakage(): + model, graph = _model() + globals_ = torch.tensor([[1.0, 2.0]]) + before = set(graph.ndata.keys()) | set(graph.edata.keys()) + logits = model(graph, globals_) + again = model(graph, globals_) + assert logits.shape == (1, 1) + assert logits.dtype == torch.float32 + assert torch.allclose(logits, again) + assert torch.equal( + model.representation(graph, globals_), model.forward_features(graph, globals_) + ) + assert model.classify is model.classifier + assert before == (set(graph.ndata.keys()) | set(graph.edata.keys())) + + +def test_edge_network_supports_multiclass_and_transfer_freezing(): + model, graph = _model(out_size=12) + transferred = FineTunedEdgeNetwork.from_pretrained(model, 1, freeze_backbone=True) + assert transferred(graph, torch.zeros(1, 2)).shape == (1, 1) + assert not any( + parameter.requires_grad for parameter in transferred.backbone.parameters() + ) + assert all( + parameter.requires_grad for parameter in transferred.classifier.parameters() + ) + assert model(graph, torch.zeros(1, 2)).shape == (1, 12) + + +@pytest.mark.parametrize("out_size", [1, 3]) +@pytest.mark.parametrize("n_proc_steps", [0, 1]) +def test_model_shape_matrix_and_state_dict_round_trip(out_size, n_proc_steps): + graph = build_dgl_graph(torch.tensor([[10.0, 0.2, 0.1], [8.0, -0.3, -0.2]])) + globals_ = torch.tensor([[1.0, 2.0]]) + torch.manual_seed(7) + model = EdgeNetwork(graph, globals_, 4, out_size, 1, n_proc_steps).eval() + with torch.no_grad(): + expected = model(graph, globals_) + fresh = EdgeNetwork(graph, globals_, 4, out_size, 1, n_proc_steps).eval() + fresh.load_state_dict(model.state_dict()) + with torch.no_grad(): + actual = fresh(graph, globals_) + torch.testing.assert_close(actual, expected) + + +def test_model_backward_produces_finite_gradients_and_update(): + model, graph = _model() + model.train() + before = [parameter.detach().clone() for parameter in model.parameters()] + loss = model(graph, torch.zeros(1, 2)).square().mean() + assert torch.isfinite(loss) + loss.backward() + gradients = [ + parameter.grad for parameter in model.parameters() if parameter.requires_grad + ] + assert gradients and all( + gradient is not None and torch.isfinite(gradient).all() + for gradient in gradients + ) + optimizer = torch.optim.SGD(model.parameters(), lr=0.01) + optimizer.step() + assert any( + not torch.equal(old, new) for old, new in zip(before, model.parameters()) + ) diff --git a/tests/unit/tasks/test_classification.py b/tests/unit/tasks/test_classification.py new file mode 100644 index 0000000000000000000000000000000000000000..5a22fe2a67b6dbdc12918a15dfd6eed9b7a5ae31 --- /dev/null +++ b/tests/unit/tasks/test_classification.py @@ -0,0 +1,84 @@ +from types import SimpleNamespace + +import pytest +import torch + +from gnn4colliders.tasks import BinaryClassificationTask, MulticlassClassificationTask + + +def _batch(labels, weights): + return SimpleNamespace( + labels=torch.tensor(labels), + metadata=SimpleNamespace(weight=torch.tensor(weights)), + ) + + +def test_binary_normalizes_column_or_vector_targets_and_uses_named_weight(): + task = BinaryClassificationTask() + batch = _batch([[0], [0], [1], [1]], [1.0, 2.0, 1.0, 3.0]) + logits = torch.tensor([[-1.0], [0.0], [0.5], [1.0]]) + expected_elementwise = torch.nn.functional.binary_cross_entropy_with_logits( + logits[:, 0], batch.labels[:, 0].float(), reduction="none" + ) + expected = ( + (expected_elementwise[:2] * torch.tensor([1.0, 2.0])).sum() / 3 + + (expected_elementwise[2:] * torch.tensor([1.0, 3.0])).sum() / 4 + ) / 2 + assert torch.allclose(task.loss(logits, batch), expected) + + +def test_binary_negative_weights_are_not_absolute_by_default(): + batch = _batch([0, 0, 1, 1], [1.0, -2.0, 1.0, 1.0]) + logits = torch.tensor([[0.0], [2.0], [0.0], [2.0]]) + task = BinaryClassificationTask() + elementwise = torch.nn.functional.binary_cross_entropy_with_logits( + logits[:, 0], batch.labels.float(), reduction="none" + ) + expected = ( + (elementwise[:2] * torch.tensor([1.0, -2.0])).sum() / -1 + + elementwise[2:].mean() + ) / 2 + assert task.loss(logits, batch).item() == pytest.approx(expected.item()) + assert BinaryClassificationTask(absolute_weights=True).loss( + logits, batch + ).item() != pytest.approx(expected.item()) + + +def test_binary_predictions_and_undefined_auc(): + task = BinaryClassificationTask() + batch = _batch([1, 1], [1.0, 2.0]) + output = task.predictions(torch.tensor([[-2.0], [2.0]])) + assert output["scores"].shape == (2,) + assert output["predictions"].tolist() == [False, True] + assert ( + task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"] + != task.metrics(torch.tensor([[-2.0], [2.0]]), batch)["roc_auc"] + ) + + +def test_multiclass_loss_and_metrics(): + task = MulticlassClassificationTask() + batch = _batch([0, 1, 2], [1.0, 1.0, 1.0]) + logits = torch.eye(3) + assert task.loss(logits, batch).item() > 0 + output = task.predict(logits) + assert output["predictions"].tolist() == [0, 1, 2] + assert task.metrics(logits, batch)["accuracy"] == 1.0 + + +def test_binary_loss_is_invariant_under_joint_event_permutation(): + task = BinaryClassificationTask() + batch = _batch([0, 1, 0, 1], [1.0, 2.0, 3.0, 4.0]) + logits = torch.tensor([[-2.0], [1.0], [0.5], [-0.5]]) + permutation = torch.tensor([2, 0, 3, 1]) + shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist()) + assert task.loss(logits, batch) == pytest.approx(task.loss(logits[permutation], shuffled).item()) + + +def test_multiclass_metrics_are_invariant_under_joint_event_permutation(): + task = MulticlassClassificationTask() + batch = _batch([0, 1, 2, 1], [1.0, 2.0, 1.0, 3.0]) + logits = torch.tensor([[3.0, 0.0, 0.0], [0.0, 2.0, 0.0], [0.0, 0.0, 4.0], [0.0, 1.0, 0.0]]) + permutation = torch.tensor([3, 0, 2, 1]) + shuffled = _batch(batch.labels[permutation].tolist(), batch.metadata.weight[permutation].tolist()) + assert task.metrics(logits, batch) == task.metrics(logits[permutation], shuffled) diff --git a/tests/unit/test_package_import.py b/tests/unit/test_package_import.py new file mode 100644 index 0000000000000000000000000000000000000000..643b103dc0c84a39ce9c217e8cc2b80942cd8bc0 --- /dev/null +++ b/tests/unit/test_package_import.py @@ -0,0 +1,39 @@ +"""Smoke tests for the installable package boundary.""" + + +def test_package_imports() -> None: + import gnn4colliders + + assert gnn4colliders.__file__ is not None + assert gnn4colliders.__version__ == "0.1.0" + + +def test_public_import_contract() -> None: + from gnn4colliders.inference import Predictor + from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + from gnn4colliders.tasks import ( + BinaryClassificationTask, + MulticlassClassificationTask, + ) + from gnn4colliders.training import CheckpointManager, Trainer + + assert all( + item is not None + for item in ( + EdgeNetwork, + FineTunedEdgeNetwork, + BinaryClassificationTask, + MulticlassClassificationTask, + Trainer, + CheckpointManager, + Predictor, + ) + ) + + +def test_hydra_config_is_discoverable_from_package() -> None: + from gnn4colliders.cli import _config + + config = _config([]) + assert config.model.name == "edge_network" + assert config.task.type == "multiclass_classification" diff --git a/tests/unit/training/test_checkpoint.py b/tests/unit/training/test_checkpoint.py new file mode 100644 index 0000000000000000000000000000000000000000..b4e36a79bb3ecafd5b41410045bee2c8a34c2ed7 --- /dev/null +++ b/tests/unit/training/test_checkpoint.py @@ -0,0 +1,127 @@ +import random + +import numpy as np +import pytest +import torch +from torch import nn + +from gnn4colliders.training import ( + CheckpointManager, + EarlyStopping, + TrainerState, + load_legacy_checkpoint, + load_model_weights, + restore_training_state, +) + + +def test_checkpoint_round_trip_and_selection(tmp_path): + model = nn.Linear(2, 1) + optimizer = torch.optim.Adam(model.parameters(), lr=0.1) + scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=0.5) + model(torch.ones(2)).sum().backward() + optimizer.step() + scheduler.step() + early = EarlyStopping(patience=2) + early.update(3.0) + state = TrainerState(epoch=9, global_step=17) + manager = CheckpointManager(tmp_path) + manager.save( + model=model, + optimizer=optimizer, + scheduler=scheduler, + early_stopping=early, + trainer_state=state, + task_config={"type": "binary_classification", "threshold": 0.5}, + monitor_name="loss", + monitor_value=0.4, + ) + state.epoch = 10 + manager.save( + model=model, + optimizer=optimizer, + scheduler=scheduler, + early_stopping=early, + trainer_state=state, + monitor_name="loss", + monitor_value=0.8, + ) + + assert manager.latest().name == "epoch_0010.pt" + assert manager.best(monitor="loss").name == "epoch_0009.pt" + payload = manager.load(manager.latest()) + assert payload["schema_version"] == 1 + assert payload["task_config"] is None # omitted explicitly for the second save + + restored_model = nn.Linear(2, 1) + restored_optimizer = torch.optim.Adam(restored_model.parameters(), lr=0.1) + restored_scheduler = torch.optim.lr_scheduler.ExponentialLR( + restored_optimizer, gamma=0.5 + ) + restored_early = EarlyStopping(patience=2) + restored = restore_training_state( + payload, + model=restored_model, + optimizer=restored_optimizer, + scheduler=restored_scheduler, + early_stopping=restored_early, + ) + assert restored == TrainerState(epoch=10, global_step=17) + assert restored_scheduler.last_epoch == scheduler.last_epoch + assert restored_early.best == early.best + assert all( + torch.equal(left, right) + for left, right in zip(model.parameters(), restored_model.parameters()) + ) + + +def test_weight_only_load_does_not_need_optimizer(tmp_path): + model = nn.Linear(1, 1) + manager = CheckpointManager(tmp_path) + path = manager.save(model=model, trainer_state=TrainerState(epoch=0)) + fresh = nn.Linear(1, 1) + load_model_weights(fresh, manager.load(path)) + assert torch.equal(model.weight, fresh.weight) + assert torch.equal(model.bias, fresh.bias) + + +def test_legacy_checkpoint_normalizes_supported_prefixes(): + state = {"module._orig_mod.linear.weight": torch.ones(1, 1)} + normalized = load_legacy_checkpoint( + {"epoch": 4, "model_state_dict": state, "early_stop": {"count": 2}} + ) + assert list(normalized["model_state_dict"]) == ["linear.weight"] + assert normalized["epoch"] == 4 + assert normalized["early_stopping_state"]["num_bad_epochs"] == 2 + + +def test_rng_state_is_restored(tmp_path): + manager = CheckpointManager(tmp_path) + model = nn.Linear(1, 1) + path = manager.save(model=model, trainer_state=TrainerState(epoch=0)) + expected = (random.random(), np.random.random(), torch.rand(1)) + payload = manager.load(path) + restore_training_state(payload, model=nn.Linear(1, 1)) + actual = (random.random(), np.random.random(), torch.rand(1)) + assert expected[0] == actual[0] + assert expected[1] == actual[1] + assert torch.equal(expected[2], actual[2]) + + +def test_checkpoint_rejects_unsupported_schema_and_missing_weights(tmp_path): + unsupported = tmp_path / "unsupported.pt" + torch.save({"schema_version": 999}, unsupported) + with pytest.raises(ValueError, match="unsupported checkpoint schema"): + CheckpointManager.load(unsupported) + + missing = tmp_path / "missing.pt" + torch.save({"schema_version": 1, "epoch": 0}, missing) + with pytest.raises(ValueError, match="incompatible"): + load_model_weights(nn.Linear(2, 1), CheckpointManager.load(missing)) + + +def test_checkpoint_loading_rejects_architecture_mismatch(tmp_path): + manager = CheckpointManager(tmp_path) + path = manager.save(model=nn.Linear(2, 1), trainer_state=TrainerState(epoch=0)) + with pytest.raises(ValueError, match="incompatible"): + load_model_weights(nn.Linear(3, 1), manager.load(path)) diff --git a/tests/unit/training/test_trainer.py b/tests/unit/training/test_trainer.py new file mode 100644 index 0000000000000000000000000000000000000000..122c6f09d1954b558a407aae308472092d79b21e --- /dev/null +++ b/tests/unit/training/test_trainer.py @@ -0,0 +1,92 @@ +import pytest +import torch +from torch import nn + +from gnn4colliders.data import BatchMetadata, EventMetadata, GraphBatch +from gnn4colliders.tasks import BinaryClassificationTask +from gnn4colliders.training import ( + EarlyStopping, + Trainer, + build_optimizer, + build_scheduler, +) + + +class TinyModel(nn.Module): + def __init__(self): + super().__init__() + self.linear = nn.Linear(1, 1) + + def forward(self, batch): + return self.linear(batch.graph) + + +def _batch(values, labels, weights=None): + values = torch.tensor(values, dtype=torch.float32).reshape(-1, 1) + if weights is None: + weights = [1.0] * len(labels) + return GraphBatch( + graph=values, + labels=torch.tensor(labels), + global_features=None, + metadata=BatchMetadata.from_events( + [ + EventMetadata(fold=0, weight=weight, sample_id=f"event:{index}") + for index, weight in enumerate(weights) + ] + ), + ) + + +def test_graph_batch_to_moves_named_tensors_without_losing_ids(): + batch = _batch([1, 2], [0, 1]) + moved = batch.to("cpu") + assert moved.labels.device.type == "cpu" + assert moved.metadata.weight.device.type == "cpu" + assert moved.metadata.sample_id == ("event:0", "event:1") + + +def test_trainer_aggregates_metrics_over_the_complete_epoch(): + torch.manual_seed(4) + model = TinyModel() + optimizer = build_optimizer(model, learning_rate=0.01) + scheduler = build_scheduler(optimizer, gamma=0.5) + trainer = Trainer(model, BinaryClassificationTask(), optimizer, scheduler) + loader = [ + _batch([1, 2], [0, 1]), + _batch([3, 4], [0, 1]), + ] + before = model.linear.weight.detach().clone() + result = trainer.train_epoch(loader) + assert result.num_samples == 4 + assert set(result.metrics) == {"accuracy", "roc_auc", "auc"} + assert torch.isfinite(torch.tensor(result.loss)) + assert not torch.equal(before, model.linear.weight.detach()) + trainer.scheduler.step() + assert optimizer.param_groups[0]["lr"] == 0.005 + + +def test_fit_uses_validation_for_early_stopping_and_records_history(): + model = TinyModel() + optimizer = build_optimizer(model, learning_rate=0.01) + trainer = Trainer( + model, + BinaryClassificationTask(), + optimizer, + early_stopping=EarlyStopping(patience=1), + ) + loader = [_batch([1, 2], [0, 1])] + history = trainer.fit(loader, loader, epochs=4) + assert len(history.train) == len(history.validation) + assert 1 <= len(history.train) <= 4 + assert all(isinstance(item.loss, float) for item in history.train) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA unavailable") +def test_graph_batch_to_cuda_preserves_named_metadata(): + batch = _batch([1], [0]) + moved = batch.to("cuda") + assert moved.labels.device.type == "cuda" + assert moved.metadata.weight.device.type == "cuda" + assert moved.metadata.sample_id == ("event:0",) diff --git a/tests/unit/validation/conftest.py b/tests/unit/validation/conftest.py new file mode 100644 index 0000000000000000000000000000000000000000..b41e065dec0fefb37304bc3b81bc8d881ca62587 --- /dev/null +++ b/tests/unit/validation/conftest.py @@ -0,0 +1,10 @@ +"""Make the repository-level validation package importable in src-layout tests.""" + +import sys +from pathlib import Path + +ROOT = str(Path(__file__).resolve().parents[3]) +sys.path.insert(0, ROOT) +validation_module = sys.modules.get("validation") +if validation_module is not None: + validation_module.__path__ = [str(Path(ROOT) / "validation")] diff --git a/tests/unit/validation/test_artifacts.py b/tests/unit/validation/test_artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..c95fb6681b578f5cecef101d1e66155de9aaaa46 --- /dev/null +++ b/tests/unit/validation/test_artifacts.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +import numpy as np + +from validation.artifacts import ValidationArtifact, load_artifact, save_artifact +from validation.compare import compare_artifacts + + +def _artifact(value: float = 1.0) -> ValidationArtifact: + return ValidationArtifact( + sample_id=np.asarray(["file.root:0", "file.root:1"]), + labels=np.asarray([0, 1]), + folds=np.asarray([0, 1]), + weights=np.asarray([1.0, -2.0], dtype=np.float32), + globals=np.asarray([[value, 2.0], [3.0, 4.0]], dtype=np.float32), + node_features_flat=np.asarray( + [[1.0, 0.0], [2.0, 1.0], [3.0, 2.0]], dtype=np.float32 + ), + node_offsets=np.asarray([0, 1, 3]), + edge_src_flat=np.asarray([0, 1]), + edge_dst_flat=np.asarray([0, 0]), + edge_features_flat=np.asarray([[0.1], [0.2]], dtype=np.float32), + edge_offsets=np.asarray([0, 0, 2]), + logits=np.asarray([[0.1], [0.2]], dtype=np.float32), + ) + + +def test_artifact_round_trip_preserves_variable_size_arrays(tmp_path): + save_artifact(_artifact(), tmp_path / "artifact") + restored = load_artifact(tmp_path / "artifact") + assert restored.event_count == 2 + np.testing.assert_array_equal(restored.sample_id, ["file.root:0", "file.root:1"]) + np.testing.assert_array_equal(restored.event_nodes(1), [[2.0, 1.0], [3.0, 2.0]]) + np.testing.assert_array_equal(restored.event_edges(1)[0], [0, 1]) + + +def test_comparator_reports_tolerance_and_exact_mismatches(): + same = compare_artifacts(_artifact(), _artifact(1.0000001)) + assert same["overall_status"] == "PASS" + changed = _artifact() + changed.labels[1] = 0 + report = compare_artifacts(_artifact(), changed) + assert report["overall_status"] == "FAIL" + assert report["stages"]["labels"]["status"] == "FAIL" + + +def test_comparator_rejects_duplicate_sample_ids(): + duplicate = _artifact() + duplicate.sample_id[1] = duplicate.sample_id[0] + try: + compare_artifacts(_artifact(), duplicate) + except ValueError as error: + assert "duplicates" in str(error) + else: + raise AssertionError("duplicate IDs must fail loudly") diff --git a/uv.lock b/uv.lock new file mode 100644 index 0000000000000000000000000000000000000000..f3d14c74815836813507f5fc206ca83d424bbe78 --- /dev/null +++ b/uv.lock @@ -0,0 +1,1107 @@ +version = 1 +revision = 3 +requires-python = "==3.12.*" +resolution-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] +supported-markers = [ + "platform_machine == 'x86_64' and sys_platform == 'linux'", +] + +[[package]] +name = "annotated-types" +version = "0.8.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/56/a8120250d128bed162cd73c76d45f6ef9991f3e068f62a8ee060afa3104a/annotated_types-0.8.0.tar.gz", hash = "sha256:13b2beaad985e05e2d6407ee4c4f35590b11f8d693a258a561055cac8f64cab7", size = 15893, upload-time = "2026-07-23T20:16:13.995Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/99/91/8acff4f5e50511b911bbccb72b8628a49c68ce14148cd9f6431094859a90/annotated_types-0.8.0-py3-none-any.whl", hash = "sha256:f072f4d804ea359e4eaf198b1af7a8b0943881a87f31bb764f8bf219bb9419e0", size = 13427, upload-time = "2026-07-23T20:16:12.938Z" }, +] + +[[package]] +name = "antlr4-python3-runtime" +version = "4.9.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/3e/38/7859ff46355f76f8d19459005ca000b6e7012f2f1ca597746cbcd1fbfe5e/antlr4-python3-runtime-4.9.3.tar.gz", hash = "sha256:f224469b4168294902bb1efa80a8bf7855f24c99aef99cbefc1bcd3cce77881b", size = 117034, upload-time = "2021-11-06T17:52:23.524Z" } + +[[package]] +name = "awkward" +version = "2.13.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "awkward-cpp", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fsspec", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/97/2728ab879ed3edaecfbd8e9daa7d88a2f89a7ea584e413c4d8971fc0414d/awkward-2.13.0.tar.gz", hash = "sha256:36f127573295e1ddf65b551bf071b803ad4af21bd14519f4c4dd34953cb3efc2", size = 6479149, upload-time = "2026-08-14T16:52:43.873Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/25/34/b8eece9a8d024defc08418ea883361b3d7e3300392719ee27535e93885a5/awkward-2.13.0-py3-none-any.whl", hash = "sha256:ff40879e7179a6f14f4c4ee5f5297e3dcd027b99b4c64188bdb7dccdf625e982", size = 990112, upload-time = "2026-08-14T16:52:42.018Z" }, +] + +[[package]] +name = "awkward-cpp" +version = "56" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/92/6325292ccc2be3ede227fd0b4bd7b23bc7597790f31227171bec4009d2df/awkward_cpp-56.tar.gz", hash = "sha256:cd3635fab926c6630c0a85d92b59a15851c4f5a881ea749984069027634f8f52", size = 1501723, upload-time = "2026-08-14T15:30:10.748Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a5/35/b33ce937f644f33a0a3c7185c525163f98152e1b29ee1a26917a3c25597a/awkward_cpp-56-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:82fa36172a714cd25907e9921c417f7a4c4f78e5151b6c0b2d82ed1ed3b8a5d8", size = 698758, upload-time = "2026-08-14T15:29:13.364Z" }, + { url = "https://files.pythonhosted.org/packages/b4/74/4dbfedd5fba6169f14f19f2637c7a40b5f1ed3d3128dde2056a5b644584c/awkward_cpp-56-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:e59f4f933052a1b97bec3ce749483fca8ebb424f0d220c103538c21730c6a9a8", size = 1745779, upload-time = "2026-08-14T15:29:16.681Z" }, +] + +[[package]] +name = "certifi" +version = "2026.7.22" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a3/c2/24167ea9858356b47a87a50d39908bfdb72ceeefe0041586e704e5376b3a/certifi-2026.7.22.tar.gz", hash = "sha256:741e2c3b351ddf169a738da9f2c048608ff7f2c5cc02f1ebc6b118bb090d5d55", size = 138112, upload-time = "2026-07-22T03:35:12.644Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0b/a7/71ac2cff56fec219ed242bb11b8efb69fcc4bec75db06fb7bfe35de520e6/certifi-2026.7.22-py3-none-any.whl", hash = "sha256:62f22742b58a1a33014a2b6b706588a8d7e2a88ae7bd1a6ebe8c992928483775", size = 136983, upload-time = "2026-07-22T03:35:11.276Z" }, +] + +[[package]] +name = "cffi" +version = "2.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "pycparser", marker = "implementation_name != 'PyPy' and platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9e/ef/008a1939e372c06329a3fce4279c02f328488f3526744906eeec3da7ad5f/cffi-2.1.1.tar.gz", hash = "sha256:dd31f52ea1086513bb9df30f8fcee9b8918323ae067a3d5b78bc826a000712be", size = 530807, upload-time = "2026-08-03T21:21:18.939Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b1/db/dceb9dd5b231e1da801793f8acc9f3c52a7e1afe40bb1aae37e02b0faad5/cffi-2.1.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:c1453022f490d2459a11819d83ad1d586e9ff65a12ac3e705ffebd46d3685dcf", size = 221822, upload-time = "2026-08-03T21:19:52.054Z" }, + { url = "https://files.pythonhosted.org/packages/cb/52/3fa190537004dd7f0ab860a6dc7c0175b8667f68d1e618a46f5498d30250/cffi-2.1.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:210019b6c7cf07f081b4c54635c8cf744377001350e29cc0f81c4377b4797735", size = 223597, upload-time = "2026-08-03T21:19:54.515Z" }, +] + +[[package]] +name = "charset-normalizer" +version = "3.5.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e5/3f/143b048436775b0f76ac3eec145c019e8173ccc2885c8f20319b996d5e83/charset_normalizer-3.5.1.tar.gz", hash = "sha256:6117b84ea48435e5356dc737f5121485c30920ba43375fa7b434fd753df0eac3", size = 171764, upload-time = "2026-08-15T08:20:44.807Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6f/89/bb5108dc6c3651dca963f2b0a3ba19bbcb370c94e1b6d3e0e844a58e6dca/charset_normalizer-3.5.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b9af956078716df40d985fb0dfeb2c2120c5ca92ba4ff4b388acfd01cdc14d08", size = 248801, upload-time = "2026-08-15T08:17:17.683Z" }, + { url = "https://files.pythonhosted.org/packages/b7/52/643d11ffd60e9ac2fd1fb87e167a19285b9eefeff4a40e63c87cbfbeab36/charset_normalizer-3.5.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:ba501e667c17d8411f98e67a022d9604ef179aff0e459b7e292c796837c13573", size = 250562, upload-time = "2026-08-15T08:17:27.971Z" }, + { url = "https://files.pythonhosted.org/packages/9f/2f/fe3f187327aac18e2d54e9d2b08e15d27bf9b642d9e51c219f130fc34d1a/charset_normalizer-3.5.1-cp37-abi3-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:a6dac12ff6b846103483683f60c5f8fee205121adc58ffd87e90a90a3af69e99", size = 253057, upload-time = "2026-08-15T08:19:52.654Z" }, + { url = "https://files.pythonhosted.org/packages/80/c2/a7379b840292d0c1ab9fbd17d1f3967aa81794dc95bc74be8999d7fedcf7/charset_normalizer-3.5.1-cp37-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:e199fb99720074809a7720f1c0b4d919eea8b87e88713e0f8f602f7bef543d9d", size = 254846, upload-time = "2026-08-15T08:20:10.727Z" }, + { url = "https://files.pythonhosted.org/packages/cc/61/d01fc49b8dea277640b55a9e15960dbca9fdc8c9fde18e572d39c59f4019/charset_normalizer-3.5.1-py3-none-any.whl", hash = "sha256:6df0ec430f9a831772c23ca5a224cba36517a58a84bb32c32bb59a9fa67c47f6", size = 68658, upload-time = "2026-08-15T08:20:43.306Z" }, +] + +[[package]] +name = "contourpy" +version = "1.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/58/01/1253e6698a07380cd31a736d248a3f2a50a7c88779a1813da27503cadc2a/contourpy-1.3.3.tar.gz", hash = "sha256:083e12155b210502d0bca491432bb04d56dc3432f95a979b429f2848c3dbe880", size = 13466174, upload-time = "2025-07-26T12:03:12.549Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cc/8f/ec6289987824b29529d0dfda0d74a07cec60e54b9c92f3c9da4c0ac732de/contourpy-1.3.3-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4d00e655fcef08aba35ec9610536bfe90267d7ab5ba944f7032549c55a146da1", size = 362601, upload-time = "2025-07-26T12:01:28.808Z" }, + { url = "https://files.pythonhosted.org/packages/33/1d/acad9bd4e97f13f3e2b18a3977fe1b4a37ecf3d38d815333980c6c72e963/contourpy-1.3.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:459c1f020cd59fcfe6650180678a9993932d80d44ccde1fa1868977438f0b411", size = 1403386, upload-time = "2025-07-26T12:01:33.947Z" }, +] + +[[package]] +name = "coverage" +version = "7.15.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/be/c3/4f2195f512fb172aa425a8803a874b2baa9ba7f80ff7b6080998761fc701/coverage-7.15.4.tar.gz", hash = "sha256:0548198fff07ccf4faf469520bce1c2eceb1ce3e62891921138dec10907f9d00", size = 936952, upload-time = "2026-08-06T13:50:24.442Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8c/3f/f0642a372f494bd0d7dad3b497083b910194a5f1c88be2c94fef707c3b59/coverage-7.15.4-cp312-cp312-manylinux1_x86_64.manylinux_2_28_x86_64.manylinux_2_5_x86_64.whl", hash = "sha256:899b9da30f3c6c336566e3707495bb23e8302d39d862f01fa78c48b99b9437e2", size = 257145, upload-time = "2026-08-06T13:47:41.931Z" }, + { url = "https://files.pythonhosted.org/packages/e9/b9/fdb92c8ae7a8bb9b850cc253b7b3b9c8526f68130002048b5671cd510d09/coverage-7.15.4-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:c2de40895718f91951b86712b4c5b694acaf9a0a49be13874896f599a1eed3f4", size = 255763, upload-time = "2026-08-06T13:47:54.296Z" }, + { url = "https://files.pythonhosted.org/packages/b4/d9/e70c286c979378f061d8266e279b686ab0b0b688e1fe0af864684f23a77d/coverage-7.15.4-py3-none-any.whl", hash = "sha256:964730a1e9de9c0cf11be6a1a3c79ce419c34882842abd256086ba4698705e84", size = 214332, upload-time = "2026-08-06T13:50:22.192Z" }, +] + +[[package]] +name = "cramjam" +version = "2.11.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/14/12/34bf6e840a79130dfd0da7badfb6f7810b8fcfd60e75b0539372667b41b6/cramjam-2.11.0.tar.gz", hash = "sha256:5c82500ed91605c2d9781380b378397012e25127e89d64f460fea6aeac4389b4", size = 99100, upload-time = "2025-07-27T21:25:07.559Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/dc/f0/5c2a5cd5711032f3b191ca50cb786c17689b4a9255f9f768866e6c9f04d9/cramjam-2.11.0-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2fa2fe41f48c4d58d923803383b0737f048918b5a0d10390de9628bb6272b107", size = 1978104, upload-time = "2025-07-27T21:22:41.106Z" }, + { url = "https://files.pythonhosted.org/packages/11/b0/4a595f01a243aec8ad272b160b161c44351190c35d98d7787919d962e9e5/cramjam-2.11.0-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:193c6488bd2f514cbc0bef5c18fad61a5f9c8d059dd56edf773b3b37f0e85496", size = 2155651, upload-time = "2025-07-27T21:22:48.46Z" }, +] + +[[package]] +name = "cryptography" +version = "50.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cffi", marker = "platform_machine == 'x86_64' and platform_python_implementation != 'PyPy' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/de/41/6cbdcf9142d00fe82836fbb51e503e58088575cf7a0fe1dbff6695bf0840/cryptography-50.0.0.tar.gz", hash = "sha256:eeac2acb5a20ed25e0ad6d1df9891a520b78b404266b6d11778f25d5d691a6c9", size = 880201, upload-time = "2026-07-31T14:25:10.11Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/d9/41/029086c34d91052fc3b88bcc8056f709a7c915c7a23b235a54eb800b1c97/cryptography-50.0.0-cp311-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:06a32a980526a6ab9a4b9bf8f7385800791e2bb960903cb6b530e4817509a3b7", size = 4702130, upload-time = "2026-07-31T14:23:37.635Z" }, + { url = "https://files.pythonhosted.org/packages/6b/72/a1116d683a6d7ece94590013882515de087edf9ef0e6292aae615a44df73/cryptography-50.0.0-cp311-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:b42a28c1844fd9de8f3f7d540e36b66f3a9c83fceac7170ebc7a6a19edd9dcae", size = 4734609, upload-time = "2026-07-31T14:23:43.139Z" }, + { url = "https://files.pythonhosted.org/packages/da/3a/f05e32c99d440c9bb891ea0e36c9091891e36be5a9a87ab2ee6ea20729f6/cryptography-50.0.0-cp311-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:82148ec5bddac30b51a5b3c1945075f896fa022cb93f8e4a01e9f6ee95292c5f", size = 4734462, upload-time = "2026-07-31T14:23:50.861Z" }, + { url = "https://files.pythonhosted.org/packages/27/20/c930314a2ab476d15dec966ec87e2e9637bb02b06106b12c0396c57bb603/cryptography-50.0.0-cp311-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:7cec5b856506da6defb290f30c9ee687d5f5e8cb0bd3f6459dde43b0b4fa40ef", size = 5004179, upload-time = "2026-07-31T14:23:54.887Z" }, + { url = "https://files.pythonhosted.org/packages/85/4f/0fa8c2f4428198f15d9ff8d63400e27afbf94ce833f6108da1eb3753f945/cryptography-50.0.0-cp39-abi3-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:a91296cb61e8df6f86d0c19cc4068228da256bf59bf86049fbd821084565327f", size = 4728939, upload-time = "2026-07-31T14:24:27.994Z" }, + { url = "https://files.pythonhosted.org/packages/46/c9/f60aed34c013f317f92817b6c171c2d22a78270fa41109bd4b08af26b194/cryptography-50.0.0-cp39-abi3-manylinux_2_28_x86_64.whl", hash = "sha256:105110f43a471dbd0060b9c9516cb8a6a79233631a04cc2ba16f28323ac6e025", size = 4762647, upload-time = "2026-07-31T14:24:34.599Z" }, + { url = "https://files.pythonhosted.org/packages/9c/f8/d97f9603efda3888187bfdb893f26c41be4735c10631d05d284ee6b047c4/cryptography-50.0.0-cp39-abi3-manylinux_2_34_x86_64.whl", hash = "sha256:37fdb0d0111f1e2ff07139dfb79f1b49531f8e213c46f1163dd7642979b58c47", size = 4762400, upload-time = "2026-07-31T14:24:43.636Z" }, + { url = "https://files.pythonhosted.org/packages/d2/1a/efcfb02f91407149a0dacffffab791f7e19bf6385f63b3666dc8b5e5c9c8/cryptography-50.0.0-cp39-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:65c2c3add92b45fd0709db8594536aea39c2a67af0e27ffcf049c498501140b7", size = 5037050, upload-time = "2026-07-31T14:24:47.697Z" }, +] + +[[package]] +name = "cycler" +version = "0.12.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/a9/95/a3dbbb5028f35eafb79008e7522a75244477d2838f38cbb722248dabc2a8/cycler-0.12.1.tar.gz", hash = "sha256:88bb128f02ba341da8ef447245a9e138fae777f6a23943da4540077d3601eb1c", size = 7615, upload-time = "2023-10-07T05:32:18.335Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/05/c19819d5e3d95294a6f5947fb9b9629efb316b96de511b418c53d245aae6/cycler-0.12.1-py3-none-any.whl", hash = "sha256:85cef7cff222d8644161529808465972e51340599459b8ac3ccbac5a854e0d30", size = 8321, upload-time = "2023-10-07T05:32:16.783Z" }, +] + +[[package]] +name = "dgl" +version = "2.4.0+cu121" +source = { registry = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html" } +dependencies = [ + { name = "networkx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pandas", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "psutil", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pydantic", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pyyaml", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "scipy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "tqdm", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://data.dgl.ai/wheels/torch-2.2/cu121/dgl-2.4.0%2Bcu121-cp312-cp312-manylinux1_x86_64.whl" }, +] + +[[package]] +name = "docutils" +version = "0.23" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/39/a4/5180d9afc57e8fca05601dd652bdff19604c218814037fe90ffc7625a50a/docutils-0.23.tar.gz", hash = "sha256:746f5060322511280a1e50eb76846ed6bf2342984b2ac04dc42caa1a8d78799e", size = 2303823, upload-time = "2026-05-27T17:41:06.934Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/91/30151a39f7570f448ed84529390628a651d7f27c87d73c9b887f8189695e/docutils-0.23-py3-none-any.whl", hash = "sha256:25d013af9bf23bc1c7b2b093dff4208166c53a94786c9e447808335ef1185fea", size = 634701, upload-time = "2026-05-27T17:40:58.442Z" }, +] + +[[package]] +name = "filelock" +version = "3.32.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/64/a02e6765de08964ed371eca577870593245afc9dfac16d037de7c10d18e6/filelock-3.32.3.tar.gz", hash = "sha256:0ffa185a3540854c95caa7fa76b76cb219d907415e2c5dc9af25fd970563487f", size = 218135, upload-time = "2026-08-13T16:00:05.577Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a7/8e/50f46a9c0ce8d2861a394c1347caae037ea0431d2f67d7feb151cbc4649a/filelock-3.32.3-py3-none-any.whl", hash = "sha256:7f0ca4bcc0e181c60dbbd8aa9ab5b120ebb99e4e064e83636340056f833a1f09", size = 98901, upload-time = "2026-08-13T16:00:03.974Z" }, +] + +[[package]] +name = "flatbuffers" +version = "25.12.19" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/2d/d2a548598be01649e2d46231d151a6c56d10b964d94043a335ae56ea2d92/flatbuffers-25.12.19-py2.py3-none-any.whl", hash = "sha256:7634f50c427838bb021c2d66a3d1168e9d199b0607e6329399f04846d42e20b4", size = 26661, upload-time = "2025-12-19T23:16:13.622Z" }, +] + +[[package]] +name = "fonttools" +version = "4.63.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/84/69/c97f2c18e0db87d2c7b15da1974dace76ae938f1cfa22e2727a648b7ed43/fonttools-4.63.0.tar.gz", hash = "sha256:caeb583deeb5168e694b65cda8b4ee62abedfa66cf88488734466f2366b9c4e0", size = 3597189, upload-time = "2026-05-14T12:04:30.958Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/77/c7/2342da9830e3e9d4870305ca5d2091d2a83284f2953079b7bdd3b5e029d8/fonttools-4.63.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:58dc6bb86a78d782f00f9190ca02c119cf5bbe2807536e361e18d42019f877d8", size = 4999800, upload-time = "2026-05-14T12:03:20.161Z" }, + { url = "https://files.pythonhosted.org/packages/f2/00/3bbab338c07c71fa56269953845e92c951a61457bbbb0f1022551ea266d9/fonttools-4.63.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:27fdc65af8da6f88b9c6121c47a464cbe359fcfff7ff6fc2d37a1f395d755b78", size = 5133598, upload-time = "2026-05-14T12:03:25.168Z" }, + { url = "https://files.pythonhosted.org/packages/2c/47/c99d5268f354002ce80f8d029cd9d7d872969da1de8b93d32de4dc56d6f4/fonttools-4.63.0-py3-none-any.whl", hash = "sha256:445af2eab030a16b9171ea8bdda7ebf7d96bda2df88ee182a464252f6e05e20d", size = 1164562, upload-time = "2026-05-14T12:04:29.092Z" }, +] + +[[package]] +name = "fsspec" +version = "2026.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/00/78/f34251dadb8f3921264a1d9b8946f5e542014ee2614b285261b4e40e6775/fsspec-2026.7.0.tar.gz", hash = "sha256:c803c40f4cf860b49dea58ee3e1c33cb9c790520e233537e1340049f89b82a88", size = 317040, upload-time = "2026-07-28T16:34:51.052Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/3c/6a2bf344106328fd04963664a60b9bb6496fc25df8e962fcdc1367285fb9/fsspec-2026.7.0-py3-none-any.whl", hash = "sha256:b57ddbafedfaef7018c1ecab32aa200a9d7ca26b77965f64e48b70061249d279", size = 206583, upload-time = "2026-07-28T16:34:49.538Z" }, +] + +[[package]] +name = "gnn4colliders" +source = { editable = "." } +dependencies = [ + { name = "awkward", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "hydra-core", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "scikit-learn", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "torch", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "uproot", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] + +[package.optional-dependencies] +onnx = [ + { name = "onnx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "onnxruntime", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +root-gnn = [ + { name = "dgl", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] + +[package.dev-dependencies] +dev = [ + { name = "matplotlib", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "onnx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "onnxruntime", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pytest", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pytest-cov", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "ruff", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "twine", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] + +[package.metadata] +requires-dist = [ + { name = "awkward", specifier = ">=2.0" }, + { name = "dgl", marker = "extra == 'root-gnn'", specifier = "==2.4.0+cu121", index = "https://data.dgl.ai/wheels/torch-2.2/cu121/repo.html" }, + { name = "hydra-core", specifier = ">=1.3" }, + { name = "numpy", specifier = ">=1.24,<2" }, + { name = "onnx", marker = "extra == 'onnx'", specifier = ">=1.16,<2" }, + { name = "onnxruntime", marker = "extra == 'onnx'", specifier = ">=1.18,<2" }, + { name = "scikit-learn", specifier = ">=1.3" }, + { name = "torch", specifier = "==2.2.2", index = "https://download.pytorch.org/whl/cu121" }, + { name = "uproot", specifier = ">=5.0" }, +] +provides-extras = ["root-gnn", "onnx"] + +[package.metadata.requires-dev] +dev = [ + { name = "matplotlib", specifier = ">=3.8,<4" }, + { name = "onnx", specifier = ">=1.16,<2" }, + { name = "onnxruntime", specifier = ">=1.18,<2" }, + { name = "pytest", specifier = ">=8" }, + { name = "pytest-cov", specifier = ">=5" }, + { name = "ruff", specifier = ">=0.6" }, + { name = "twine", specifier = ">=6,<7" }, +] + +[[package]] +name = "hydra-core" +version = "1.3.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "omegaconf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/3e/e4/69a522676faf88994d93d8a5e69e0666c61cae7f73d1bbcc483222023e74/hydra_core-1.3.5.tar.gz", hash = "sha256:71c441eabbde086062045e4d3fce9e26015244f1a4ac721cf3e444c7edf10633", size = 3264337, upload-time = "2026-08-05T18:33:21.394Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9c/97/f9d463a6f3c7d0955753eca5cbbf35b596ac471dd13fe357211a53fd37be/hydra_core-1.3.5-py3-none-any.whl", hash = "sha256:a3ff35b4ea6794e4c83d993016f4bde4ac35797ebe7a08f30e83ed9341880331", size = 155768, upload-time = "2026-08-05T18:33:19.834Z" }, +] + +[[package]] +name = "id" +version = "1.6.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "urllib3", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6d/04/c2156091427636080787aac190019dc64096e56a23b7364d3c1764ee3a06/id-1.6.1.tar.gz", hash = "sha256:d0732d624fb46fd4e7bc4e5152f00214450953b9e772c182c1c22964def1a069", size = 18088, upload-time = "2026-02-04T16:19:41.26Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/77/de194443bf38daed9452139e960c632b0ef9f9a5dd9ce605fdf18ca9f1b1/id-1.6.1-py3-none-any.whl", hash = "sha256:f5ec41ed2629a508f5d0988eda142e190c9c6da971100612c4de9ad9f9b237ca", size = 14689, upload-time = "2026-02-04T16:19:40.051Z" }, +] + +[[package]] +name = "idna" +version = "3.19" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5f/f7/abb373e5757eaec4b922b92f97ec8d6d7e057cf06778247604fbc4e7c3f3/idna-3.19.tar.gz", hash = "sha256:5e0811a4383b21dc5838069f801c4fb62113b7447663d2530d2bd6e77b49bf15", size = 215237, upload-time = "2026-08-18T05:14:24.27Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/57/b0/0e52c878c53f245edd3a11020f20979b3f490f245af532c7cae3027754b5/idna-3.19-py3-none-any.whl", hash = "sha256:815e7be7a7806d54abb586dc943addc79e8b2ee16915059658cbeff4b1b43bf4", size = 68550, upload-time = "2026-08-18T05:14:22.343Z" }, +] + +[[package]] +name = "iniconfig" +version = "2.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/72/34/14ca021ce8e5dfedc35312d08ba8bf51fdd999c576889fc2c24cb97f4f10/iniconfig-2.3.0.tar.gz", hash = "sha256:c76315c77db068650d49c5b56314774a7804df16fee4402c1f19d6d15d8c4730", size = 20503, upload-time = "2025-10-18T21:55:43.219Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/cb/b1/3846dd7f199d53cb17f49cba7e651e9ce294d8497c8c150530ed11865bb8/iniconfig-2.3.0-py3-none-any.whl", hash = "sha256:f631c04d2c48c52b84d0d0549c99ff3859c98df65b3101406327ecc7d53fbf12", size = 7484, upload-time = "2025-10-18T21:55:41.639Z" }, +] + +[[package]] +name = "jaraco-classes" +version = "3.4.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/c0/ed4a27bc5571b99e3cff68f8a9fa5b56ff7df1c2251cc715a652ddd26402/jaraco.classes-3.4.0.tar.gz", hash = "sha256:47a024b51d0239c0dd8c8540c6c7f484be3b8fcf0b2d85c13825780d3b3f3acd", size = 11780, upload-time = "2024-03-31T07:27:36.643Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/66/b15ce62552d84bbfcec9a4873ab79d993a1dd4edb922cbfccae192bd5b5f/jaraco.classes-3.4.0-py3-none-any.whl", hash = "sha256:f662826b6bed8cace05e7ff873ce0f9283b5c924470fe664fff1c2f00f581790", size = 6777, upload-time = "2024-03-31T07:27:34.792Z" }, +] + +[[package]] +name = "jaraco-context" +version = "6.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/af/50/4763cd07e722bb6285316d390a164bc7e479db9d90daa769f22578f698b4/jaraco_context-6.1.2.tar.gz", hash = "sha256:f1a6c9d391e661cc5b8d39861ff077a7dc24dc23833ccee564b234b81c82dfe3", size = 16801, upload-time = "2026-03-20T22:13:33.922Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f2/58/bc8954bda5fcda97bd7c19be11b85f91973d67a706ed4a3aec33e7de22db/jaraco_context-6.1.2-py3-none-any.whl", hash = "sha256:bf8150b79a2d5d91ae48629d8b427a8f7ba0e1097dd6202a9059f29a36379535", size = 7871, upload-time = "2026-03-20T22:13:32.808Z" }, +] + +[[package]] +name = "jaraco-functools" +version = "4.6.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "more-itertools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/6c/1f/c23395957d41ccf27c4e535c3d334c4051e5395b3752057ba4cbaec35c56/jaraco_functools-4.6.0.tar.gz", hash = "sha256:880c577ec9720b3a052d5bc611fb9f2269b3d87902ef42440df443b88e443280", size = 20837, upload-time = "2026-07-14T01:28:02.544Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/02/36/ecc85bc96c273dc8a11273ed4782272975e6338d4a3e9228621175edf0e3/jaraco_functools-4.6.0-py3-none-any.whl", hash = "sha256:99e3dc0060c5cbe8fcd1cdb36258e2a65ca40f1566b2033b12abb1bb44dd3c30", size = 11677, upload-time = "2026-07-14T01:28:01.59Z" }, +] + +[[package]] +name = "jeepney" +version = "0.9.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7b/6f/357efd7602486741aa73ffc0617fb310a29b588ed0fd69c2399acbb85b0c/jeepney-0.9.0.tar.gz", hash = "sha256:cf0e9e845622b81e4a28df94c40345400256ec608d0e55bb8a3feaa9163f5732", size = 106758, upload-time = "2025-02-27T18:51:01.684Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b2/a3/e137168c9c44d18eff0376253da9f1e9234d0239e0ee230d2fee6cea8e55/jeepney-0.9.0-py3-none-any.whl", hash = "sha256:97e5714520c16fc0a45695e5365a2e11b81ea79bba796e26f9f1d178cb182683", size = 49010, upload-time = "2025-02-27T18:51:00.104Z" }, +] + +[[package]] +name = "jinja2" +version = "3.1.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markupsafe", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115, upload-time = "2025-03-05T20:05:02.478Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899, upload-time = "2025-03-05T20:05:00.369Z" }, +] + +[[package]] +name = "joblib" +version = "1.5.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/41/f2/d34e8b3a08a9cc79a50b2208a93dce981fe615b64d5a4d4abee421d898df/joblib-1.5.3.tar.gz", hash = "sha256:8561a3269e6801106863fd0d6d84bb737be9e7631e33aaed3fb9ce5953688da3", size = 331603, upload-time = "2025-12-15T08:41:46.427Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7b/91/984aca2ec129e2757d1e4e3c81c3fcda9d0f85b74670a094cc443d9ee949/joblib-1.5.3-py3-none-any.whl", hash = "sha256:5fc3c5039fc5ca8c0276333a188bbd59d6b7ab37fe6632daa76bc7f9ec18e713", size = 309071, upload-time = "2025-12-15T08:41:44.973Z" }, +] + +[[package]] +name = "keyring" +version = "25.7.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "jaraco-classes", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "jaraco-context", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "jaraco-functools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "jeepney", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "secretstorage", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/43/4b/674af6ef2f97d56f0ab5153bf0bfa28ccb6c3ed4d1babf4305449668807b/keyring-25.7.0.tar.gz", hash = "sha256:fe01bd85eb3f8fb3dd0405defdeac9a5b4f6f0439edbb3149577f244a2e8245b", size = 63516, upload-time = "2025-11-16T16:26:09.482Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/81/db/e655086b7f3a705df045bf0933bdd9c2f79bb3c97bfef1384598bb79a217/keyring-25.7.0-py3-none-any.whl", hash = "sha256:be4a0b195f149690c166e850609a477c532ddbfbaed96a404d4e43f8d5e2689f", size = 39160, upload-time = "2025-11-16T16:26:08.402Z" }, +] + +[[package]] +name = "kiwisolver" +version = "1.5.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d0/67/9c61eccb13f0bdca9307614e782fec49ffdde0f7a2314935d489fa93cd9c/kiwisolver-1.5.0.tar.gz", hash = "sha256:d4193f3d9dc3f6f79aaed0e5637f45d98850ebf01f7ca20e69457f3e8946b66a", size = 103482, upload-time = "2026-03-09T13:15:53.382Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c4/13/680c54afe3e65767bed7ec1a15571e1a2f1257128733851ade24abcefbcc/kiwisolver-1.5.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:bb5136fb5352d3f422df33f0c879a1b0c204004324150cc3b5e3c4f310c9049f", size = 1477934, upload-time = "2026-03-09T13:13:27.166Z" }, + { url = "https://files.pythonhosted.org/packages/e4/34/8aefdd0be9cfd00a44509251ba864f5caf2991e36772e61c408007e7f417/kiwisolver-1.5.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:1d9daea4ea6b9be74fe2f01f7fbade8d6ffab263e781274cffca0dba9be9eec9", size = 2294947, upload-time = "2026-03-09T13:13:43.343Z" }, + { url = "https://files.pythonhosted.org/packages/fa/06/7399a607f434119c6e1fdc8ec89a8d51ccccadf3341dee4ead6bd14caaf5/kiwisolver-1.5.0-graalpy312-graalpy250_312_native-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:c31c13da98624f957b0fb1b5bae5383b2333c2c3f6793d9825dd5ce79b525cb7", size = 194295, upload-time = "2026-03-09T13:15:38.22Z" }, +] + +[[package]] +name = "markdown-it-py" +version = "4.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mdurl", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/06/ff/7841249c247aa650a76b9ee4bbaeae59370dc8bfd2f6c01f3630c35eb134/markdown_it_py-4.2.0.tar.gz", hash = "sha256:04a21681d6fbb623de53f6f364d352309d4094dd4194040a10fd51833e418d49", size = 82454, upload-time = "2026-05-07T12:08:28.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/81/4da04ced5a082363ecfa159c010d200ecbd959ae410c10c0264a38cac0f5/markdown_it_py-4.2.0-py3-none-any.whl", hash = "sha256:9f7ebbcd14fe59494226453aed97c1070d83f8d24b6fc3a3bcf9a38092641c4a", size = 91687, upload-time = "2026-05-07T12:08:27.182Z" }, +] + +[[package]] +name = "markupsafe" +version = "3.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7e/99/7690b6d4034fffd95959cbe0c02de8deb3098cc577c67bb6a24fe5d7caa7/markupsafe-3.0.3.tar.gz", hash = "sha256:722695808f4b6457b320fdc131280796bdceb04ab50fe1795cd540799ebe1698", size = 80313, upload-time = "2025-09-27T18:37:40.426Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3c/2e/8d0c2ab90a8c1d9a24f0399058ab8519a3279d1bd4289511d74e909f060e/markupsafe-3.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d6dd0be5b5b189d31db7cda48b91d7e0a9795f31430b7f271219ab30f1d3ac9d", size = 22947, upload-time = "2025-09-27T18:36:33.86Z" }, + { url = "https://files.pythonhosted.org/packages/89/e0/4486f11e51bbba8b0c041098859e869e304d1c261e59244baa3d295d47b7/markupsafe-3.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:77f0643abe7495da77fb436f50f8dab76dbc6e5fd25d39589a0f1fe6548bfa2b", size = 23015, upload-time = "2025-09-27T18:36:37.868Z" }, +] + +[[package]] +name = "matplotlib" +version = "3.11.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "contourpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "cycler", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fonttools", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "kiwisolver", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pillow", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pyparsing", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "python-dateutil", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/49/64/f9a391af28f518b11ad45a8a712353c94a0aefce09d3703200e5c54b610a/matplotlib-3.11.1.tar.gz", hash = "sha256:69647db5746941c793d6e445a4cd349323ffb87d9cc958c2ad84a659b4832d30", size = 32612045, upload-time = "2026-07-18T03:39:46.63Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c9/cf/b4ad2cc81b6672ea29ea04e64e350a9f9b493b0908ccd884c67eeff8f7b2/matplotlib-3.11.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:ba8f811b8ddfac493734d6af0b2dff96919d0c28ca0d641858dab4262777c6ea", size = 10035615, upload-time = "2026-07-18T03:38:24.315Z" }, + { url = "https://files.pythonhosted.org/packages/88/eb/799612d0f8cd3e816a10fec59329fca52cd2353264df80378dfc541ae855/matplotlib-3.11.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:8b14eb22961fe865efb0e4ff167e333e428908b00115a8d800ccb65ee108e481", size = 10927532, upload-time = "2026-07-18T03:38:28.532Z" }, +] + +[[package]] +name = "mdurl" +version = "0.1.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/d6/54/cfe61301667036ec958cb99bd3efefba235e65cdeb9c84d24a8293ba1d90/mdurl-0.1.2.tar.gz", hash = "sha256:bb413d29f5eea38f31dd4754dd7377d4465116fb207585f97bf925588687c1ba", size = 8729, upload-time = "2022-08-14T12:40:10.846Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b3/38/89ba8ad64ae25be8de66a6d463314cf1eb366222074cfda9ee839c56a4b4/mdurl-0.1.2-py3-none-any.whl", hash = "sha256:84008a41e51615a49fc9966191ff91509e3c40b939176e643fd50a5c2196b8f8", size = 9979, upload-time = "2022-08-14T12:40:09.779Z" }, +] + +[[package]] +name = "ml-dtypes" +version = "0.5.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/0e/4a/c27b42ed9b1c7d13d9ba8b6905dece787d6259152f2309338aed29b2447b/ml_dtypes-0.5.4.tar.gz", hash = "sha256:8ab06a50fb9bf9666dd0fe5dfb4676fa2b0ac0f31ecff72a6c3af8e22c063453", size = 692314, upload-time = "2025-11-17T22:32:31.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/cb/28ce52eb94390dda42599c98ea0204d74799e4d8047a0eb559b6fd648056/ml_dtypes-0.5.4-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9ad459e99793fa6e13bd5b7e6792c8f9190b4e5a1b45c63aba14a4d0a7f1d5ff", size = 5009002, upload-time = "2025-11-17T22:31:52.001Z" }, +] + +[[package]] +name = "more-itertools" +version = "11.1.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/de/1d/f4da6f02cdffe04d6362210b807146a26044c88d839208aec273bb0d9184/more_itertools-11.1.0.tar.gz", hash = "sha256:48e8f4d9e7e5878571ecf6f2b4e57634f93cd474cc8cfbd2376f2d11b396e30d", size = 145772, upload-time = "2026-05-22T14:14:29.909Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e8/3d/1087453384dbde46a8c7f9356eead2c58be8a7bf156bca40243377c85715/more_itertools-11.1.0-py3-none-any.whl", hash = "sha256:4b65538ae22f6fed0ce4874efd317463a7489796a0939fa66824dd542125a192", size = 72226, upload-time = "2026-05-22T14:14:28.824Z" }, +] + +[[package]] +name = "mpmath" +version = "1.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/e0/47/dd32fa426cc72114383ac549964eecb20ecfd886d1e5ccf5340b55b02f57/mpmath-1.3.0.tar.gz", hash = "sha256:7a28eb2a9774d00c7bc92411c19a89209d5da7c4c9a9e227be8330a23a25b91f", size = 508106, upload-time = "2023-03-07T16:47:11.061Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" }, +] + +[[package]] +name = "narwhals" +version = "2.24.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/2b/1d/58946e5aab18393e793bd4add6985b95d0e01c3a2d832f38f54468b10dcd/narwhals-2.24.0.tar.gz", hash = "sha256:b5c0f684ccd9d7475b564111e319a4964abcf2baf79d3cf6b1003d06ac9b828d", size = 661143, upload-time = "2026-07-13T10:49:19.086Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/85/a5bfaebfd305ac18b57b0854d74e37e586809061a91fda62f0bd50c8518e/narwhals-2.24.0-py3-none-any.whl", hash = "sha256:42fdedf44e5b2ca7505630d45b4ac3058f38d8485cba9fe1652ca23152df7489", size = 461030, upload-time = "2026-07-13T10:49:17.571Z" }, +] + +[[package]] +name = "networkx" +version = "3.6.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/6a/51/63fe664f3908c97be9d2e4f1158eb633317598cfa6e1fc14af5383f17512/networkx-3.6.1.tar.gz", hash = "sha256:26b7c357accc0c8cde558ad486283728b65b6a95d85ee1cd66bafab4c8168509", size = 2517025, upload-time = "2025-12-08T17:02:39.908Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9e/c9/b2622292ea83fbb4ec318f5b9ab867d0a28ab43c5717bb85b0a5f6b3b0a4/networkx-3.6.1-py3-none-any.whl", hash = "sha256:d47fbf302e7d9cbbb9e2555a0d267983d2aa476bac30e90dfbe5669bd57f3762", size = 2068504, upload-time = "2025-12-08T17:02:38.159Z" }, +] + +[[package]] +name = "nh3" +version = "0.3.6" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/5e/1b/ef84624f14954d270f74060a19fc550dd4f06656399447569afb584d8c06/nh3-0.3.6.tar.gz", hash = "sha256:f3736c9dd3d1856f80cd031715b84ca75cda2bbb1ac802c3da26bfce590838d7", size = 24684, upload-time = "2026-06-22T00:47:02.008Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5b/67/314f6151bad77a93d751978a344033e1fc890822f05f0416079338e34231/nh3-0.3.6-cp38-abi3-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:905f877dc66dd7aea4a76e54bcb26acb5ff8216f720c0017ccf63e0e6035698e", size = 806699, upload-time = "2026-06-22T00:46:45.99Z" }, + { url = "https://files.pythonhosted.org/packages/17/0c/6cdb5ee1e127be50dc8391e54bddc1f64e87bf4bfad0c55633320e2e02db/nh3-0.3.6-cp38-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:36d06341bd501240d320f5942481ed5e6846136b666e1ba4faf802b78ebc875f", size = 1033829, upload-time = "2026-06-22T00:46:56.258Z" }, +] + +[[package]] +name = "numpy" +version = "1.26.4" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/65/6e/09db70a523a96d25e115e71cc56a6f9031e7b8cd166c1ac8438307c14058/numpy-1.26.4.tar.gz", hash = "sha256:2a02aba9ed12e4ac4eb3ea9421c420301a0c6460d9830d74a9df87efa4912010", size = 15786129, upload-time = "2024-02-06T00:26:44.495Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0f/50/de23fde84e45f5c4fda2488c759b69990fd4512387a8632860f3ac9cd225/numpy-1.26.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:675d61ffbfa78604709862923189bad94014bef562cc35cf61d3a07bba02a7ed", size = 17950613, upload-time = "2024-02-05T23:56:56.054Z" }, + { url = "https://files.pythonhosted.org/packages/76/8c/2ba3902e1a0fc1c74962ea9bb33a534bb05984ad7ff9515bf8d07527cadd/numpy-1.26.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:1dda2e7b4ec9dd512f84935c5f126c8bd8b9f2fc001e9f54af255e8c5f16b0e0", size = 17786643, upload-time = "2024-02-05T23:57:56.585Z" }, +] + +[[package]] +name = "nvidia-cublas-cu12" +version = "12.1.3.1" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/37/6d/121efd7382d5b0284239f4ab1fc1590d86d34ed4a4a2fdb13b30ca8e5740/nvidia_cublas_cu12-12.1.3.1-py3-none-manylinux1_x86_64.whl", hash = "sha256:ee53ccca76a6fc08fb9701aa95b6ceb242cdaab118c3bb152af4e579af792728", size = 410594774, upload-time = "2023-04-19T15:50:03.519Z" }, +] + +[[package]] +name = "nvidia-cuda-cupti-cu12" +version = "12.1.105" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7e/00/6b218edd739ecfc60524e585ba8e6b00554dd908de2c9c66c1af3e44e18d/nvidia_cuda_cupti_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:e54fde3983165c624cb79254ae9818a456eb6e87a7fd4d56a2352c24ee542d7e", size = 14109015, upload-time = "2023-04-19T15:47:32.502Z" }, +] + +[[package]] +name = "nvidia-cuda-nvrtc-cu12" +version = "12.1.105" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b6/9f/c64c03f49d6fbc56196664d05dba14e3a561038a81a638eeb47f4d4cfd48/nvidia_cuda_nvrtc_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:339b385f50c309763ca65456ec75e17bbefcbbf2893f462cb8b90584cd27a1c2", size = 23671734, upload-time = "2023-04-19T15:48:32.42Z" }, +] + +[[package]] +name = "nvidia-cuda-runtime-cu12" +version = "12.1.105" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/eb/d5/c68b1d2cdfcc59e72e8a5949a37ddb22ae6cade80cd4a57a84d4c8b55472/nvidia_cuda_runtime_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:6e258468ddf5796e25f1dc591a31029fa317d97a0a94ed93468fc86301d61e40", size = 823596, upload-time = "2023-04-19T15:47:22.471Z" }, +] + +[[package]] +name = "nvidia-cudnn-cu12" +version = "8.9.2.26" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/74/a2e2be7fb83aaedec84f391f082cf765dfb635e7caa9b49065f73e4835d8/nvidia_cudnn_cu12-8.9.2.26-py3-none-manylinux1_x86_64.whl", hash = "sha256:5ccb288774fdfb07a7e7025ffec286971c06d8d7b4fb162525334616d7629ff9", size = 731725872, upload-time = "2023-06-01T19:24:57.328Z" }, +] + +[[package]] +name = "nvidia-cufft-cu12" +version = "11.0.2.54" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/86/94/eb540db023ce1d162e7bea9f8f5aa781d57c65aed513c33ee9a5123ead4d/nvidia_cufft_cu12-11.0.2.54-py3-none-manylinux1_x86_64.whl", hash = "sha256:794e3948a1aa71fd817c3775866943936774d1c14e7628c74f6f7417224cdf56", size = 121635161, upload-time = "2023-04-19T15:50:46Z" }, +] + +[[package]] +name = "nvidia-curand-cu12" +version = "10.3.2.106" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/44/31/4890b1c9abc496303412947fc7dcea3d14861720642b49e8ceed89636705/nvidia_curand_cu12-10.3.2.106-py3-none-manylinux1_x86_64.whl", hash = "sha256:9d264c5036dde4e64f1de8c50ae753237c12e0b1348738169cd0f8a536c0e1e0", size = 56467784, upload-time = "2023-04-19T15:51:04.804Z" }, +] + +[[package]] +name = "nvidia-cusolver-cu12" +version = "11.4.5.107" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/bc/1d/8de1e5c67099015c834315e333911273a8c6aaba78923dd1d1e25fc5f217/nvidia_cusolver_cu12-11.4.5.107-py3-none-manylinux1_x86_64.whl", hash = "sha256:8a7ec542f0412294b15072fa7dab71d31334014a69f953004ea7a118206fe0dd", size = 124161928, upload-time = "2023-04-19T15:51:25.781Z" }, +] + +[[package]] +name = "nvidia-cusparse-cu12" +version = "12.1.0.106" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "nvidia-nvjitlink-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/65/5b/cfaeebf25cd9fdec14338ccb16f6b2c4c7fa9163aefcf057d86b9cc248bb/nvidia_cusparse_cu12-12.1.0.106-py3-none-manylinux1_x86_64.whl", hash = "sha256:f3b50f42cf363f86ab21f720998517a659a48131e8d538dc02f8768237bd884c", size = 195958278, upload-time = "2023-04-19T15:51:49.939Z" }, +] + +[[package]] +name = "nvidia-nccl-cu12" +version = "2.19.3" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/38/00/d0d4e48aef772ad5aebcf70b73028f88db6e5640b36c38e90445b7a57c45/nvidia_nccl_cu12-2.19.3-py3-none-manylinux1_x86_64.whl", hash = "sha256:a9734707a2c96443331c1e48c717024aa6678a0e2a4cb66b2c364d18cee6b48d", size = 165987969, upload-time = "2023-10-24T16:16:24.789Z" }, +] + +[[package]] +name = "nvidia-nvjitlink-cu12" +version = "12.9.86" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/46/0c/c75bbfb967457a0b7670b8ad267bfc4fffdf341c074e0a80db06c24ccfd4/nvidia_nvjitlink_cu12-12.9.86-py3-none-manylinux2010_x86_64.manylinux_2_12_x86_64.whl", hash = "sha256:e3f1171dbdc83c5932a45f0f4c99180a70de9bd2718c1ab77d14104f6d7147f9", size = 39748338, upload-time = "2025-06-05T20:10:25.613Z" }, +] + +[[package]] +name = "nvidia-nvtx-cu12" +version = "12.1.105" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/da/d3/8057f0587683ed2fcd4dbfbdfdfa807b9160b809976099d36b8f60d08f03/nvidia_nvtx_cu12-12.1.105-py3-none-manylinux1_x86_64.whl", hash = "sha256:dc21cf308ca5691e7c04d962e213f8a4aa9bbfa23d95412f452254c2caeb09e5", size = 99138, upload-time = "2023-04-19T15:48:43.556Z" }, +] + +[[package]] +name = "omegaconf" +version = "2.3.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "antlr4-python3-runtime", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pyyaml", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ce/3d/e4b57b8d9008c6ebe0d5eff901f91d5700cf7bdb8c8863df817463a7fd5e/omegaconf-2.3.1.tar.gz", hash = "sha256:e5e7de64aeebeddaf8e6d3f7a783b32ac2a01c0fbd9c878012caecb891a1f42a", size = 3298472, upload-time = "2026-06-11T05:05:12.885Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a4/0e/152509871bf30df6fc38569f52a2db9b55dd41aae957adae50a053ac7778/omegaconf-2.3.1-py3-none-any.whl", hash = "sha256:3d701d14e9a8828f1edd28bb70b725908b34277cdd72cf7d6a83f94dadc6b6a0", size = 79502, upload-time = "2026-06-11T05:05:09.954Z" }, +] + +[[package]] +name = "onnx" +version = "1.22.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "ml-dtypes", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "protobuf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/04/19/8ea73a64b368b75fe339771a20a02bc61ea1f551484c9e3d9d0bfbd0450f/onnx-1.22.0.tar.gz", hash = "sha256:ef40c0aaf0b643857ea9306fc7eddce17eaf9fb0407e4801f1fc5758443a38e0", size = 12024721, upload-time = "2026-06-15T12:50:05.354Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/09/a6/bd32357e6cc1ecb473afd78193d7231724f284435d2db25696ecfaaa1503/onnx-1.22.0-cp312-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:955e02e1f6d385b53d52f9cd7b9cdf5caf417c300bcfe3c64c6d542be763845b", size = 19106514, upload-time = "2026-06-15T12:49:37.424Z" }, + { url = "https://files.pythonhosted.org/packages/5a/9d/3af461ac6c714b8b369cb71499659932f4f12cfb066250b62f7567c3d530/onnx-1.22.0-cp312-abi3-pyemscripten_2025_0_wasm32.whl", hash = "sha256:82e9f27fc1223cb06d68a56bed6f9d3caf3d0dad1b61bce45006d529b15bd94c", size = 16966387, upload-time = "2026-06-15T12:49:40.918Z" }, +] + +[[package]] +name = "onnxruntime" +version = "1.29.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "flatbuffers", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "protobuf", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/96/eb/e6968f5e41aac3125f2ff5708855f09cb0b70d85ed3115b625b0b58305ba/onnxruntime-1.29.0-cp312-cp312-manylinux_2_28_x86_64.whl", hash = "sha256:2b80d8c7ec2cc7438e4da3760b88c24568cba72c9ace96d668800a6c79419acb", size = 23136745, upload-time = "2026-08-17T22:53:53.92Z" }, +] + +[[package]] +name = "packaging" +version = "26.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/7d/fa/3944b40b07da9ce895c0e6303a5ab7d53da063554f534556b134a54d6093/packaging-26.3.tar.gz", hash = "sha256:94edc256424af38762eb31306eed28beb9f0efc50a8837492c9d6fd6004aed79", size = 313412, upload-time = "2026-08-04T18:15:28.737Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/34/ba1c580383c9eada3711951fef0795c80b829a078d72188184bcab9dd527/packaging-26.3-py3-none-any.whl", hash = "sha256:d7193f7c8e4e93f444fde0262bf90af30e16fa0ad0ad44cb553c87339b23cd1c", size = 129956, upload-time = "2026-08-04T18:15:27.159Z" }, +] + +[[package]] +name = "pandas" +version = "3.0.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "python-dateutil", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/be/4f/5f3422a2afec5ffc46308b79e53291365a93748b498ac2e58bead0197916/pandas-3.0.5.tar.gz", hash = "sha256:dca3734d6ab7c906e6730f0788b0a1dbb9f2467731f9711f77995c8e9d62d712", size = 4658219, upload-time = "2026-07-22T22:19:28.819Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/12/98/e976540c1addf70442be7842a18cf70884a964abbf69442504f4d2939989/pandas-3.0.5-cp312-cp312-manylinux_2_24_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d373ce03ffd84010ed9839fa73672a9c8256990532e158440c0085db7d914b34", size = 10998091, upload-time = "2026-07-22T22:18:15.209Z" }, + { url = "https://files.pythonhosted.org/packages/9d/e2/bd9c98ad2df7b38bde002adde4cdf353519da51881634323b126c55997f9/pandas-3.0.5-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:a5ad3b02ed6bc7d7ae9b70804b2c6aa31827489d150f8e623ce82491b82085d7", size = 12060559, upload-time = "2026-07-22T22:18:20.147Z" }, + { url = "https://files.pythonhosted.org/packages/f3/9a/ffbd852d58bd74a617fe2f8ee6a58a96982271ce41cf981eab22190b4a4b/pandas-3.0.5-cp312-cp312-pyemscripten_2024_0_wasm32.whl", hash = "sha256:b2acb4650527eec6822c3dadb2b771277b65e7dae7a267d4bccf65fd1bb3fbce", size = 7197652, upload-time = "2026-07-22T22:18:22.502Z" }, +] + +[[package]] +name = "pillow" +version = "12.3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1c/3d/bb7fca845737cf9d7dbde16ed1843984665ff2e0a518f5db43e77ec540b9/pillow-12.3.0.tar.gz", hash = "sha256:3b8182a766685eaa002637e28b4ec8d6b18819a0c71f579bf0dbaa5830297cce", size = 47025035, upload-time = "2026-07-01T11:56:38.965Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/84/21/a35af28dcc61f37ed850a2d64c65c701321dfbf25085e469d5559360cbbf/pillow-12.3.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:78cb2c6865a35ab8ff8b75fd122f6033b92a62c82801110e48ddd6c936a45d91", size = 6940830, upload-time = "2026-07-01T11:54:13.732Z" }, + { url = "https://files.pythonhosted.org/packages/1d/72/cf78ac9780bb93c28328f408973845a309d4d145041665f734572ced1b52/pillow-12.3.0-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:0dd2064cbc55aaec028ef5fbb60fa47bb6c3e7918e07ff17935284b227a9d2df", size = 7052934, upload-time = "2026-07-01T11:54:17.721Z" }, +] + +[[package]] +name = "pluggy" +version = "1.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f9/e2/3e91f31a7d2b083fe6ef3fa267035b518369d9511ffab804f839851d2779/pluggy-1.6.0.tar.gz", hash = "sha256:7dcc130b76258d33b90f61b658791dede3486c3e6bfb003ee5c9bfb396dd22f3", size = 69412, upload-time = "2025-05-15T12:30:07.975Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/54/20/4d324d65cc6d9205fabedc306948156824eb9f0ee1633355a8f7ec5c66bf/pluggy-1.6.0-py3-none-any.whl", hash = "sha256:e920276dd6813095e9377c0bc5566d94c932c33b27a3e3945d8389c374dd4746", size = 20538, upload-time = "2025-05-15T12:30:06.134Z" }, +] + +[[package]] +name = "protobuf" +version = "7.35.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/da/01/9ef0afd7999eb9badb3a768b4aedd78c86d4c65cfaf1958ab276199e76b4/protobuf-7.35.1.tar.gz", hash = "sha256:ce115a26fe0c39a2c29973d914d327e516a6455464489fe3cd1e51a1b354f81a", size = 458717, upload-time = "2026-06-11T21:55:40.257Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e4/be/5b3cfe508bfab6761414ff944e3366eb13be4fd71efcd69450f89ba39f43/protobuf-7.35.1-cp310-abi3-manylinux2014_x86_64.whl", hash = "sha256:74758715c53d7158fb76caf4f0cfdacc5329a4b1bb994f865d6cf302d413a1c4", size = 327130, upload-time = "2026-06-11T21:55:35.921Z" }, + { url = "https://files.pythonhosted.org/packages/19/c7/5f7c636ec43e0c545e28d1f1db71990108306f7bdcb89f069ba97e428e7f/protobuf-7.35.1-py3-none-any.whl", hash = "sha256:4bc97768d8fe4ad6743c8a19403e314511ed9f6d13205b687e52421c023ac1b9", size = 171659, upload-time = "2026-06-11T21:55:39.155Z" }, +] + +[[package]] +name = "psutil" +version = "7.2.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/aa/c6/d1ddf4abb55e93cebc4f2ed8b5d6dbad109ecb8d63748dd2b20ab5e57ebe/psutil-7.2.2.tar.gz", hash = "sha256:0746f5f8d406af344fd547f1c8daa5f5c33dbc293bb8d6a16d80b4bb88f59372", size = 493740, upload-time = "2026-01-28T18:14:54.428Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b5/70/5d8df3b09e25bce090399cf48e452d25c935ab72dad19406c77f4e828045/psutil-7.2.2-cp36-abi3-manylinux2010_x86_64.manylinux_2_12_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:076a2d2f923fd4821644f5ba89f059523da90dc9014e85f8e45a5774ca5bc6f9", size = 155560, upload-time = "2026-01-28T18:15:25.976Z" }, + { url = "https://files.pythonhosted.org/packages/04/78/0acd37ca84ce3ddffaa92ef0f571e073faa6d8ff1f0559ab1272188ea2be/psutil-7.2.2-cp36-abi3-musllinux_1_2_x86_64.whl", hash = "sha256:b58fabe35e80b264a4e3bb23e6b96f9e45a3df7fb7eed419ac0e5947c61e47cc", size = 148266, upload-time = "2026-01-28T18:15:31.597Z" }, +] + +[[package]] +name = "pycparser" +version = "3.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/1b/7d/92392ff7815c21062bea51aa7b87d45576f649f16458d78b7cf94b9ab2e6/pycparser-3.0.tar.gz", hash = "sha256:600f49d217304a5902ac3c37e1281c9fe94e4d0489de643a9504c5cdfdfc6b29", size = 103492, upload-time = "2026-01-21T14:26:51.89Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/0c/c3/44f3fbbfa403ea2a7c779186dc20772604442dde72947e7d01069cbe98e3/pycparser-3.0-py3-none-any.whl", hash = "sha256:b727414169a36b7d524c1c3e31839a521725078d7b2ff038656844266160a992", size = 48172, upload-time = "2026-01-21T14:26:50.693Z" }, +] + +[[package]] +name = "pydantic" +version = "2.13.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "annotated-types", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pydantic-core", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-inspection", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/18/a5/b60d21ac674192f8ab0ba4e9fd860690f9b4a6e51ca5df118733b487d8d6/pydantic-2.13.4.tar.gz", hash = "sha256:c40756b57adaa8b1efeeced5c196f3f3b7c435f90e84ea7f443901bec8099ef6", size = 844775, upload-time = "2026-05-06T13:43:05.343Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/fd/7b/122376b1fd3c62c1ed9dc80c931ace4844b3c55407b6fb2d199377c9736f/pydantic-2.13.4-py3-none-any.whl", hash = "sha256:45a282cde31d808236fd7ea9d919b128653c8b38b393d1c4ab335c62924d9aba", size = 472262, upload-time = "2026-05-06T13:43:02.641Z" }, +] + +[[package]] +name = "pydantic-core" +version = "2.46.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/9d/56/921726b776ace8d8f5db44c4ef961006580d91dc52b803c489fafd1aa249/pydantic_core-2.46.4.tar.gz", hash = "sha256:62f875393d7f270851f20523dd2e29f082bcc82292d66db2b64ea71f64b6e1c1", size = 471464, upload-time = "2026-05-06T13:37:06.98Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5f/97/2aab507d3d00ca626e8e57c1eac6a79e4e5fbcc63eb99733ff55d1717f65/pydantic_core-2.46.4-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:926c9541b14b12b1681dca8a0b75feb510b06c6341b70a8e500c2fdcff837cce", size = 2094516, upload-time = "2026-05-06T13:39:10.577Z" }, + { url = "https://files.pythonhosted.org/packages/d0/cf/c873d91679f3a30bcf5e7ac280ce5573483e72295307685120d0d5ad3416/pydantic_core-2.46.4-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:fbdb89b3e1c94a30cc5edfce477c6e6a5dc4d8f84665b455c27582f211a1c72c", size = 2374464, upload-time = "2026-05-06T13:38:06.976Z" }, + { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" }, +] + +[[package]] +name = "pygments" +version = "2.21.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/49/2e/ced460408999b33da6b31b0021b0f37d329e202d4169aeb164493778f25b/pygments-2.21.0.tar.gz", hash = "sha256:610ca751c9bc2492b38eb9a38a7fbc93edbbb2d7182edaf34e66ae493dee5c8c", size = 5005329, upload-time = "2026-08-17T08:02:48.824Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/46/17f022dd3e953bf20a04a028a21ec746d942f8d2af30fa0f124fa0e6a684/pygments-2.21.0-py3-none-any.whl", hash = "sha256:2363c69b61c4a97c838da3b130dcd6468f4848992b21a82f2a63ec34377137d9", size = 1250147, upload-time = "2026-08-17T08:02:44.912Z" }, +] + +[[package]] +name = "pyparsing" +version = "3.3.2" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f3/91/9c6ee907786a473bf81c5f53cf703ba0957b23ab84c264080fb5a450416f/pyparsing-3.3.2.tar.gz", hash = "sha256:c777f4d763f140633dcb6d8a3eda953bf7a214dc4eff598413c070bcdc117cbc", size = 6851574, upload-time = "2026-01-21T03:57:59.36Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/10/bd/c038d7cc38edc1aa5bf91ab8068b63d4308c66c4c8bb3cbba7dfbc049f9c/pyparsing-3.3.2-py3-none-any.whl", hash = "sha256:850ba148bd908d7e2411587e247a1e4f0327839c40e2e5e6d05a007ecc69911d", size = 122781, upload-time = "2026-01-21T03:57:55.912Z" }, +] + +[[package]] +name = "pytest" +version = "9.1.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "iniconfig", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pluggy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pygments", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e4/47/b9efed96c114afcfa3c9d3fe98a76a1d14c74a9e266d397cf6eb64be5e01/pytest-9.1.1.tar.gz", hash = "sha256:1088fbde8f2b49d95a549a195707afa7a76a3ce9bcadc26b6d71f0ffda5fe313", size = 1636369, upload-time = "2026-06-19T10:58:32.857Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/24/25/1de2678b631f5a49215c6c96fff41ba892b0a34df68d6d80292b1b48aa7f/pytest-9.1.1-py3-none-any.whl", hash = "sha256:37a86b45efb9a47a61a36449063e8e18d0cab3161329fc099eb21783169c4f0c", size = 386536, upload-time = "2026-06-19T10:58:31.347Z" }, +] + +[[package]] +name = "pytest-cov" +version = "7.1.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "coverage", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pluggy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pytest", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/b1/51/a849f96e117386044471c8ec2bd6cfebacda285da9525c9106aeb28da671/pytest_cov-7.1.0.tar.gz", hash = "sha256:30674f2b5f6351aa09702a9c8c364f6a01c27aae0c1366ae8016160d1efc56b2", size = 55592, upload-time = "2026-03-21T20:11:16.284Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/9d/7a/d968e294073affff457b041c2be9868a40c1c71f4a35fcc1e45e5493067b/pytest_cov-7.1.0-py3-none-any.whl", hash = "sha256:a0461110b7865f9a271aa1b51e516c9a95de9d696734a2f71e3e78f46e1d4678", size = 22876, upload-time = "2026-03-21T20:11:14.438Z" }, +] + +[[package]] +name = "python-dateutil" +version = "2.9.0.post0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "six", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/66/c0/0c8b6ad9f17a802ee498c46e004a0eb49bc148f2fd230864601a86dcf6db/python-dateutil-2.9.0.post0.tar.gz", hash = "sha256:37dd54208da7e1cd875388217d5e00ebd4179249f90fb72437e91a35459a0ad3", size = 342432, upload-time = "2024-03-01T18:36:20.211Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ec/57/56b9bcc3c9c6a792fcbaf139543cee77261f3651ca9da0c93f5c1221264b/python_dateutil-2.9.0.post0-py2.py3-none-any.whl", hash = "sha256:a8b2bc7bffae282281c8140a97d3aa9c14da0b136dfe83f850eea9a5f7470427", size = 229892, upload-time = "2024-03-01T18:36:18.57Z" }, +] + +[[package]] +name = "pyyaml" +version = "6.0.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/05/8e/961c0007c59b8dd7729d542c61a4d537767a59645b82a0b521206e1e25c2/pyyaml-6.0.3.tar.gz", hash = "sha256:d76623373421df22fb4cf8817020cbb7ef15c725b9d5e45f17e189bfc384190f", size = 130960, upload-time = "2025-09-25T21:33:16.546Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/8b/9d/b3589d3877982d4f2329302ef98a8026e7f4443c765c46cfecc8858c6b4b/pyyaml-6.0.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ba1cc08a7ccde2d2ec775841541641e4548226580ab850948cbfda66a1befcdc", size = 807870, upload-time = "2025-09-25T21:32:16.431Z" }, + { url = "https://files.pythonhosted.org/packages/be/8e/98435a21d1d4b46590d5459a22d88128103f8da4c2d4cb8f14f2a96504e1/pyyaml-6.0.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:41715c910c881bc081f1e8872880d3c650acf13dfa8214bad49ed4cede7c34ea", size = 790181, upload-time = "2025-09-25T21:32:18.834Z" }, +] + +[[package]] +name = "readme-renderer" +version = "45.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "docutils", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nh3", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pygments", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/02/51/d3a6ea424652c60f05600d8c2e01a55c913755e7cdad64afabbd1aa16f44/readme_renderer-45.0.tar.gz", hash = "sha256:030a8fac74904f8fba11ad1bb6964e3f76e896dc7e5e71f16af190c9056696d1", size = 36172, upload-time = "2026-06-09T21:05:17.37Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/97/1b/295bf2fa3e740131778065e5ffa2c481f0e7210182d408e9a2c244ff5b0c/readme_renderer-45.0-py3-none-any.whl", hash = "sha256:3385ed220117104a2bceb4a9dac8c5fdf6d1f96890d7ea2a9c7174fd5c84091f", size = 14134, upload-time = "2026-06-09T21:05:15.85Z" }, +] + +[[package]] +name = "requests" +version = "2.34.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "certifi", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "charset-normalizer", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "idna", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "urllib3", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/ac/c3/e2a2b89f2d3e2179abd6d00ebd70bff6273f37fb3e0cc209f48b39d00cbf/requests-2.34.2.tar.gz", hash = "sha256:f288924cae4e29463698d6d60bc6a4da69c89185ad1e0bcc4104f584e960b9ed", size = 142856, upload-time = "2026-05-14T19:25:27.735Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a0/f4/c67b0b3f1b9245e8d266f0f112c500d50e5b4e83cb6f3b71b6528104182a/requests-2.34.2-py3-none-any.whl", hash = "sha256:2a0d60c172f83ac6ab31e4554906c0f3b3588d37b5cb939b1c061f4907e278e0", size = 73075, upload-time = "2026-05-14T19:25:26.443Z" }, +] + +[[package]] +name = "requests-toolbelt" +version = "1.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "requests", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/f3/61/d7545dafb7ac2230c70d38d31cbfe4cc64f7144dc41f6e4e4b78ecd9f5bb/requests-toolbelt-1.0.0.tar.gz", hash = "sha256:7681a0a3d047012b5bdc0ee37d7f8f07ebe76ab08caeccfc3921ce23c88d5bc6", size = 206888, upload-time = "2023-05-01T04:11:33.229Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3f/51/d4db610ef29373b879047326cbf6fa98b6c1969d6f6dc423279de2b1be2c/requests_toolbelt-1.0.0-py2.py3-none-any.whl", hash = "sha256:cccfdd665f0a24fcf4726e690f65639d272bb0637b9b92dfd91a5568ccf6bd06", size = 54481, upload-time = "2023-05-01T04:11:28.427Z" }, +] + +[[package]] +name = "rfc3986" +version = "2.0.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/85/40/1520d68bfa07ab5a6f065a186815fb6610c86fe957bc065754e47f7b0840/rfc3986-2.0.0.tar.gz", hash = "sha256:97aacf9dbd4bfd829baad6e6309fa6573aaf1be3f6fa735c8ab05e46cecb261c", size = 49026, upload-time = "2022-01-10T00:52:30.832Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/9a/9afaade874b2fa6c752c36f1548f718b5b83af81ed9b76628329dab81c1b/rfc3986-2.0.0-py2.py3-none-any.whl", hash = "sha256:50b1502b60e289cb37883f3dfd34532b8873c7de9f49bb546641ce9cbd256ebd", size = 31326, upload-time = "2022-01-10T00:52:29.594Z" }, +] + +[[package]] +name = "rich" +version = "15.0.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "markdown-it-py", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "pygments", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/c0/8f/0722ca900cc807c13a6a0c696dacf35430f72e0ec571c4275d2371fca3e9/rich-15.0.0.tar.gz", hash = "sha256:edd07a4824c6b40189fb7ac9bc4c52536e9780fbbfbddf6f1e2502c31b068c36", size = 230680, upload-time = "2026-04-12T08:24:00.75Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/82/3b/64d4899d73f91ba49a8c18a8ff3f0ea8f1c1d75481760df8c68ef5235bf5/rich-15.0.0-py3-none-any.whl", hash = "sha256:33bd4ef74232fb73fe9279a257718407f169c09b78a87ad3d296f548e27de0bb", size = 310654, upload-time = "2026-04-12T08:24:02.83Z" }, +] + +[[package]] +name = "ruff" +version = "0.16.3" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/61/b3/3213589383f8f1b3938781bd1278713f6d18621a14992b3e81fefb8a5ef9/ruff-0.16.3.tar.gz", hash = "sha256:e76d33a347661a84b5be6d043d0347fdc745dfdcf825a8f4fed64b5e26eebdf2", size = 4891904, upload-time = "2026-08-13T15:17:13.381Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/c7/30/4c468429ac195addc5ee1b717b6ab1b66632786737ca3b2ed3443fb0c26a/ruff-0.16.3-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:294b95c4ae0cda9388525c2047778aa758d6b8d4bb876fd4e9eaa3ebc92343eb", size = 11058736, upload-time = "2026-08-13T15:16:48.823Z" }, + { url = "https://files.pythonhosted.org/packages/35/85/7795f6e817af050e7517bf3e7aa9b061cce70ef33d280aad902c956c1ecf/ruff-0.16.3-py3-none-musllinux_1_2_x86_64.whl", hash = "sha256:a2d85c02f9b8e165d85e6779184d38c4132de12603dab59c51c28e22584f9e4d", size = 11626910, upload-time = "2026-08-13T15:17:03.299Z" }, +] + +[[package]] +name = "scikit-learn" +version = "1.9.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "joblib", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "narwhals", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "scipy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "threadpoolctl", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/fa/6f/37092bdb25f712817231799fc5674d8e704066a8a70c1d2d40517e18b4ab/scikit_learn-1.9.0.tar.gz", hash = "sha256:8833266989d3a5110178a9fae30783675460724d0e1efb13b14901d2c660c557", size = 7750767, upload-time = "2026-06-02T11:54:32.706Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/6c/c2/63fdda36c56437eeb44aaf9493c8bcd62ce230ab1598924fc626ffbfa943/scikit_learn-1.9.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:056c92bb67ad4c28463c2f2653d9701449201e7e7a9e94e321be0f71c4fef2b8", size = 9132097, upload-time = "2026-06-02T11:53:33.456Z" }, +] + +[[package]] +name = "scipy" +version = "1.17.1" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/7a/97/5a3609c4f8d58b039179648e62dd220f89864f56f7357f5d4f45c29eb2cc/scipy-1.17.1.tar.gz", hash = "sha256:95d8e012d8cb8816c226aef832200b1d45109ed4464303e997c5b13122b297c0", size = 30573822, upload-time = "2026-02-23T00:26:24.851Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/01/8e/1e35281b8ab6d5d72ebe9911edcdffa3f36b04ed9d51dec6dd140396e220/scipy-1.17.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:02ae3b274fde71c5e92ac4d54bc06c42d80e399fec704383dcd99b301df37458", size = 35235890, upload-time = "2026-02-23T00:18:49.188Z" }, + { url = "https://files.pythonhosted.org/packages/65/94/7698add8f276dbab7a9de9fb6b0e02fc13ee61d51c7c3f85ac28b65e1239/scipy-1.17.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f590cd684941912d10becc07325a3eeb77886fe981415660d9265c4c418d0bea", size = 37625856, upload-time = "2026-02-23T00:19:00.307Z" }, +] + +[[package]] +name = "secretstorage" +version = "3.5.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "cryptography", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "jeepney", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/1c/03/e834bcd866f2f8a49a85eaff47340affa3bfa391ee9912a952a1faa68c7b/secretstorage-3.5.0.tar.gz", hash = "sha256:f04b8e4689cbce351744d5537bf6b1329c6fc68f91fa666f60a380edddcd11be", size = 19884, upload-time = "2025-11-23T19:02:53.191Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/46/f5af3402b579fd5e11573ce652019a67074317e18c1935cc0b4ba9b35552/secretstorage-3.5.0-py3-none-any.whl", hash = "sha256:0ce65888c0725fcb2c5bc0fdb8e5438eece02c523557ea40ce0703c266248137", size = 15554, upload-time = "2025-11-23T19:02:51.545Z" }, +] + +[[package]] +name = "six" +version = "1.17.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031, upload-time = "2024-12-04T17:35:28.174Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050, upload-time = "2024-12-04T17:35:26.475Z" }, +] + +[[package]] +name = "sympy" +version = "1.14.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "mpmath", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/a2/09/77d55d46fd61b4a135c444fc97158ef34a095e5681d0a6c10b75bf356191/sympy-1.14.0-py3-none-any.whl", hash = "sha256:e091cc3e99d2141a0ba2847328f5479b05d94a6635cb96148ccb3f34671bd8f5", size = 6299353, upload-time = "2025-04-27T18:04:59.103Z" }, +] + +[[package]] +name = "threadpoolctl" +version = "3.6.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/4d/08c89e34946fce2aec4fbb45c9016efd5f4d7f24af8e5d93296e935631d8/threadpoolctl-3.6.0.tar.gz", hash = "sha256:8ab8b4aa3491d812b623328249fab5302a68d2d71745c8a4c719a2fcaba9f44e", size = 21274, upload-time = "2025-03-13T13:49:23.031Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/32/d5/f9a850d79b0851d1d4ef6456097579a9005b31fea68726a4ae5f2d82ddd9/threadpoolctl-3.6.0-py3-none-any.whl", hash = "sha256:43a0b8fd5a2928500110039e43a5eed8480b918967083ea48dc3ab9f13c4a7fb", size = 18638, upload-time = "2025-03-13T13:49:21.846Z" }, +] + +[[package]] +name = "torch" +version = "2.2.2+cu121" +source = { registry = "https://download.pytorch.org/whl/cu121" } +dependencies = [ + { name = "filelock", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fsspec", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "jinja2", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "networkx", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cublas-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cudnn-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cufft-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-curand-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusolver-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-cusparse-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nccl-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "nvidia-nvtx-cu12", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "sympy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +wheels = [ + { url = "https://download-r2.pytorch.org/whl/cu121/torch-2.2.2%2Bcu121-cp312-cp312-linux_x86_64.whl", hash = "sha256:badc14d413ff1847d15021a1ec0affa479d24dfc83e6d51b9b4b9fbfaad1b14c", upload-time = "2024-04-24T17:07:56Z" }, +] + +[[package]] +name = "tqdm" +version = "4.70.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/21/3b/6c24bec5be5e743ffd99576daa5cc077722fc7d5bbc00bd133fa0c698dc6/tqdm-4.70.0.tar.gz", hash = "sha256:55b0b0dbd97462d06ebee91e4dac24ed4d4702be82b24f07e6c1d27e08cea220", size = 795438, upload-time = "2026-07-27T11:33:15.271Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl", hash = "sha256:7f585706bfddbdebf89daac705b2dfcc16890130727d3197ca62c732b4310953", size = 80184, upload-time = "2026-07-27T11:33:13.167Z" }, +] + +[[package]] +name = "twine" +version = "6.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "id", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "keyring", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "readme-renderer", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "requests", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "requests-toolbelt", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "rfc3986", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "rich", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "urllib3", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/e0/a8/949edebe3a82774c1ec34f637f5dd82d1cf22c25e963b7d63771083bbee5/twine-6.2.0.tar.gz", hash = "sha256:e5ed0d2fd70c9959770dce51c8f39c8945c574e18173a7b81802dab51b4b75cf", size = 172262, upload-time = "2025-09-04T15:43:17.255Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/3a/7a/882d99539b19b1490cac5d77c67338d126e4122c8276bf640e411650c830/twine-6.2.0-py3-none-any.whl", hash = "sha256:418ebf08ccda9a8caaebe414433b0ba5e25eb5e4a927667122fbe8f829f985d8", size = 42727, upload-time = "2025-09-04T15:43:15.994Z" }, +] + +[[package]] +name = "typing-extensions" +version = "4.16.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/cc/6253133b5bb138fc3306cebfbda2c520f545d36b5be2c7255cc528bb45d6/typing_extensions-4.16.0.tar.gz", hash = "sha256:dc983d19a509c94dba722ee6abd33940f7c05a89e243c47e907eb4db6f1a43e5", size = 113555, upload-time = "2026-07-02T08:40:05.92Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/49/d3/b8441a820a491ddfc024b0b0cf0393375b75ea13866d9c66727e54c2fc80/typing_extensions-4.16.0-py3-none-any.whl", hash = "sha256:481caa481374e813c1b176ada14e97f1f67a4539ce9cfeb3f350d78d6370c2e8", size = 45571, upload-time = "2026-07-02T08:40:04.659Z" }, +] + +[[package]] +name = "typing-inspection" +version = "0.4.4" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/a3/26/b09b8010994eccc3c09092e6b34058f36a460eea2d4c3e8b910c695975a0/typing_inspection-0.4.4.tar.gz", hash = "sha256:547274fa6b0a561ccf549cc9524b999a578e737d015d8709d021f9d0d13bea47", size = 76928, upload-time = "2026-08-12T12:37:25.997Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/67/81/4add07e5172b7ac40d8ed5ff580409a7801a4fe26d529bdd915401dabfbe/typing_inspection-0.4.4-py3-none-any.whl", hash = "sha256:65b8397ba37ccbce054456aaccddfc91e6e3083c92824df348d96ca832f3f147", size = 14750, upload-time = "2026-08-12T12:37:24.648Z" }, +] + +[[package]] +name = "uproot" +version = "5.7.5" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "awkward", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "cramjam", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "fsspec", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "numpy", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "packaging", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, + { name = "xxhash", marker = "platform_machine == 'x86_64' and sys_platform == 'linux'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/d6/e6/9739727cc9c9ed95c40c23993eeecd9ea8367d5c36cbea20d97605a49de5/uproot-5.7.5.tar.gz", hash = "sha256:c9a30c8b39ffd30d9f87d7a8ec49a82c3d71fe163161b28e193c83d771119cd6", size = 1025717, upload-time = "2026-07-04T10:43:21.14Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7a/2c/2fb9b2d992bd9a9c4aed623c0ace167988027ccc7421b178a26a9a40ef50/uproot-5.7.5-py3-none-any.whl", hash = "sha256:25aee18f2d11dda6b26c6d76d84a9bc0f0aa0adeeba60548e47391964edae760", size = 401195, upload-time = "2026-07-04T10:43:19.434Z" }, +] + +[[package]] +name = "urllib3" +version = "2.7.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, +] + +[[package]] +name = "xxhash" +version = "4.0.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/f6/a5/1386f35da1475fcaeef42581deae73417c6d2a6a0b2d2e8914de18844dcd/xxhash-4.0.1.tar.gz", hash = "sha256:d55bf4ef10eb09b8b6866790e083d26d087d84caa3cc0946ba87c3ca7ecaf7b7", size = 101513, upload-time = "2026-08-17T08:24:08.557Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/83/97/31bd8b8279e6935a0719f6910ced15e9d5a2cd554b253f6027ce1b5a1c2c/xxhash-4.0.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:237b8f63a2a0fcfb1ffc06e21dad23add44e6d354b2b014364a1d41e419a4dee", size = 261812, upload-time = "2026-08-17T08:22:00.469Z" }, + { url = "https://files.pythonhosted.org/packages/6d/32/c6148d39a49efa95f39b4cf0d41ef35a487f3b30f6fb1fc8fe8d8eab577e/xxhash-4.0.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:96d8de55029d42251945531f6aa7590c32b48163c66a43bf29d8657d7446a377", size = 258174, upload-time = "2026-08-17T08:35:21.18Z" }, + { url = "https://files.pythonhosted.org/packages/8f/fb/0b04b68d6c5bc71c7a2c344f1287327b67e607f28fbcfd937697caca64b6/xxhash-4.0.1-cp312-cp312-pyemscripten_2024_0_wasm32.whl", hash = "sha256:0163b5d259de23ae9e07b7eabf435ce4704f6f205589a2b154e6af4be985ce1b", size = 20767, upload-time = "2026-08-17T08:21:00.806Z" }, + { url = "https://files.pythonhosted.org/packages/ad/23/2d549e5d5d7759eaf9ac2d2d2ab81ff60f1bb2b52cdaae8e5ec5c6524354/xxhash-4.0.1-graalpy312-graalpy250_312_native-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:deca2a30d983d240b8375ec2ee0a4288e72042827fc61df2f7671f8467e4cb2f", size = 38206, upload-time = "2026-08-17T08:36:32.193Z" }, +] diff --git a/validation/README.md b/validation/README.md new file mode 100644 index 0000000000000000000000000000000000000000..9b0016b2b7a00e6e42c18289b69fbbd60fde4138 --- /dev/null +++ b/validation/README.md @@ -0,0 +1,50 @@ +# End-to-end validation + +Task 21 uses a portable event-level artifact. Legacy and rewrite extraction +run independently and each writes `manifest.json` and `artifact.npz`; the +comparison environment joins events by `sample_id`. Variable-size arrays use +flattened storage plus offsets. Exact fields (IDs, labels, folds, offsets, +topology) are exact-compared; floating fields use `tolerances.yaml`. + +```bash +uv run python -m validation.run_full_validation --smoke +uv run python -m validation.run_full_validation \ + --legacy-artifact validation_output/legacy \ + --rewrite-artifact validation_output/rewrite \ + --output validation_output/reports +``` + +Smoke validates serialization and comparison plumbing only; it does not claim +legacy parity. Full scientific validation uses the same ROOT sample in both +runtimes, with `GNN4COLLIDERS_E2E_FIXTURE` selecting an external fixture. Keep +generated output out of Git. + +## Current HF fixture campaign + +The small fixture can be extracted in both environments: + +```bash +ROOT=/global/cfs/projectdirs/atlas/joshua/gnn_data/hf_validation/testing/ttH_NLO_64.root +HF_SHA=89d69eae9cd4d28a5d414d77bbc07185ea5b7de03481fdafe28899e2f3a09dec + +conda run -n dgl env PYTHONPATH=.:legacy/root_gnn_dgl:src \ + python validation/extract_legacy.py /tmp/legacy_ttree.root validation_output/legacy \ + --source-sha256 "$HF_SHA" + +PYTHONPATH=.:src uv run python validation/extract_rewrite.py "$ROOT" \ + validation_output/rewrite --source-sha256 "$HF_SHA" +``` + +The legacy environment is Python 3.8/Uproot 5.3 and cannot read the HF +RNTuple directly. Convert the fixture to a temporary TTree with a modern +Uproot process; this changes only the ROOT container representation, not event +values. The manifest retains the original HF SHA-256. Then run `forward.py` +with `model_epoch_71.pt`, `train_step.py`, `compare_step.py`, and +`compare_tasks.py` to generate fixed-forward, one-step, and task reports. + +For multiclass validation, `validation/manifests/multiclass_fixture.json` +defines the 96-event composite (eight entries from each of the 12 active +legacy classes). The source files are the HF-derived `stats_100K` samples; +the temporary composite adds only a `validation_label` branch containing the +legacy config class labels. This is the campaign fixture used for the final +12-class report. diff --git a/validation/__init__.py b/validation/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..4587c82da081a4317e110e5e5019454ba44e890a --- /dev/null +++ b/validation/__init__.py @@ -0,0 +1,17 @@ +"""Staged end-to-end parity validation utilities. + +The validation package is intentionally outside the production package. Its +interchange format contains NumPy arrays and JSON metadata only, so extraction +can happen in separate legacy and rewrite environments. +""" + +from .artifacts import ValidationArtifact, load_artifact, save_artifact +from .compare import compare_artifacts, write_reports + +__all__ = [ + "ValidationArtifact", + "compare_artifacts", + "load_artifact", + "save_artifact", + "write_reports", +] diff --git a/validation/artifacts.py b/validation/artifacts.py new file mode 100644 index 0000000000000000000000000000000000000000..8fec746317fc9cfde607ecd1bce19b9530c5aa9d --- /dev/null +++ b/validation/artifacts.py @@ -0,0 +1,140 @@ +"""Portable event-level artifacts used by the parity campaign.""" + +from __future__ import annotations + +import hashlib +import json +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import numpy as np + +SCHEMA_VERSION = 1 + + +@dataclass +class ValidationArtifact: + """Normalized representation of one extraction run. + + Variable-sized node and edge arrays are stored flattened with offsets. + ``sample_id`` is the join key and is never inferred from row position. + """ + + sample_id: np.ndarray + labels: np.ndarray + folds: np.ndarray + weights: np.ndarray + globals: np.ndarray + node_features_flat: np.ndarray + node_offsets: np.ndarray + edge_src_flat: np.ndarray + edge_dst_flat: np.ndarray + edge_features_flat: np.ndarray + edge_offsets: np.ndarray + logits: np.ndarray | None = None + scores: np.ndarray | None = None + predictions: np.ndarray | None = None + manifest: dict[str, Any] = field(default_factory=dict) + + @property + def event_count(self) -> int: + return int(self.sample_id.shape[0]) + + def validate(self) -> None: + n = self.event_count + for name in ("labels", "folds", "weights", "globals"): + if getattr(self, name).shape[0] != n: + raise ValueError(f"{name} does not contain one row per sample") + for name in ("node_offsets", "edge_offsets"): + offsets = getattr(self, name) + if offsets.shape != (n + 1,) or offsets[0] != 0: + raise ValueError(f"{name} must have shape ({n + 1},) and start at zero") + if np.any(offsets[1:] < offsets[:-1]): + raise ValueError(f"{name} must be monotonic") + if self.node_offsets[-1] != len(self.node_features_flat): + raise ValueError("node offsets do not describe node_features_flat") + if self.edge_offsets[-1] != len(self.edge_src_flat): + raise ValueError("edge offsets do not describe edge_src_flat") + if len(self.edge_src_flat) != len(self.edge_dst_flat): + raise ValueError("edge source and destination arrays differ in length") + if self.edge_offsets[-1] != len(self.edge_features_flat): + raise ValueError("edge offsets do not describe edge_features_flat") + if len(np.unique(self.sample_id)) != n: + raise ValueError("sample_id contains duplicates") + + def event_nodes(self, index: int) -> np.ndarray: + start, stop = self.node_offsets[index : index + 2] + return self.node_features_flat[start:stop] + + def event_edges(self, index: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + start, stop = self.edge_offsets[index : index + 2] + return ( + self.edge_src_flat[start:stop], + self.edge_dst_flat[start:stop], + self.edge_features_flat[start:stop], + ) + + +def _optional_arrays(artifact: ValidationArtifact) -> dict[str, np.ndarray]: + return { + name: value + for name in ("logits", "scores", "predictions") + if (value := getattr(artifact, name)) is not None + } + + +def save_artifact(artifact: ValidationArtifact, directory: str | Path) -> Path: + artifact.validate() + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + arrays = { + "sample_id": artifact.sample_id, + "labels": artifact.labels, + "folds": artifact.folds, + "weights": artifact.weights, + "globals": artifact.globals, + "node_features_flat": artifact.node_features_flat, + "node_offsets": artifact.node_offsets, + "edge_src_flat": artifact.edge_src_flat, + "edge_dst_flat": artifact.edge_dst_flat, + "edge_features_flat": artifact.edge_features_flat, + "edge_offsets": artifact.edge_offsets, + **_optional_arrays(artifact), + } + np.savez_compressed(directory / "artifact.npz", **arrays) + manifest = { + "schema_version": SCHEMA_VERSION, + "event_count": artifact.event_count, + "arrays": sorted(arrays), + **artifact.manifest, + } + (directory / "manifest.json").write_text( + json.dumps(manifest, indent=2, sort_keys=True) + "\n" + ) + return directory + + +def load_artifact(directory: str | Path) -> ValidationArtifact: + directory = Path(directory) + manifest = json.loads((directory / "manifest.json").read_text()) + if manifest.get("schema_version") != SCHEMA_VERSION: + raise ValueError( + f"unsupported validation artifact schema: {manifest.get('schema_version')}" + ) + with np.load(directory / "artifact.npz", allow_pickle=False) as data: + values = {key: data[key] for key in data.files} + optional = { + name: values.pop(name, None) for name in ("logits", "scores", "predictions") + } + artifact = ValidationArtifact(**values, **optional, manifest=manifest) + artifact.validate() + return artifact + + +def file_sha256(path: str | Path) -> str: + digest = hashlib.sha256() + with Path(path).open("rb") as stream: + for chunk in iter(lambda: stream.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() diff --git a/validation/batching.py b/validation/batching.py new file mode 100644 index 0000000000000000000000000000000000000000..e8ef06fe12c4fefb1907f0701da675d13ea35f06 --- /dev/null +++ b/validation/batching.py @@ -0,0 +1,115 @@ +#!/usr/bin/env python3 +"""Dump deterministic batch membership and aggregate sizes for one implementation.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch + +from validation.artifacts import load_artifact +from validation.forward import _graph + + +class _LegacyDataset(torch.utils.data.Dataset): + def __init__(self, artifact): + self.artifact = artifact + + def __len__(self): + return self.artifact.event_count + + def __getitem__(self, index): + graph = _graph(self.artifact, index) + tracking = torch.tensor( + [ + self.artifact.folds[index], + self.artifact.weights[index], + index, + ], + dtype=torch.float32, + ) + return ( + graph, + torch.tensor(self.artifact.labels[index]), + tracking, + torch.empty(0), + ) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", choices=("legacy", "rewrite"), required=True + ) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--batch-size", type=int, default=8) + args = parser.parse_args() + artifact = load_artifact(args.artifact) + if args.implementation == "legacy": + from dgl.dataloading import GraphDataLoader + + loader = GraphDataLoader( + _LegacyDataset(artifact), + batch_size=args.batch_size, + shuffle=False, + drop_last=False, + num_workers=0, + ) + else: + from gnn4colliders.data import ( + EventMetadata, + GraphDataLoader, + GraphDataset, + GraphSample, + ) + + samples = [] + for index in range(artifact.event_count): + samples.append( + GraphSample( + _graph(artifact, index), + torch.tensor(artifact.labels[index]), + None, + EventMetadata( + int(artifact.folds[index]), + float(artifact.weights[index]), + str(artifact.sample_id[index]), + {"index": index}, + ), + ) + ) + loader = GraphDataLoader(GraphDataset(samples), args.batch_size, shuffle=False) + batches = [] + for batch in loader: + if args.implementation == "legacy": + graph, labels, tracking, _ = batch + indices = tracking[:, 2].to(torch.long).tolist() + ids = [str(artifact.sample_id[index]) for index in indices] + weights = tracking[:, 1].tolist() + else: + graph = batch.graph + labels = batch.labels + ids = list(batch.metadata.sample_id) + weights = batch.metadata.weight.tolist() + batches.append( + { + "sample_id": ids, + "num_graphs": len(ids), + "nodes": int(graph.num_nodes()), + "edges": int(graph.num_edges()), + "labels": labels.reshape(-1).tolist(), + "weights": weights, + } + ) + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text( + json.dumps({"batch_size": args.batch_size, "batches": batches}, indent=2) + "\n" + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/compare.py b/validation/compare.py new file mode 100644 index 0000000000000000000000000000000000000000..40a4519761e037db2d7b54c52334cf612873058a --- /dev/null +++ b/validation/compare.py @@ -0,0 +1,181 @@ +"""Stage-aware comparison of normalized legacy and rewrite artifacts.""" + +from __future__ import annotations + +import json +from pathlib import Path +from typing import Any + +import numpy as np + +from .artifacts import ValidationArtifact, load_artifact + +STATUS = ("PASS", "PASS_WITH_EXPECTED_DIFFERENCE", "FAIL", "NOT_APPLICABLE", "NOT_RUN") + + +def _numeric( + left: np.ndarray, right: np.ndarray, atol: float, rtol: float +) -> dict[str, Any]: + if left.shape != right.shape: + return { + "status": "FAIL", + "reason": "shape mismatch", + "left_shape": left.shape, + "right_shape": right.shape, + } + delta = np.abs(left.astype(np.float64) - right.astype(np.float64)) + scale = np.maximum( + np.abs(left.astype(np.float64)), np.abs(right.astype(np.float64)) + ) + relative = np.divide(delta, scale, out=np.zeros_like(delta), where=scale != 0) + passed = np.isclose(left, right, atol=atol, rtol=rtol) + location = ( + np.unravel_index(int(delta.argmax()), delta.shape) if delta.size else None + ) + return { + "status": "PASS" if bool(np.all(passed)) else "FAIL", + "max_abs": float(delta.max()) if delta.size else 0.0, + "max_rel": float(relative.max()) if relative.size else 0.0, + "outside_tolerance": int((~passed).sum()), + "elements": int(delta.size), + "location": list(location) if location is not None else None, + "atol": atol, + "rtol": rtol, + } + + +def _exact(left: np.ndarray, right: np.ndarray) -> dict[str, Any]: + return { + "status": "PASS" if np.array_equal(left, right) else "FAIL", + "elements": int(left.size), + "shape": left.shape, + "reason": "exact mismatch" if not np.array_equal(left, right) else None, + } + + +def _ids(left: ValidationArtifact, right: ValidationArtifact) -> dict[str, Any]: + left_ids, right_ids = left.sample_id.tolist(), right.sample_id.tolist() + left_set, right_set = set(left_ids), set(right_ids) + duplicates = sorted( + {item for item in left_ids if left_ids.count(item) > 1} + | {item for item in right_ids if right_ids.count(item) > 1} + ) + return { + "status": "PASS" if left_set == right_set and not duplicates else "FAIL", + "left_count": len(left_ids), + "right_count": len(right_ids), + "missing_in_rewrite": sorted(left_set - right_set), + "missing_in_legacy": sorted(right_set - left_set), + "duplicates": duplicates, + "order_equal": left_ids == right_ids, + } + + +def compare_artifacts( + left: ValidationArtifact | str | Path, + right: ValidationArtifact | str | Path, + *, + tolerances: dict[str, float] | None = None, +) -> dict[str, Any]: + left = load_artifact(left) if not isinstance(left, ValidationArtifact) else left + right = load_artifact(right) if not isinstance(right, ValidationArtifact) else right + left.validate() + right.validate() + tol = { + "feature_atol": 1e-6, + "feature_rtol": 1e-6, + "logit_atol": 1e-5, + "logit_rtol": 1e-5, + **(tolerances or {}), + } + report: dict[str, Any] = { + "schema_version": 1, + "dataset": _ids(left, right), + "stages": {}, + } + stages = report["stages"] + stages["labels"] = _exact(left.labels, right.labels) + stages["folds"] = _exact(left.folds, right.folds) + stages["weights"] = _numeric( + left.weights, right.weights, tol["feature_atol"], tol["feature_rtol"] + ) + stages["globals"] = _numeric( + left.globals, right.globals, tol["feature_atol"], tol["feature_rtol"] + ) + stages["node_offsets"] = _exact(left.node_offsets, right.node_offsets) + stages["node_features"] = _numeric( + left.node_features_flat, + right.node_features_flat, + tol["feature_atol"], + tol["feature_rtol"], + ) + stages["edge_offsets"] = _exact(left.edge_offsets, right.edge_offsets) + stages["edge_src"] = _exact(left.edge_src_flat, right.edge_src_flat) + stages["edge_dst"] = _exact(left.edge_dst_flat, right.edge_dst_flat) + stages["edge_features"] = _numeric( + left.edge_features_flat, + right.edge_features_flat, + tol["feature_atol"], + tol["feature_rtol"], + ) + if left.logits is None or right.logits is None: + stages["logits"] = {"status": "NOT_APPLICABLE"} + else: + stages["logits"] = _numeric( + left.logits, right.logits, tol["logit_atol"], tol["logit_rtol"] + ) + if left.scores is None or right.scores is None: + stages["scores"] = {"status": "NOT_APPLICABLE"} + else: + stages["scores"] = _numeric( + left.scores, right.scores, tol["logit_atol"], tol["logit_rtol"] + ) + if left.predictions is None or right.predictions is None: + stages["predictions"] = {"status": "NOT_APPLICABLE"} + else: + stages["predictions"] = _exact(left.predictions, right.predictions) + report["overall_status"] = ( + "PASS" + if all( + item["status"] + in ("PASS", "PASS_WITH_EXPECTED_DIFFERENCE", "NOT_APPLICABLE") + for item in [report["dataset"], *stages.values()] + ) + else "FAIL" + ) + return report + + +def write_reports(report: dict[str, Any], directory: str | Path) -> tuple[Path, Path]: + directory = Path(directory) + directory.mkdir(parents=True, exist_ok=True) + json_path = directory / "comparison.json" + md_path = directory / "comparison.md" + json_path.write_text( + json.dumps(report, indent=2, sort_keys=True, default=_json_default) + "\n" + ) + lines = [ + "# End-to-end parity report", + "", + f"Overall status: **{report['overall_status']}**", + "", + "| Stage | Status | Max abs | Outside tolerance |", + "|---|---|---:|---:|", + ] + for name, result in report.get("stages", {}).items(): + lines.append( + f"| {name} | {result['status']} | " + f"{result.get('max_abs', '—')} | " + f"{result.get('outside_tolerance', '—')} |" + ) + md_path.write_text("\n".join(lines) + "\n") + return json_path, md_path + + +def _json_default(value: Any) -> Any: + """Serialize NumPy scalar/shape values present in diagnostic reports.""" + if isinstance(value, np.generic): + return value.item() + if isinstance(value, tuple): + return list(value) + raise TypeError(f"cannot serialize {type(value).__name__}") diff --git a/validation/compare_batches.py b/validation/compare_batches.py new file mode 100644 index 0000000000000000000000000000000000000000..37decdfb76afd204137e2b43498f43e27b3c9fd1 --- /dev/null +++ b/validation/compare_batches.py @@ -0,0 +1,32 @@ +#!/usr/bin/env python3 +"""Compare deterministic batch manifests.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("legacy", type=Path) + parser.add_argument("rewrite", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + left = json.loads(args.legacy.read_text()) + right = json.loads(args.rewrite.read_text()) + report = { + "schema_version": 1, + "status": "PASS" if left == right else "FAIL", + "legacy": left, + "rewrite": right, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(report, indent=2) + "\n") + print(json.dumps({"overall_status": report["status"]}, indent=2)) + return 0 if report["status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/compare_step.py b/validation/compare_step.py new file mode 100644 index 0000000000000000000000000000000000000000..b0f877d2b36b1bda163c88bb5907da9f1edbc03d --- /dev/null +++ b/validation/compare_step.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Compare serialized one-step legacy/rewrite diagnostics.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + + +def _compare(left: np.ndarray, right: np.ndarray, atol: float, rtol: float): + delta = np.abs(left.astype(np.float64) - right.astype(np.float64)) + passed = np.isclose(left, right, atol=atol, rtol=rtol) + return { + "status": "PASS" if bool(np.all(passed)) else "FAIL", + "shape": list(left.shape), + "max_abs": float(delta.max()) if delta.size else 0.0, + "mean_abs": float(delta.mean()) if delta.size else 0.0, + "outside_tolerance": int((~passed).sum()), + "elements": int(delta.size), + "atol": atol, + "rtol": rtol, + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("legacy", type=Path) + parser.add_argument("rewrite", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + with np.load(args.legacy) as left, np.load(args.rewrite) as right: + stages = { + "logits": _compare(left["logits"], right["logits"], 1e-5, 1e-5), + "loss": _compare(left["loss"], right["loss"], 1e-5, 1e-5), + "gradients": _compare(left["gradients"], right["gradients"], 1e-4, 1e-5), + "parameters_after_step": _compare( + left["parameters"], right["parameters"], 1e-5, 1e-5 + ), + } + if "history" in left.files and "history" in right.files: + stages["history"] = _compare(left["history"], right["history"], 1e-5, 1e-5) + if "exp_avg" in left.files and "exp_avg" in right.files: + stages["optimizer_exp_avg"] = _compare( + left["exp_avg"], right["exp_avg"], 1e-5, 1e-5 + ) + stages["optimizer_exp_avg_sq"] = _compare( + left["exp_avg_sq"], right["exp_avg_sq"], 1e-5, 1e-5 + ) + stages["optimizer_steps"] = _compare( + left["optimizer_steps"], right["optimizer_steps"], 0.0, 0.0 + ) + report = {"schema_version": 1, "stages": stages} + report["overall_status"] = ( + "PASS" + if all(stage["status"] == "PASS" for stage in stages.values()) + else "FAIL" + ) + args.output.mkdir(parents=True, exist_ok=True) + (args.output / "training_step.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n" + ) + lines = [ + "# One-step training parity", + "", + f"Overall status: **{report['overall_status']}**", + "", + "| Stage | Status | Max abs | Outside tolerance |", + "|---|---|---:|---:|", + ] + for name, stage in stages.items(): + lines.append( + f"| {name} | {stage['status']} | {stage['max_abs']} | " + f"{stage['outside_tolerance']} |" + ) + (args.output / "training_step.md").write_text("\n".join(lines) + "\n") + print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) + return 0 if report["overall_status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/compare_tasks.py b/validation/compare_tasks.py new file mode 100644 index 0000000000000000000000000000000000000000..ca1e60823796e3b50a86d9d778b79e328b47b254 --- /dev/null +++ b/validation/compare_tasks.py @@ -0,0 +1,93 @@ +#!/usr/bin/env python3 +"""Compare task loss/score/metric semantics for two forward artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from types import SimpleNamespace + +import torch +from sklearn.metrics import roc_auc_score + +from validation.artifacts import load_artifact + + +def _evaluate(artifact, *, rewrite: bool): + logits = torch.from_numpy(artifact.logits).to(torch.float32) + labels = torch.from_numpy(artifact.labels).to(torch.long) + weights = torch.from_numpy(artifact.weights).to(torch.float32) + elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none") + loss = torch.zeros(()) + for label in torch.unique(labels): + mask = labels == label + loss = loss + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum() + loss = loss / len(torch.unique(labels)) + scores = torch.softmax(logits, dim=1) + predictions = scores.argmax(dim=1) + positive = weights > 0 + one_hot = torch.nn.functional.one_hot(labels, num_classes=scores.shape[1]).numpy() + positive_numpy = positive.numpy() + legacy_auc = float( + roc_auc_score( + one_hot[positive_numpy], + scores.detach().numpy()[positive_numpy], + multi_class="ovr", + sample_weight=weights.numpy()[positive_numpy], + ) + ) + if rewrite: + from gnn4colliders.tasks import MulticlassClassificationTask + + batch = SimpleNamespace( + labels=labels, + metadata=SimpleNamespace(weight=weights), + ) + auc = MulticlassClassificationTask().metrics(logits, batch)["roc_auc"] + else: + auc = legacy_auc + return { + "loss": float(loss), + "accuracy": float((predictions == labels).float().mean()), + "roc_auc": auc, + "positive_weight_events": int(positive.sum()), + "negative_weight_events": int((weights < 0).sum()), + "zero_weight_events": int((weights == 0).sum()), + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("legacy", type=Path) + parser.add_argument("rewrite", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + left = _evaluate(load_artifact(args.legacy), rewrite=False) + right = _evaluate(load_artifact(args.rewrite), rewrite=True) + differences = { + name: abs(left[name] - right[name]) for name in ("loss", "accuracy", "roc_auc") + } + report = { + "schema_version": 1, + "legacy": left, + "rewrite": right, + "absolute_differences": differences, + "overall_status": "PASS" + if all(value <= 1e-5 for value in differences.values()) + else "FAIL", + "note": ( + "Legacy uses sklearn weighted OVR AUC; rewrite uses its task-owned " + "weighted OVR implementation." + ), + } + args.output.mkdir(parents=True, exist_ok=True) + (args.output / "task_metrics.json").write_text( + json.dumps(report, indent=2, sort_keys=True, allow_nan=True) + "\n" + ) + print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) + return 0 if report["overall_status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/compare_transfer.py b/validation/compare_transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..62aca26da432b128987cedee7055b1f6a18bd3c0 --- /dev/null +++ b/validation/compare_transfer.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +"""Compare frozen/trainable transfer-forward artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + + +def compare(left: np.ndarray, right: np.ndarray) -> dict[str, object]: + delta = np.abs(left.astype(np.float64) - right.astype(np.float64)) + passed = np.isclose(left, right, atol=1e-5, rtol=1e-5) + return { + "status": "PASS" if bool(np.all(passed)) else "FAIL", + "shape": list(left.shape), + "max_abs": float(delta.max()), + "outside_tolerance": int((~passed).sum()), + } + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("legacy_frozen", type=Path) + parser.add_argument("rewrite_frozen", type=Path) + parser.add_argument("legacy_trainable", type=Path) + parser.add_argument("rewrite_trainable", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + reports = {} + for name, left_path, right_path in ( + ("frozen_backbone", args.legacy_frozen, args.rewrite_frozen), + ("trainable_backbone", args.legacy_trainable, args.rewrite_trainable), + ): + with np.load(left_path) as left, np.load(right_path) as right: + reports[name] = compare(left["logits"], right["logits"]) + report = { + "schema_version": 1, + "stages": reports, + "overall_status": "PASS" + if all(item["status"] == "PASS" for item in reports.values()) + else "FAIL", + } + args.output.mkdir(parents=True, exist_ok=True) + (args.output / "transfer.json").write_text( + json.dumps(report, indent=2, sort_keys=True) + "\n" + ) + print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) + return 0 if report["overall_status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/consolidate.py b/validation/consolidate.py new file mode 100644 index 0000000000000000000000000000000000000000..ae4a9c98bb226a4d877d9d832a071577bfa6725b --- /dev/null +++ b/validation/consolidate.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python3 +"""Combine staged reports into one machine- and human-readable campaign report.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + + +def _read(path: Path, name: str): + return json.loads((path / name).read_text()) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--preprocessing", type=Path, required=True) + parser.add_argument("--forward", type=Path, required=True) + parser.add_argument("--step", type=Path, required=True) + parser.add_argument("--short", type=Path, required=True) + parser.add_argument("--tasks", type=Path, required=True) + parser.add_argument("--batch", type=Path, required=True) + parser.add_argument("--transfer", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + args = parser.parse_args() + report = { + "schema_version": 1, + "preprocessing": _read(args.preprocessing, "comparison.json"), + "forward": _read(args.forward, "comparison.json"), + "training_step": _read(args.step, "training_step.json"), + "short_training": _read(args.short, "training_step.json"), + "tasks": _read(args.tasks, "task_metrics.json"), + "batching": json.loads(args.batch.read_text()), + "transfer": _read(args.transfer, "transfer.json"), + } + statuses = [ + report["preprocessing"]["overall_status"], + report["forward"]["overall_status"], + report["training_step"]["overall_status"], + report["short_training"]["overall_status"], + report["tasks"]["overall_status"], + report["batching"]["status"], + report["transfer"]["overall_status"], + ] + report["overall_status"] = ( + "PASS" if all(status == "PASS" for status in statuses) else "FAIL" + ) + args.output.mkdir(parents=True, exist_ok=True) + (args.output / "comparison.json").write_text( + json.dumps(report, indent=2, sort_keys=True, allow_nan=True) + "\n" + ) + rows = [ + ("Preprocessing", report["preprocessing"]["overall_status"]), + ("Fixed forward", report["forward"]["overall_status"]), + ("One optimizer step", report["training_step"]["overall_status"]), + ("Three-step training trajectory", report["short_training"]["overall_status"]), + ("Loss and metrics", report["tasks"]["overall_status"]), + ("Batch membership", report["batching"]["status"]), + ("Binary transfer forward", report["transfer"]["overall_status"]), + ] + lines = [ + "# GNN4Colliders end-to-end parity report", + "", + f"Overall status: **{report['overall_status']}**", + "", + "| Stage | Status |", + "|---|---|", + *(f"| {name} | {status} |" for name, status in rows), + "", + "The campaign uses the HF-derived 96-event, 12-class composite fixture", + "(eight entries per active legacy class).", + ] + (args.output / "comparison.md").write_text("\n".join(lines) + "\n") + print(json.dumps({"overall_status": report["overall_status"]}, indent=2)) + return 0 if report["overall_status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/convert_rntuple.py b/validation/convert_rntuple.py new file mode 100644 index 0000000000000000000000000000000000000000..f463a0847f98942218f80c38ca369d8020dc41e2 --- /dev/null +++ b/validation/convert_rntuple.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python3 +"""Make a temporary TTree copy for legacy Uproot versions.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import awkward as ak +import uproot + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("source", type=Path) + parser.add_argument("output", type=Path) + args = parser.parse_args() + with uproot.open(args.source) as root_file: + tree = root_file["output"] + arrays = tree.arrays(list(tree.keys()), library="ak") + branch_types = { + name: str(ak.type(arrays[name])).split("* ", 1)[1] for name in arrays.fields + } + with uproot.recreate(args.output) as root_file: + root_file.mktree("output", branch_types) + root_file["output"].extend(arrays) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/extract_common.py b/validation/extract_common.py new file mode 100644 index 0000000000000000000000000000000000000000..d0ecf50868dfd248161395ac5899d117db68dab4 --- /dev/null +++ b/validation/extract_common.py @@ -0,0 +1,151 @@ +"""Shared portable extraction helpers; no legacy or DGL classes cross the boundary.""" + +from __future__ import annotations + +from pathlib import Path +from typing import Any, Callable + +import awkward as ak +import numpy as np +import torch +import uproot + +from .artifacts import ValidationArtifact, file_sha256, save_artifact + +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 = [1e-1, 1, 1, 1e-1, 1, 1, 1] +BRANCHES = ( + "jet_pt", + "ele_pt", + "mu_pt", + "ph_pt", + "MET_met", + "jet_eta", + "ele_eta", + "mu_eta", + "ph_eta", + "jet_phi", + "ele_phi", + "mu_phi", + "ph_phi", + "MET_phi", + "jet_btag", + "ele_charge", + "mu_charge", + "Number", + "weight", +) + + +def _row(value: Any) -> Any: + """Convert an Awkward scalar while retaining a jagged vector.""" + converted = ak.to_numpy(value) if isinstance(value, ak.Array) else value + array = np.asarray(converted) + return array.item() if array.ndim == 0 else converted + + +def extract_root( + path: str | Path, + output: str | Path, + *, + build_features: Callable[[dict[str, Any]], torch.Tensor], + sample_name: str = "testing/ttH_NLO_64.root", + label: int = 0, + filter_zero_pt: bool = False, + source_sha256: str | None = None, + label_branch: str | None = None, +) -> Path: + """Extract one ROOT tree into the implementation-neutral artifact format.""" + path = Path(path) + requested_branches = list(BRANCHES) + if label_branch and label_branch not in requested_branches: + requested_branches.append(label_branch) + with uproot.open(path) as root_file: + tree = root_file["output"] + arrays = tree.arrays(requested_branches, library="ak") + event_count = int(tree.num_entries) + + sample_ids: list[str] = [] + labels: list[int] = [] + folds: list[int] = [] + weights: list[float] = [] + node_parts: list[np.ndarray] = [] + node_offsets = [0] + edge_src_parts: list[np.ndarray] = [] + edge_dst_parts: list[np.ndarray] = [] + edge_feature_parts: list[np.ndarray] = [] + edge_offsets = [0] + + for index in range(event_count): + event = {name: _row(arrays[name][index]) for name in requested_branches} + nodes = build_features(event) + if filter_zero_pt: + nodes = nodes[nodes[:, 0] != 0] + node_count = int(nodes.shape[0]) + source = torch.arange(node_count).repeat_interleave(node_count) + destination = torch.arange(node_count).repeat(node_count) + if node_count > 1: + keep = source != destination + source, destination = source[keep], destination[keep] + eta = nodes[:, 1] + phi = nodes[:, 2] + deta = eta[source] - eta[destination] + dphi = phi[source] - phi[destination] + dphi = torch.where(dphi > np.pi, dphi - 2 * np.pi, dphi) + dphi = torch.where(dphi < -np.pi, dphi + 2 * np.pi, dphi) + edge_features = torch.stack( + (deta, dphi, torch.sqrt(deta.square() + dphi.square())), dim=1 + ) + + sample_ids.append(f"{sample_name}:output:{index}") + labels.append(int(event[label_branch]) if label_branch else label) + folds.append(int(event["Number"])) + weights.append(float(event["weight"])) + node_parts.append(nodes.detach().cpu().numpy()) + edge_src_parts.append(source.numpy()) + edge_dst_parts.append(destination.numpy()) + edge_feature_parts.append(edge_features.detach().cpu().numpy()) + node_offsets.append(node_offsets[-1] + node_count) + edge_offsets.append(edge_offsets[-1] + int(source.numel())) + + artifact = ValidationArtifact( + sample_id=np.asarray(sample_ids), + labels=np.asarray(labels), + folds=np.asarray(folds), + weights=np.asarray(weights, dtype=np.float32), + globals=np.empty((event_count, 0), dtype=np.float32), + node_features_flat=np.concatenate(node_parts, axis=0), + node_offsets=np.asarray(node_offsets, dtype=np.int64), + edge_src_flat=np.concatenate(edge_src_parts), + edge_dst_flat=np.concatenate(edge_dst_parts), + edge_features_flat=np.concatenate(edge_feature_parts, axis=0), + edge_offsets=np.asarray(edge_offsets, dtype=np.int64), + manifest={ + "source": "huggingface://datasets/HWresearch/Delphes", + "repository_path": sample_name, + "tree": "output", + "events": event_count, + "sha256": file_sha256(path), + "source_sha256": source_sha256 or file_sha256(path), + "input_path": str(path), + "feature_schema": [ + "pt", + "eta", + "phi", + "energy", + "btag", + "charge", + "node_type", + ], + }, + ) + return save_artifact(artifact, output) diff --git a/validation/extract_legacy.py b/validation/extract_legacy.py new file mode 100644 index 0000000000000000000000000000000000000000..edc0e6bcd843557228cb857ea2b34b5286015042 --- /dev/null +++ b/validation/extract_legacy.py @@ -0,0 +1,49 @@ +#!/usr/bin/env python3 +"""Extract normalized artifacts with the frozen legacy feature builder.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import torch + +from validation.extract_common import ( + FEATURE_BRANCHES, + OBJECT_TYPES, + SCALES, + extract_root, +) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("root_file", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--source-sha256") + parser.add_argument("--label-branch") + args = parser.parse_args() + from root_gnn_base.dataset import node_features_from_tree + + def build(event): + features, _ = node_features_from_tree( + event, + FEATURE_BRANCHES, + OBJECT_TYPES, + torch.tensor(SCALES, dtype=torch.float32), + ) + return features + + extract_root( + args.root_file, + args.output, + build_features=build, + filter_zero_pt=True, + source_sha256=args.source_sha256, + label_branch=args.label_branch, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/extract_rewrite.py b/validation/extract_rewrite.py new file mode 100644 index 0000000000000000000000000000000000000000..7d8aaec1aa5091949c3c85d4002a92b336ea2655 --- /dev/null +++ b/validation/extract_rewrite.py @@ -0,0 +1,41 @@ +#!/usr/bin/env python3 +"""Extract normalized artifacts with the rewritten shared feature builder.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from validation.extract_common import ( + FEATURE_BRANCHES, + OBJECT_TYPES, + SCALES, + extract_root, +) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("root_file", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--source-sha256") + parser.add_argument("--label-branch") + args = parser.parse_args() + from gnn4colliders.features import build_node_features + + def build(event): + features, _ = build_node_features(event, FEATURE_BRANCHES, OBJECT_TYPES, SCALES) + return features + + extract_root( + args.root_file, + args.output, + build_features=build, + source_sha256=args.source_sha256, + label_branch=args.label_branch, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/forward.py b/validation/forward.py new file mode 100644 index 0000000000000000000000000000000000000000..234d331d4a4cf73d38c1bdb25a2698864e0be3f8 --- /dev/null +++ b/validation/forward.py @@ -0,0 +1,96 @@ +#!/usr/bin/env python3 +"""Run fixed-weight ROOT-GNN inference on a normalized validation artifact.""" + +from __future__ import annotations + +import argparse +from dataclasses import replace +from pathlib import Path + +import dgl +import numpy as np +import torch + +from validation.artifacts import load_artifact, save_artifact + + +def _graph(artifact, index: int): + nodes = torch.from_numpy(artifact.event_nodes(index)).to(torch.float32) + src, dst, edges = artifact.event_edges(index) + graph = dgl.graph( + (torch.from_numpy(src), torch.from_numpy(dst)), num_nodes=nodes.shape[0] + ) + graph.ndata["features"] = nodes + graph.edata["features"] = torch.from_numpy(edges).to(torch.float32) + return graph + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", choices=("legacy", "rewrite"), required=True + ) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--events", type=int) + args = parser.parse_args() + artifact = load_artifact(args.artifact) + event_count = ( + artifact.event_count + if args.events is None + else min(args.events, artifact.event_count) + ) + graphs = [_graph(artifact, index) for index in range(event_count)] + batch_graph = dgl.batch(graphs) + first = graphs[0] + empty_globals = torch.empty((1, 0), dtype=torch.float32) + if args.implementation == "legacy": + from models.GCN import Edge_Network + + model = Edge_Network(first, empty_globals, 64, 12, 4, 4, dropout=0.0) + payload = torch.load(args.checkpoint, map_location="cpu") + model.load_state_dict(payload["model_state_dict"]) + else: + from gnn4colliders.models.root_gnn import ( + EdgeNetwork, + load_legacy_edge_network_state_dict, + ) + + model = EdgeNetwork(first, empty_globals, 64, 12, 4, 4, dropout=0.0) + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + load_legacy_edge_network_state_dict(model, payload) + model.eval() + + with torch.inference_mode(): + logits = model(batch_graph, None).cpu().numpy() + args.output.mkdir(parents=True, exist_ok=True) + reload_path = args.output / "checkpoint_reload.pt" + torch.save({"model_state_dict": model.state_dict()}, reload_path) + fresh = type(model)(first, empty_globals, 64, 12, 4, 4, dropout=0.0) + fresh.load_state_dict( + torch.load(reload_path, map_location="cpu")["model_state_dict"] + ) + fresh.eval() + with torch.inference_mode(): + reloaded_logits = fresh(batch_graph, None).cpu().numpy() + scores = torch.softmax(torch.from_numpy(logits), dim=1).numpy() + predictions = scores.argmax(axis=1) + updated = replace( + artifact, + logits=logits, + scores=scores, + predictions=predictions, + manifest={ + **artifact.manifest, + "implementation": args.implementation, + "checkpoint": str(args.checkpoint), + "reload_max_abs": float(np.abs(logits - reloaded_logits).max()), + }, + ) + save_artifact(updated, args.output) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/manifests/config_mapping.md b/validation/manifests/config_mapping.md new file mode 100644 index 0000000000000000000000000000000000000000..59d94d4e0c8e19d5e1bdab9d925a9f7358c9c2ce --- /dev/null +++ b/validation/manifests/config_mapping.md @@ -0,0 +1,11 @@ +# Canonical configuration mapping + +| Legacy setting | Rewrite setting | Classification | +|---|---|---| +| `models.GCN.Edge_Network` | `model.type: root_gnn` | implementation-only | +| `args.hid_size`, `out_size`, `n_layers`, `n_proc_steps`, `dropout` | corresponding `model.*` fields | scientific | +| `tracking[:,0]` | `EventMetadata.fold` | representation-only | +| `tracking[:,1]` | `EventMetadata.weight` | representation-only | +| legacy batch size | `data.batch_size` | scientific | +| legacy fold selection | `data.split.*` and named metadata | scientific | +| input/output paths, workers, device | invocation/environment | environmental | diff --git a/validation/manifests/dataset.json b/validation/manifests/dataset.json new file mode 100644 index 0000000000000000000000000000000000000000..b046738a67439e710f996f4a04635b706b2e2f1f --- /dev/null +++ b/validation/manifests/dataset.json @@ -0,0 +1,21 @@ +{ + "source": "https://huggingface.co/datasets/HWresearch/Delphes", + "revision": "main", + "files": [ + { + "repository_path": "testing/ttH_NLO_64.root", + "sha256": "89d69eae9cd4d28a5d414d77bbc07185ea5b7de03481fdafe28899e2f3a09dec", + "size_bytes": 15050 + } + ], + "tree": "output", + "events": 64, + "selection": "none", + "branches": [ + "jet_pt", "jet_eta", "jet_phi", "jet_btag", "ph_pt", "ph_eta", + "ph_phi", "ele_pt", "ele_eta", "ele_phi", "ele_charge", "mu_pt", + "mu_eta", "mu_phi", "mu_charge", "MET_met", "MET_phi", "weight", + "Number" + ], + "note": "Single-class preprocessing fixture; not a complete multiclass campaign." +} diff --git a/validation/manifests/environments.json b/validation/manifests/environments.json new file mode 100644 index 0000000000000000000000000000000000000000..251a4c41f37ba9a1a00ee9553491f11b7b3779b7 --- /dev/null +++ b/validation/manifests/environments.json @@ -0,0 +1,20 @@ +{ + "legacy": { + "conda_environment": "dgl", + "python": "3.8.17", + "torch": "2.0.1", + "dgl": "1.1.1+cu118", + "uproot": "5.3.7", + "awkward": "2.6.7", + "device": "cpu" + }, + "rewrite": { + "python": "3.12.13", + "torch": "2.2.2+cu121", + "dgl": "2.4.0+cu121", + "device": "cpu" + }, + "checkpoint": "legacy/root_gnn_dgl/Pretrained_GNN/multiclass_pretrained_model_12/model_epoch_71.pt", + "seed_policy": "fixed checkpoint; dropout=0; explicit torch seed 20260818 for transfer head", + "source_identity": "working-tree validation campaign; legacy tree unmodified" +} diff --git a/validation/manifests/legacy_config.yaml b/validation/manifests/legacy_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..d3127d49301a95f7e1a8f8e89838ec4151eb3d05 --- /dev/null +++ b/validation/manifests/legacy_config.yaml @@ -0,0 +1,4 @@ +# Portable scientific settings from the active legacy pretraining workflow. +source: legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml +model: {type: root_gnn, hid_size: 64, in_size: 7, out_size: 12, n_layers: 4, n_proc_steps: 4, dropout: 0.0} +data: {tree_name: output, batch_size: 1024, fold_count: 4, feature_schema: [pt, eta, phi, energy, btag, charge, node_type]} diff --git a/validation/manifests/multiclass_fixture.json b/validation/manifests/multiclass_fixture.json new file mode 100644 index 0000000000000000000000000000000000000000..4d5b4d42b5f83eda4aa741028161e758c6e180ec --- /dev/null +++ b/validation/manifests/multiclass_fixture.json @@ -0,0 +1,28 @@ +{ + "source": "https://huggingface.co/datasets/HWresearch/Delphes", + "revision": "main", + "tree": "output", + "events_per_class": 8, + "events": 96, + "classes": [ + "ttH_NLO_inc", "tHjb_NLO_inc", "ggF_NLO_inc", "VBF_NLO_inc", + "WH_NLO_inc", "ZH_NLO_inc", "ttyy", "tttt", "SingleT_schan", + "ttbar", "ttW", "ttt" + ], + "repository_paths": [ + "samples/higgs/top-associated/tth/ttH_NLO_inc.root", + "samples/higgs/top-associated/thjb/tHjb_NLO_inc.root", + "samples/higgs/ggf/ggF_NLO_inc.root", + "samples/higgs/vbf/VBF_NLO_inc.root", + "samples/higgs/vh/wh/WH_NLO_inc.root", + "samples/higgs/vh/zh/ZH_NLO_inc.root", + "samples/photon/ttyy/ttyy.root", + "samples/top/multitop/tttt.root", + "samples/top/single-top/SingleT_schan.root", + "samples/top/ttbar/ttbar.root", + "samples/top/ttv/ttW.root", + "samples/top/multitop/ttt.root" + ], + "selection": "first 8 entries from each source file", + "construction": "temporary local composite with validation_label branch; source values are unchanged" +} diff --git a/validation/manifests/rewrite_config.yaml b/validation/manifests/rewrite_config.yaml new file mode 100644 index 0000000000000000000000000000000000000000..6641b27260ded4cb3805bd657ed2bae4bb28a572 --- /dev/null +++ b/validation/manifests/rewrite_config.yaml @@ -0,0 +1,5 @@ +# Semantic rewrite counterpart; paths and device are environmental. +model: {type: root_gnn, hid_size: 64, in_size: 7, out_size: 12, n_layers: 4, n_proc_steps: 4, dropout: 0.0} +task: {type: pretraining_multiclass} +data: {batch_size: 1024, split: {fold_count: 4}} +environment: {device: cpu, num_workers: 0} diff --git a/validation/run_full_validation.py b/validation/run_full_validation.py new file mode 100644 index 0000000000000000000000000000000000000000..92b3d049de7b70c0a4ae4fc32d4127450646d3e6 --- /dev/null +++ b/validation/run_full_validation.py @@ -0,0 +1,86 @@ +#!/usr/bin/env python3 +"""Run or compare staged end-to-end validation artifacts.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +from validation.artifacts import ValidationArtifact, save_artifact +from validation.compare import compare_artifacts, write_reports + + +def _smoke_artifact() -> ValidationArtifact: + nodes = np.asarray( + [ + [10, 0.2, 0.1, 10.2, 0, 1, 0], + [8, -0.3, -0.2, 8.4, 1, -1, 1], + [4, 0.1, 2.8, 4.1, 0, 0, 2], + ], + dtype=np.float32, + ) + edges = np.asarray([[0, 1], [0, 2], [1, 0], [1, 2], [2, 0], [2, 1]], dtype=np.int64) + edge_features = np.asarray( + [ + [0.1, 0.3, 0.4], + [0.1, -2.7, 2.8], + [-0.1, -0.3, 0.4], + [-0.4, 2.7, 2.8], + [-0.1, 2.7, 2.8], + [0.4, -2.7, 2.8], + ], + dtype=np.float32, + ) + return ValidationArtifact( + sample_id=np.asarray(["smoke:0"]), + labels=np.asarray([1]), + folds=np.asarray([0]), + weights=np.asarray([1.0], dtype=np.float32), + globals=np.asarray([[0.5, 1.0]], dtype=np.float32), + node_features_flat=nodes, + node_offsets=np.asarray([0, 3]), + edge_src_flat=edges[:, 0], + edge_dst_flat=edges[:, 1], + edge_features_flat=edge_features, + edge_offsets=np.asarray([0, 6]), + logits=np.asarray([[0.25]], dtype=np.float32), + scores=np.asarray([[0.562]], dtype=np.float32), + predictions=np.asarray([1]), + manifest={"mode": "smoke", "legacy": "not_run"}, + ) + + +def main() -> int: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--legacy-artifact", type=Path) + parser.add_argument("--rewrite-artifact", type=Path) + parser.add_argument( + "--output", type=Path, default=Path("validation_output/reports") + ) + parser.add_argument( + "--smoke", action="store_true", help="write a tiny synthetic rewrite artifact" + ) + args = parser.parse_args() + if args.smoke: + path = args.output.parent / "rewrite" + save_artifact(_smoke_artifact(), path) + print(f"wrote smoke artifact: {path}") + return 0 + if not args.legacy_artifact or not args.rewrite_artifact: + parser.error("comparison requires --legacy-artifact and --rewrite-artifact") + report = compare_artifacts(args.legacy_artifact, args.rewrite_artifact) + write_reports(report, args.output) + print( + json.dumps( + {"overall_status": report["overall_status"], "report": str(args.output)}, + indent=2, + ) + ) + return 0 if report["overall_status"] == "PASS" else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/tolerances.yaml b/validation/tolerances.yaml new file mode 100644 index 0000000000000000000000000000000000000000..a3fd03600bd563d72feb6eed8d119a02db83b0fe --- /dev/null +++ b/validation/tolerances.yaml @@ -0,0 +1,13 @@ +# Strict CPU defaults for the staged campaign. +feature_atol: 1.0e-6 +feature_rtol: 1.0e-6 +edge_atol: 1.0e-6 +edge_rtol: 1.0e-6 +logit_atol: 1.0e-5 +logit_rtol: 1.0e-5 +gradient_atol: 1.0e-5 +gradient_rtol: 1.0e-5 +parameter_atol: 1.0e-5 +parameter_rtol: 1.0e-5 +metric_atol: 1.0e-6 +metric_rtol: 1.0e-6 diff --git a/validation/train_step.py b/validation/train_step.py new file mode 100644 index 0000000000000000000000000000000000000000..9227d8ec95073d8587ad62863090060929643261 --- /dev/null +++ b/validation/train_step.py @@ -0,0 +1,111 @@ +#!/usr/bin/env python3 +"""Run one deterministic Adam step for legacy or rewritten ROOT-GNN.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import dgl +import numpy as np +import torch + +from validation.artifacts import load_artifact +from validation.forward import _graph + + +def _weighted_loss(logits, labels, weights): + elementwise = torch.nn.functional.cross_entropy(logits, labels, reduction="none") + result = logits.new_zeros(()) + for label in torch.unique(labels): + mask = labels == label + result = ( + result + (weights[mask] * elementwise[mask]).sum() / weights[mask].sum() + ) + return result / len(torch.unique(labels)) + + +def _flat_state(model, *, gradients: bool = False): + values = [] + for name, parameter in sorted(model.named_parameters()): + if gradients: + if parameter.grad is None: + raise RuntimeError(f"missing gradient for {name}") + values.append(parameter.grad.detach().cpu().reshape(-1)) + else: + values.append(parameter.detach().cpu().reshape(-1)) + return torch.cat(values).numpy() + + +def _flat_optimizer_state(model, optimizer, key): + values = [] + steps = [] + for name, parameter in sorted(model.named_parameters()): + state = optimizer.state[parameter] + values.append(state[key].detach().cpu().reshape(-1)) + steps.append(float(state["step"])) + return torch.cat(values).numpy(), np.asarray(steps, dtype=np.float64) + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", choices=("legacy", "rewrite"), required=True + ) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--events", type=int, default=8) + parser.add_argument("--steps", type=int, default=1) + args = parser.parse_args() + artifact = load_artifact(args.artifact) + count = min(args.events, artifact.event_count) + graphs = [_graph(artifact, index) for index in range(count)] + graph = dgl.batch(graphs) + empty_globals = torch.empty((count, 0), dtype=torch.float32) + if args.implementation == "legacy": + from models.GCN import Edge_Network + + model = Edge_Network(graphs[0], empty_globals[:1], 64, 12, 4, 4, dropout=0.0) + payload = torch.load(args.checkpoint, map_location="cpu") + model.load_state_dict(payload["model_state_dict"]) + else: + from gnn4colliders.models.root_gnn import ( + EdgeNetwork, + load_legacy_edge_network_state_dict, + ) + + model = EdgeNetwork(graphs[0], empty_globals[:1], 64, 12, 4, 4, dropout=0.0) + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + load_legacy_edge_network_state_dict(model, payload) + model.train() + labels = torch.from_numpy(artifact.labels[:count]).to(torch.long) + weights = torch.from_numpy(artifact.weights[:count]).to(torch.float32) + optimizer = torch.optim.Adam(model.parameters(), lr=1e-4) + histories = [] + for _ in range(args.steps): + optimizer.zero_grad() + logits = model(graph, None) + loss = _weighted_loss(logits, labels, weights) + loss.backward() + gradients = _flat_state(model, gradients=True) + optimizer.step() + histories.append(float(loss.detach())) + exp_avg, optimizer_steps = _flat_optimizer_state(model, optimizer, "exp_avg") + exp_avg_sq, _ = _flat_optimizer_state(model, optimizer, "exp_avg_sq") + np.savez_compressed( + args.output, + logits=logits.detach().cpu().numpy(), + loss=np.asarray(float(loss.detach())), + history=np.asarray(histories, dtype=np.float64), + gradients=gradients, + parameters=_flat_state(model), + optimizer_steps=optimizer_steps, + exp_avg=exp_avg, + exp_avg_sq=exp_avg_sq, + ) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/validation/transfer.py b/validation/transfer.py new file mode 100644 index 0000000000000000000000000000000000000000..f8cf2e9485845e8c91220db94840b8d2004edf69 --- /dev/null +++ b/validation/transfer.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Compare legacy/new transfer-learning forward with a fixed binary head.""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import dgl +import numpy as np +import torch + +from validation.artifacts import load_artifact +from validation.forward import _graph + +LEGACY_CONFIG = { + "module": "models.GCN", + "class": "Edge_Network", + "args": { + "hid_size": 64, + "out_size": 12, + "n_layers": 4, + "n_proc_steps": 4, + "dropout": 0.0, + }, +} + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument( + "--implementation", choices=("legacy", "rewrite"), required=True + ) + parser.add_argument("--artifact", type=Path, required=True) + parser.add_argument("--checkpoint", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--events", type=int, default=96) + parser.add_argument("--trainable-backbone", action="store_true") + args = parser.parse_args() + artifact = load_artifact(args.artifact) + count = min(args.events, artifact.event_count) + graphs = [_graph(artifact, index) for index in range(count)] + batch_graph = dgl.batch(graphs) + empty_globals = torch.empty((1, 0), dtype=torch.float32) + if args.implementation == "legacy": + from models.GCN import Transferred_Learning_Finetuning + + model = Transferred_Learning_Finetuning( + str(args.checkpoint), + LEGACY_CONFIG, + graphs[0], + empty_globals, + 64, + 1, + 4, + 4, + dropout=0.0, + frozen_pretraining=not args.trainable_backbone, + ) + classifier = model.classify + else: + from gnn4colliders.models.root_gnn import EdgeNetwork, FineTunedEdgeNetwork + + backbone = EdgeNetwork(graphs[0], empty_globals, 64, 12, 4, 4, dropout=0.0) + payload = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + from gnn4colliders.models.root_gnn import load_legacy_edge_network_state_dict + + load_legacy_edge_network_state_dict(backbone, payload) + model = FineTunedEdgeNetwork( + backbone, 1, freeze_backbone=not args.trainable_backbone + ) + classifier = model.classifier + torch.manual_seed(20260818) + classifier.load_state_dict(torch.nn.Linear(64, 1).state_dict()) + model.eval() + with torch.inference_mode(): + logits = model(batch_graph, None).cpu().numpy() + args.output.parent.mkdir(parents=True, exist_ok=True) + np.savez_compressed(args.output, logits=logits) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())