ho22joshua commited on
Commit
a76eccd
Β·
1 Parent(s): 082be01

Document root GNN architecture and migration

Browse files
Files changed (3) hide show
  1. AGENTS.md +15 -4
  2. docs/architecture.md +213 -17
  3. docs/migration.md +134 -12
AGENTS.md CHANGED
@@ -2,10 +2,21 @@
2
 
3
  ## Goal
4
 
5
- Rewrite the ML system in `legacy/` into a clean implementation in `src/`.
 
 
 
6
 
7
- Do not modify `legacy/` unless explicitly asked. It is the behavioral
8
- reference implementation.
 
 
 
 
 
 
 
 
9
 
10
  ## Priorities
11
 
@@ -63,4 +74,4 @@ For non-trivial work:
63
  6. Run them.
64
  7. Summarize changes and remaining migration work.
65
 
66
- Do not rewrite the entire repository in one pass.
 
2
 
3
  ## Goal
4
 
5
+ Rewrite the main ML system in `legacy/root_gnn_dgl/` into a clean
6
+ implementation in `src/`. The sibling `legacy/physicsnemo/` tree is a prior
7
+ rewrite attempt: it may provide design inspiration, but it is not the target
8
+ behavioral reference.
9
 
10
+ Do not modify `legacy/` unless explicitly asked. `legacy/root_gnn_dgl/` is the
11
+ behavioral reference implementation for this project.
12
+
13
+ ## Reverse-engineering scope
14
+
15
+ The primary executable pipeline is `legacy/root_gnn_dgl/scripts/` plus
16
+ `legacy/root_gnn_dgl/root_gnn_base/`, `models/`, and the YAML configurations.
17
+ Documentation and parity work should cite those files and symbols directly.
18
+ Treat `legacy/physicsnemo/` as out of scope except when explicitly comparing
19
+ its abstractions as rewrite inspiration.
20
 
21
  ## Priorities
22
 
 
74
  6. Run them.
75
  7. Summarize changes and remaining migration work.
76
 
