# 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_<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.