rewriting codebase (#7)

916755e 3 months ago

preview code

Raw History Blame

11.2 kB

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:

AGENTS.md
docs/architecture.md
docs/migration.md

Then inspect:

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:

node features
    ->
graph topology
    ->
edge indices
    ->
edge features
    ->
DGLGraph

for the active ROOT-GNN path.

The graph layer should live under:

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:

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:

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:

def fully_connected_edges(
    num_nodes: int,
    *,
    self_loops: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
    ...

and:

def build_edge_features(
    node_features: torch.Tensor,
    src: torch.Tensor,
    dst: torch.Tensor,
    *,
    eta_index: int,
    phi_index: int,
) -> torch.Tensor:
    ...

and eventually:

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:

num_edges = N * (N - 1)

where supported.

Explicitly test:

N = 1
N = 2
N = 3

For example, a two-node graph should contain:

0 -> 1
1 -> 0

and no:

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:

deta
dphi
dR

in exactly that order.

The expected semantics should be verified against the Task 3 characterization.

For each directed edge:

src -> dst

calculate the same signed deta and signed wrapped dphi convention used by the legacy implementation.

Then calculate:

dR = sqrt(deta^2 + dphi^2)

if confirmed by parity.

7. Phi wrapping

Implement phi differences explicitly.

Do not rely on naive:

phi_src - phi_dst

when crossing -pi / +pi.

Use the exact wrapping convention characterized from legacy behavior.

Add unit tests for cases including:

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:

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:

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:

topology calculation
edge-feature calculation
DGL object construction

Do not make every helper require a DGLGraph.

The DGL-specific builder should:

create graph topology
attach node features
attach edge features

using the active legacy key names if compatibility requires them.

For example, if legacy behavior uses:

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:

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:

N = 1

If the legacy graph has:

1 node
0 edges

ensure the new implementation behaves the same.

Verify the shape of empty edge features.

For example, expected behavior may be:

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:

tests/unit/graphs/

Suggested organization:

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:

legacy behavior

and:

new implementation matches legacy

where appropriate.

16. Numerical tolerances

Use tight explicit tolerances for floating-point calculations.

Prefer something around:

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:

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:

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:

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:

gnn4colliders.features

contains architecture-independent collider features, while:

gnn4colliders.graphs

contains graph-specific representation logic used by ROOT-GNN.

Keep documentation changes concise.

Validation

Run:

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:

uv run ruff check src tests
uv run ruff format --check src tests

Inspect:

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:

Graph topology is implemented under src/gnn4colliders/graphs/.
Directed fully connected no-self-loop behavior matches legacy.
Edge ordering matches legacy where required.
[deta, dphi, dR] matches legacy.
Phi wrapping matches legacy in both directions.
Node and edge feature attachment matches compatibility requirements.
Single-node behavior is covered.
Unit tests pass.
Legacy/new parity tests pass.
No DGL tests are skipped because DGL is unavailable.
No dataset, training, or model implementation has been added.
Legacy code remains untouched.

Completion report

Report:

files created
files modified
public API added
topology implemented
observed edge ordering
phi wrapping convention
DGL compatibility behavior
unit test results
parity results
numerical tolerances used
ambiguities or deferred behavior
validation commands and results

After validation succeeds, create one Git commit containing only Task 5 changes.

Use:

feat: implement ROOT-GNN graph construction

Before committing, inspect the final diff and ensure no unrelated files are included.