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