rewriting codebase (#7)

916755e 3 months ago

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:

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

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