docs/migration.md

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, dataset.py:471-482, utils.py:121-143). 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). 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);
• string/tuple selections and cutflow (check_selection, selection_mask, compute_cutflow, dataset.py:75-158);
• fold masks and cache suffixes (fold_selection, fold_selection_name, utils.py:121-143);
• deterministic chunk partitioning (hash_partition, batched_dataset.py:27-31).

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).

Then implement DGL .bin serialization, lazy chunk loading, pre-batching, and padding. Compare against RootDataset.save/load, LazyDataset, and PreBatchedDataset (dataset.py:396-469, batched_dataset.py:129-174). 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, GCN.py:182-251).

Next port Transferred_Learning_Finetuning, including pretrained model_state_dict loading, removal of the final classifier, and new classifier initialization (GCN.py:884-997). 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). Add parity cases for positive, zero, and negative weights and binary versus multiclass shapes.

Port metric behavior from training_script.py:438-510: 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). Support legacy DDP/compiled prefixes (module. and _orig_mod.) as exercised by checkpoint lookup and inference (utils.py:145-248, inference.py:274-290). Port EarlyStop state and log parsing separately (utils.py:325-390). 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

• Load a checked-in or generated multiclass pretrained checkpoint.
• Load a legacy fine-tuning checkpoint after prefix normalization.
• Resume optimizer and early-stop state.
• Produce equivalent logits on a deterministic graph fixture.
• Produce equivalent .npz score, label, and metadata fields.
• Preserve ROOT scalar/vector score branch conventions in the Python adapter.

Known risks are documented in 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 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

• active ROOT data path and graph cache
• validated feature, graph, model, task, and metric behavior
• train from scratch and fine-tune a pretrained backbone
• resume new-stack checkpoints and load supported historical weights
• evaluate and predict named outputs
• single-process and DDP application boundaries
• Perlmutter/Slurm launcher examples and profiling guidance
• 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.