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