77
+ Do not rewrite the entire repository in one pass.
docs/architecture.md CHANGED
@@ -1,17 +1,213 @@
1
- # Target architecture
2
-
3
- The new implementation is organized by responsibility:
4
-
5
- - `src/gnn4colliders/`: importable production package
6
- - `configs/`: declarative experiment and data configuration
7
- - `scripts/`: thin command-line entry points
8
- - `tests/unit/`: focused component tests
9
- - `tests/integration/`: end-to-end component integration tests
10
- - `tests/parity/`: comparisons between the new implementation and `legacy/`
11
- - `notebooks/`: exploratory analysis only
12
- - `data/fixtures/`: small deterministic test inputs
13
-
14
- Dataset loading, preprocessing, model definition, training, evaluation, and
15
- serialization should remain separate modules. Randomness must be explicit and
16
- seedable, and configuration should be passed into application code rather than
17
- hardcoded in training logic.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Reverse-engineered architecture: `root_gnn_dgl`
2
+
3
+ This document covers the target system in `legacy/root_gnn_dgl/`. The sibling
4
+ `legacy/physicsnemo/` tree is a prior rewrite attempt and is not a behavioral
5
+ target.
6
+
7
+ ## 1. High-level system description
8
+
9
+ `root_gnn_dgl` is a ROOT-to-DGL graph classification system. YAML selects
10
+ dataset, model, loss, and finish-function classes by import path. The dataset
11
+ reads ROOT trees, converts collider objects to fully connected DGL graphs, and
12
+ saves graph chunks. Training loads those chunks, applies fold selection and
13
+ optional pre-batching/padding, trains a graph network, writes one PyTorch
14
+ checkpoint per epoch, and reports weighted loss, accuracy, and ROC AUC.
15
+
16
+ The primary model is `models.GCN.Edge_Network` ([`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
17
+ It encodes node, edge, and global features, repeats edge -> node -> global
18
+ message passing `n_proc_steps` times, decodes the global state, and applies
19
+ `classify`. Fine-tuning uses `models.GCN.Transferred_Learning_Finetuning`
20
+ ([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)), which loads a
21
+ pretrained `Edge_Network`, removes its final classifier, and applies a new one.
22
+ The active configs use output size 12 for multiclass pretraining and output
23
+ size 1 for binary tasks ([`configs/stats_100K/pretraining_multiclass.yaml:1-45`](../legacy/root_gnn_dgl/configs/stats_100K/pretraining_multiclass.yaml),
24
+ [`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)).
25
+
26
+ ### Entry points and flows
27
+
28
+ - `scripts/training_script.py:main` and its CLI parser load YAML, create
29
+ loaders, construct the model, and call `train`; `--evaluate` calls
30
+ `evaluate` ([`training_script.py:638-843`](../legacy/root_gnn_dgl/scripts/training_script.py)).
31
+ - `scripts/prep_data.py:main` creates configured graph caches
32
+ ([`prep_data.py:68-110`](../legacy/root_gnn_dgl/scripts/prep_data.py)).
33
+ - `scripts/inference.py:main` reconstructs an unlazy dataset, loads one or
34
+ more checkpoints, and writes `.npz` or ROOT scores
35
+ ([`inference.py:163-387`](../legacy/root_gnn_dgl/scripts/inference.py)).
36
+ - `scripts/export_onnx.py:main` exports an ONNX-friendly model
37
+ ([`export_onnx.py:979-1035`](../legacy/root_gnn_dgl/scripts/export_onnx.py)).
38
+ - `selections.py:main`, `check_dataset_files.py:main`, and
39
+ `plot_config_distributions.py:main` are diagnostic entry points. `run_demo.sh`
40
+ sequences pretraining, binary training, fine-tuning, and inference
41
+ ([`run_demo.sh:3-59`](../legacy/root_gnn_dgl/run_demo.sh)).
42
+
43
+ The training flow is:
44
+
45
+ ```text
46
+ YAML -> load_config/buildFromConfig -> RootDataset/LazyDataset
47
+ -> ROOT/Awkward -> DGL graph + labels/tracking/globals -> .bin cache
48
+ -> fold_selection -> prebatch/padding -> GraphDataLoader
49
+ -> Edge_Network or transfer model -> weighted loss/metrics
50
+ -> model_epoch_N.pt, logs, evaluation/inference output
51
+ ```
52
+
53
+ `training_script.train` is the lifecycle implementation
54
+ ([`training_script.py:143-614`](../legacy/root_gnn_dgl/scripts/training_script.py));
55
+ distributed paths use NCCL/DDP ([`training_script.py:616-839`](../legacy/root_gnn_dgl/scripts/training_script.py)).
56
+ Inference uses `CustomPreBatchedDataset`, applies a configured finish function,
57
+ and collects `scores`, `labels`, and `tracking_info`
58
+ ([`inference.py:20-76`](../legacy/root_gnn_dgl/scripts/inference.py),
59
+ [`inference.py:223-325`](../legacy/root_gnn_dgl/scripts/inference.py)).
60
+
61
+ ## 2. Dependency and data-flow map
62
+
63
+ ```text
64
+ ROOT files (raw_dir/file_names, tree_name)
65
+ -> RootDataset.process (uproot/Awkward; dataset.py:221-320)
66
+ -> node_features_from_tree (dataset.py:15-50)
67
+ -> full_connected_graph (dataset.py:52-59)
68
+ -> EdgeDataset.make_graph: [deta, dphi, dR] (dataset.py:471-482)
69
+ -> DGL .bin cache / LazyDataset / PreBatchedDataset
70
+ -> GraphDataLoader -> models.GCN -> loss/metrics -> outputs
71
+ ```
72
+
73
+ `RootDataset` provides `process`, `save`, `load`, `__getitem__`, and `__len__`
74
+ ([`dataset.py:160-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
75
+ `LazyDataset` loads one chunk through a ring buffer
76
+ ([`dataset.py:525-578`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
77
+ `PreBatchedDataset.process` selects, shuffles, batches, pads, and caches
78
+ ([`batched_dataset.py:34-146`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
79
+
80
+ `load_config` uses PyYAML `FullLoader` and shallow `include` merging, while
81
+ `buildFromConfig` dynamically imports `module`, resolves `class`, merges extra
82
+ keys into `args`, converts list-valued weights to tensors, and injects runtime
83
+ arguments ([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
84
+ This reflection shape is a de facto interface for configured components.
85
+
86
+ ### Data and preprocessing
87
+
88
+ The active node schema is seven columns: `pt`, `eta`, `phi`, `energy`, `btag`,
89
+ `charge`, and `node_type`. `CALC_E` is `pt*cosh(eta)`, constants are broadcast
90
+ per object type, `NODE_TYPE` is an integer type code, and feature scales are
91
+ applied columnwise ([`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
92
+ `full_connected_graph` makes directed all-pairs edges; `EdgeDataset` requests
93
+ no self-loops and stores `[deta, dphi, dR]`
94
+ ([`dataset.py:52-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
95
+ [`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
96
+
97
+ Selections are strings evaluated with builtins disabled or
98
+ `(variable, cut, operator)` triples (`check_selection`, `selection_mask`;
99
+ [`dataset.py:75-145`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)). Fold
100
+ selection uses `tracking[:,0] % n_folds`; tracking column 0 is fold and column 1
101
+ is weight ([`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
102
+ [`dataset.py:176-182`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
103
+ `hash_partition` and a seeded Torch generator control pre-batch order
104
+ ([`batched_dataset.py:27-99`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
105
+ Padding modes are `NONE`, `STEPS`, `FIXED`, and `NODE`; `FIXED` is hardcoded to
106
+ 16,000 nodes and 104,000 edges ([`batched_dataset.py:100-125`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
107
+
108
+ ### Model, losses, and metrics
109
+
110
+ `Make_MLP` builds linear/ReLU/dropout blocks followed by LayerNorm
111
+ ([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py)). Each `Edge_Network`
112
+ step encodes inputs, copies source/destination states to edges, updates edges,
113
+ sums edge messages into nodes, updates nodes, then mean-pools nodes/edges to
114
+ update globals ([`GCN.py:195-249`](../legacy/root_gnn_dgl/models/GCN.py)). It
115
+ returns logits `[graphs, out_size]` without sigmoid/softmax.
116
+
117
+ The default objective is elementwise `BCEWithLogitsLoss`, multiplied by
118
+ `tracking[:,1]`, averaged separately per unique label, then averaged across
119
+ labels ([`training_script.py:143-185`](../legacy/root_gnn_dgl/scripts/training_script.py),
120
+ [`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
121
+ `--abs` makes weights positive. Binary metrics use sigmoid threshold 0.5 and
122
+ weighted ROC AUC; multiclass metrics use argmax and one-vs-rest ROC AUC
123
+ ([`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py)).
124
+ Additional configurable losses/finishers live in `models/loss.py`
125
+ ([`loss.py:6-310`](../legacy/root_gnn_dgl/models/loss.py)).
126
+
127
+ ### Checkpoints and outputs
128
+
129
+ Training writes `Training_Directory/model_epoch_<epoch>.pt` containing `epoch`,
130
+ `model_state_dict`, `optimizer_state_dict`, and serialized `early_stop`
131
+ ([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
132
+ Keys strip `module.` and compiled models save the underlying `_orig_mod` state;
133
+ `get_last_epoch`, `get_specific_epoch`, and `get_best_epoch` load the files
134
+ ([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
135
+ `evaluate` writes `evaluation_<epoch>.npz`; inference writes `.npz` fields
136
+ `scores`, `labels`, `tracking_info`, or adds score branches and `selection_pass`
137
+ to a cloned ROOT tree ([`training_script.py:57-140`](../legacy/root_gnn_dgl/scripts/training_script.py),
138
+ [`inference.py:328-385`](../legacy/root_gnn_dgl/scripts/inference.py)).
139
+
140
+ ## Randomness, external services, and coupling
141
+
142
+ The CLI exposes `--seed`, but `main` passes it to model construction rather
143
+ than globally seeding Python, NumPy, or Torch
144
+ ([`training_script.py:638-753`](../legacy/root_gnn_dgl/scripts/training_script.py)).
145
+ Pre-batching has explicit seeds, but `AugmentedDataset` mutates the process-wide
146
+ NumPy seed ([`dataset.py:716-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
147
+ clustering uses unseeded `torch.randperm`/`randint` (`loss.py:259-295`), and
148
+ model reset/fine-tuning hardcodes `torch.manual_seed(2)`
149
+ ([`GCN.py:58-65`](../legacy/root_gnn_dgl/models/GCN.py),
150
+ [`GCN.py:900-915`](../legacy/root_gnn_dgl/models/GCN.py)). CUDA kernels, DDP,
151
+ and DataLoader behavior are not made deterministic.
152
+
153
+ Implicit coupling includes repository-relative `sys.path` insertion
154
+ ([`training_script.py:14-20`](../legacy/root_gnn_dgl/scripts/training_script.py)),
155
+ dynamic imports, mutable default lists/dicts, global `FEATURE_DTYPE`
156
+ ([`dataset.py:13`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
157
+ in-place `tracking_info` mutation ([`dataset.py:176-181`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)),
158
+ and in-place DGL graph mutation during forward.
159
+
160
+ The environment assumes Python 3.8, PyTorch 2.0.1, CUDA 11.8, DGL 1.1.1,
161
+ ROOT, Awkward, Uproot, PyYAML, and scikit-learn
162
+ ([`setup/environment.yml:1-10`](../legacy/root_gnn_dgl/setup/environment.yml),
163
+ [`setup/environment.yml:240-295`](../legacy/root_gnn_dgl/setup/environment.yml)).
164
+ Standard configs assume `/global/cfs/` and `/pscratch/` paths, CUDA/NCCL,
165
+ Slurm, and optionally Podman-HPC. `setup/download_data.sh` downloads the
166
+ external Hugging Face dataset `HWresearch/Delphes`
167
+ ([`download_data.sh:13-67`](../legacy/root_gnn_dgl/setup/download_data.sh)).
168
+
169
+ ## Apparent unused or secondary code
170
+
171
+ Not selected by the standard stats/Delphes configs, or only reachable from
172
+ optional workflows, are `GCN_global`, `GCN_global_2way`, most transfer variants,
173
+ attention models, `MultiModel`, and `Clustering`
174
+ ([`GCN.py:122-1933`](../legacy/root_gnn_dgl/models/GCN.py)); `UprootDataset`,
175
+ `tHbbEdgeDataset`, `AugmentedDataset`, and photon-ID paths
176
+ ([`uproot_dataset.py:10-31`](../legacy/root_gnn_dgl/root_gnn_base/uproot_dataset.py),
177
+ [`dataset.py:484-827`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
178
+ [`photon_ID_dataset.py:1-33`](../legacy/root_gnn_dgl/root_gnn_base/photon_ID_dataset.py));
179
+ optional loss and similarity utilities; and the no-op
180
+ `root_gnn_base.utils.graph_augmentation` ([`utils.py:393-395`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)).
181
+ The main path evaluates `test_loaders`; validation loaders are only assembled
182
+ when a config has a validation fold ([`training_script.py:682-747`](../legacy/root_gnn_dgl/scripts/training_script.py)).
183
+
184
+ ## 3. De facto interfaces to preserve
185
+
186
+ 1. YAML `module`, `class`, `args`, plus runtime `sample_graph` and
187
+ `sample_global` injection.
188
+ 2. Dataset items `(DGLGraph, label, tracking, global_features)`.
189
+ 3. `ndata['features']`, `edata['features']`, seven node columns, and three edge
190
+ columns in `[deta, dphi, dR]` order.
191
+ 4. Tracking column 0 fold and column 1 weight semantics.
192
+ 5. `model(graph, global_feats)`, logits shape `[batch, out_size]`, and
193
+ `representation` where used.
194
+ 6. Weighted per-label loss, metric thresholds, checkpoint keys/prefix cleanup,
195
+ epoch filenames, and `.npz`/ROOT output fields.
196
+
197
+ ## 4. Ambiguous behavior
198
+
199
+ - Historical edge order/self-loop expectations; empty and padding graph inputs.
200
+ - Whether negative weights are meaningful or should always be absolute.
201
+ - Whether β€œvalidation” is intended to differ from the active test-loader path.
202
+ - Shape semantics of multi-label finishers and experimental transfer classes.
203
+ - Whether chunk IDs must match historical `np.array_split` boundaries.
204
+ - Required behavior for missing branches and dynamic selection expressions.
205
+
206
+ ## 5. Recommended rewrite boundaries
207
+
208
+ Separate typed configuration; ROOT/Awkward I/O; selections/folds/features/
209
+ edges; DGL cache/lazy loading/batching; active models and checkpoint adapters;
210
+ objectives/metrics; training lifecycle; and inference/ONNX applications.
211
+ Establish parity for the active `LazyDataset -> PreBatchedDataset ->
212
+ Edge_Network` binary/multiclass path first. Add experimental classes only when
213
+ a config or consumer proves they are required.
docs/migration.md CHANGED
@@ -1,16 +1,138 @@
1
- # Migration plan
2
 
3
- `legacy/` is the frozen behavioral reference. Do not modify it during the
4
- rewrite. Migrate one coherent capability at a time into `src/gnn4colliders`.
 
5
 
6
- For each migrated capability:
 
 
7
 
8
- 1. Identify the legacy inputs, outputs, shapes, dtypes, and metric semantics.
9
- 2. Define a small typed interface in the new package.
10
- 3. Add unit tests for the new implementation.
11
- 4. Add parity tests comparing representative deterministic fixtures with the
12
- legacy behavior.
13
- 5. Record any intentional incompatibility and checkpoint implications here.
14
 
15
- The initial migration order is dataset loading and preprocessing, graph/model
16
- components, training and evaluation, then checkpoint and command-line support.
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ # Incremental migration plan: `root_gnn_dgl`
2
 
3
+ The target is `legacy/root_gnn_dgl/`. `legacy/physicsnemo/` is a prior rewrite
4
+ attempt and may inspire abstractions, but it is not a parity target. Neither
5
+ legacy tree should be modified during migration.
6
 
7
+ Each phase should add focused unit tests, a deterministic fixture in
8
+ `data/fixtures/`, and parity tests under `tests/parity/` before moving upward.
9
+ Record intentional differences and checkpoint consequences here.
10
 
11
+ ## Phase 0 β€” freeze observations and fixtures
 
 
 
 
 
12
 
13
+ Capture a small representative ROOT-equivalent fixture containing the seven
14
+ active node features, three edge features, labels, fold values, weights, and
15
+ globals. Record outputs of `node_features_from_tree`, `full_connected_graph`,
16
+ `EdgeDataset.make_graph`, and `fold_selection`
17
+ ([`dataset.py:15-59`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
18
+ [`dataset.py:471-482`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
19
+ [`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Preserve
20
+ one `.bin`, one `model_epoch_N.pt`, one evaluation `.npz`, and one inference
21
+ `.npz` fixture if available.
22
+
23
+ ## Phase 1 β€” configuration boundary
24
+
25
+ Implement a typed configuration layer that reads `Training`, `Model`,
26
+ optional `Loss`, and `Datasets`. Initially retain a compatibility adapter for
27
+ `module`/`class`/`args` and runtime injection of `sample_graph` and
28
+ `sample_global`, matching `buildFromConfig`
29
+ ([`utils.py:10-43`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Keep
30
+ dynamic imports isolated at this boundary rather than spreading reflection
31
+ through new code.
32
+
33
+ ## Phase 2 β€” pure preprocessing parity
34
+
35
+ Port and test, in isolation:
36
+
37
+ - branch-to-node conversion, `CALC_E`, `NODE_TYPE`, constants, scaling, empty
38
+ objects, and dtypes (`node_features_from_tree`,
39
+ [`dataset.py:15-50`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
40
+ - string/tuple selections and cutflow (`check_selection`, `selection_mask`,
41
+ `compute_cutflow`, [`dataset.py:75-158`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py));
42
+ - fold masks and cache suffixes (`fold_selection`, `fold_selection_name`,
43
+ [`utils.py:121-143`](../legacy/root_gnn_dgl/root_gnn_base/utils.py));
44
+ - deterministic chunk partitioning (`hash_partition`,
45
+ [`batched_dataset.py:27-31`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
46
+
47
+ This is the highest-value parity layer: model parity is invalid if graph inputs
48
+ differ.
49
+
50
+ ## Phase 3 β€” graph construction and cache format
51
+
52
+ Implement graph construction with tests for node/edge counts, directed edge
53
+ ordering, self-loop policy, `[deta, dphi, dR]` order, metadata, and empty graphs.
54
+ Preserve the dataset item contract `(graph, label, tracking, global_features)`
55
+ from `RootDataset.__getitem__`
56
+ ([`dataset.py:465-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py)).
57
+
58
+ Then implement DGL `.bin` serialization, lazy chunk loading, pre-batching, and
59
+ padding. Compare against `RootDataset.save/load`, `LazyDataset`, and
60
+ `PreBatchedDataset` ([`dataset.py:396-469`](../legacy/root_gnn_dgl/root_gnn_base/dataset.py),
61
+ [`batched_dataset.py:129-174`](../legacy/root_gnn_dgl/root_gnn_base/batched_dataset.py)).
62
+ Treat `NONE`, `STEPS`, `FIXED`, and `NODE` as explicit features; do not hide
63
+ the hardcoded fixed padding sizes.
64
+
65
+ ## Phase 4 β€” active model parity
66
+
67
+ Port `models.GCN.Edge_Network` first. Preserve constructor parameters,
68
+ `forward(graph, global_feats)`, feature keys, processor order, MLP LayerNorm
69
+ placement, and logits shape. Compare intermediate and final tensors on fixed
70
+ graphs using the legacy architecture
71
+ ([`GCN.py:18-35`](../legacy/root_gnn_dgl/models/GCN.py),
72
+ [`GCN.py:182-251`](../legacy/root_gnn_dgl/models/GCN.py)).
73
+
74
+ Next port `Transferred_Learning_Finetuning`, including pretrained
75
+ `model_state_dict` loading, removal of the final classifier, and new classifier
76
+ initialization ([`GCN.py:884-997`](../legacy/root_gnn_dgl/models/GCN.py)). Test
77
+ both frozen and unfrozen modes. Defer other model classes until an active
78
+ config or consumer proves they are needed.
79
+
80
+ ## Phase 5 β€” objectives and metrics
81
+
82
+ Implement the default objective exactly: elementwise configured loss,
83
+ tracking-column weights, per-unique-label normalization, and averaging across
84
+ labels ([`training_script.py:320-359`](../legacy/root_gnn_dgl/scripts/training_script.py)).
85
+ Add parity cases for positive, zero, and negative weights and binary versus
86
+ multiclass shapes.
87
+
88
+ Port metric behavior from
89
+ [`training_script.py:438-510`](../legacy/root_gnn_dgl/scripts/training_script.py):
90
+ sigmoid threshold 0.5, argmax, weight masking, weighted ROC AUC, one-vs-rest
91
+ multiclass AUC, and NaN behavior when AUC is undefined. Add `models/loss.py`
92
+ classes only with dedicated tests; do not substitute their reductions.
93
+
94
+ ## Phase 6 β€” checkpoint and lifecycle
95
+
96
+ Create a checkpoint adapter preserving `model_epoch_<epoch>.pt` and keys
97
+ `epoch`, `model_state_dict`, `optimizer_state_dict`, and `early_stop`
98
+ ([`training_script.py:565-604`](../legacy/root_gnn_dgl/scripts/training_script.py)).
99
+ Support legacy DDP/compiled prefixes (`module.` and `_orig_mod.`) as exercised
100
+ by checkpoint lookup and inference
101
+ ([`utils.py:145-248`](../legacy/root_gnn_dgl/root_gnn_base/utils.py),
102
+ [`inference.py:274-290`](../legacy/root_gnn_dgl/scripts/inference.py)). Port
103
+ `EarlyStop` state and log parsing separately
104
+ ([`utils.py:325-390`](../legacy/root_gnn_dgl/root_gnn_base/utils.py)). Verify
105
+ resume, restart, early termination, and `.npz` fields before distributed work.
106
+
107
+ ## Phase 7 β€” CLI, inference, and export
108
+
109
+ Build thin new applications around tested library interfaces in this order:
110
+
111
+ 1. preprocessing/cache generation (`scripts/prep_data.py`);
112
+ 2. training/evaluation (`scripts/training_script.py`);
113
+ 3. inference to `.npz` and ROOT (`scripts/inference.py`);
114
+ 4. ONNX export after PyTorch parity (`scripts/export_onnx.py`).
115
+
116
+ Use subprocess integration tests with tiny fixtures. Preserve CLI options only
117
+ where they serve an active workflow; document removed diagnostic/cluster-only
118
+ options.
119
+
120
+ ## Phase 8 β€” reproducibility and deployment
121
+
122
+ Introduce one explicit seed policy covering Python, NumPy, Torch, sampling,
123
+ augmentation, and clustering. Add deterministic-mode tests and document GPU
124
+ nondeterminism. Isolate Slurm/NCCL, Podman-HPC, ROOT, and Hugging Face data
125
+ download integrations behind adapters only after local behavior is stable.
126
+
127
+ ## Checkpoint compatibility checklist
128
+
129
+ - [ ] Load the checked-in multiclass pretrained checkpoint.
130
+ - [ ] Load a legacy fine-tuning checkpoint after prefix normalization.
131
+ - [ ] Resume optimizer and early-stop state.
132
+ - [ ] Produce equivalent logits on a deterministic graph fixture.
133
+ - [ ] Produce equivalent `.npz` score, label, and tracking fields.
134
+ - [ ] Preserve ROOT scalar/vector score branch conventions if ROOT output stays.
135
+
136
+ Known risks are documented in [`architecture.md`](architecture.md): edge order,
137
+ self-loops, weight semantics, validation/test naming, padding, dynamic
138
+ selection evaluation, reproducibility, and the experimental model/loss surface.