rewriting codebase (#7)

916755e 2 months ago

8.26 kB

# Task 4: Implement Shared Physics Feature Construction

Implement the first production component of the GNN4Colliders rewrite: the shared collider-object feature construction layer.

This task should implement new code under:

```text
src/gnn4colliders/features/
```

Do not implement graph construction, datasets, training, or models yet.

Before making changes, read:

```text
AGENTS.md
docs/architecture.md
docs/migration.md
```

Then inspect:

```text
tests/parity/
legacy/root_gnn_dgl/
```

Use the Task 3 characterization tests as the behavioral specification.

---

## Goal

Implement architecture-independent collider-object feature construction that reproduces the active legacy behavior where compatibility is required.

The output of this layer should be plain feature arrays/tensors.

It should not depend on:

```text
DGL
graph topology
training code
model code
```

The same feature layer should be reusable by future model families such as:

```text
root_gnn
root_transformer
```

---

# Scope

Implement the active node/object feature behavior characterized in Task 3.

At minimum, support the active feature schema:

```text
pt
eta
phi
energy
btag
charge
node_type
```

Verify all details against the parity tests and legacy implementation.

Implement only behavior needed by the active standard workflow.

Do not migrate experimental or unused feature paths unless required by tests.

---

# Package structure

Prefer a small structure such as:

```text
src/gnn4colliders/features/
    __init__.py
    objects.py
```

If selection logic is genuinely needed for this task, a separate module such as:

```text
selections.py
```

is acceptable.

Do not create many speculative modules.

---

# API design

Design a small explicit API.

For example, something conceptually similar to:

```python
def build_object_features(...):
    ...
```

or:

```python
def build_node_features(...):
    ...
```

The exact function name and arguments should be chosen based on the actual legacy inputs and characterization tests.

Prefer:

* explicit inputs
* explicit outputs
* typed signatures
* no hidden global configuration
* no mutation of caller-owned data
* no dependency on repository-relative paths

Do not expose legacy dynamic-import configuration behavior through this API.

---

# Feature behavior

Implement and test:

## pt

Preserve active legacy behavior.

## eta

Preserve active legacy behavior.

## phi

Preserve active legacy behavior.

Do not normalize or wrap phi unless the legacy node-feature path actually does so.

## energy

Support the active derived-energy behavior.

If Task 3 confirmed:

```text
CALC_E = pt * cosh(eta)
```

implement that behavior directly and test it.

Do not add alternative physics definitions unless required.

## btag

Preserve active legacy constants/branch behavior.

## charge

Preserve active legacy constants/branch behavior.

## node_type
Preserve active legacy node-type coding and object ordering.
Do not redesign node-type identifiers in this task.
---
# Feature ordering
The output column order must match the characterized active behavior.
Expected order is likely:
```text
pt
eta
phi
energy
btag
charge
node_type
```
but verify against tests.
Add a named constant or schema representation if that makes the ordering explicit and reduces magic indices.
For example:
```python
NODE_FEATURE_NAMES = (
    "pt",
    "eta",
    "phi",
    "energy",
    "btag",
    "charge",
    "node_type",
)
```
Keep it simple.

---

# Feature scaling

If Task 3 characterized legacy feature scaling, implement that behavior.

Scaling should be explicit.

Avoid hidden mutable global scale state.

Prefer passing scale information into the feature builder or representing it through a small typed configuration object.

Do not introduce Hydra integration yet.

---

# Multiple object types

If the active legacy path combines multiple collider-object collections:

* preserve concatenation order
* preserve node-type identifiers
* preserve feature ordering
* preserve dtype

Characterize any additional behavior with unit tests.

Do not generalize prematurely to every possible collider object type.

Implement only what the active workflow requires.

---

# Data representation

Prefer the natural representation used by the implementation and tests.

Awkward Arrays may be accepted at the boundary where useful.

The feature layer should produce a predictable dense representation suitable for later graph or sequence construction.

Do not return a DGL graph.

Do not make DGL a dependency of this module.

---

# Dtype behavior

Preserve the characterized dtype where compatibility matters.

Avoid relying on implicit dtype conversion.

Use explicit dtype handling where appropriate.

Add tests that assert dtype.

---

# Tests

Add new unit tests under:

```text
tests/unit/features/
```

or an equivalent clear location.

Unit tests should test the new implementation directly.

Then adapt or extend parity testing so that the new implementation is compared against the legacy behavior.

The important distinction is:

```text
tests/unit/
    tests the new implementation in isolation
tests/parity/
    compares new behavior with legacy/reference behavior
```

Do not remove the legacy characterization tests.

---

# Parity

For deterministic feature construction, parity should be as strict as practical.

Verify:

* shape
* dtype
* column ordering
* numeric values
* object ordering
* node-type values
* scaling
* derived energy

Use explicit tolerances only for floating-point calculations.

Discrete values should use exact equality.

---

# No graph implementation

Do not implement:

```text
fully connected graphs
DGLGraph construction
edge indices
deta
dphi
dR
graph padding
```

Those belong to the next task.

Even if graph construction is simple, keep this task scoped to shared feature construction.

---

# No ROOT I/O rewrite

Do not implement the new ROOT reader yet unless absolutely necessary for a small interface boundary.

Tests may continue to construct small in-memory arrays/events.

The new feature implementation should not need to open ROOT files itself.

ROOT/Awkward I/O will be implemented separately.

---

# No configuration system

Do not implement Hydra schemas or YAML loading in this task.

The feature API should be usable directly from Python.

Configuration will wrap it later.

---

# Code quality

Use:

* type annotations
* docstrings for public functions
* small pure functions where practical
* explicit constants for schema/order
* no hidden global state

Avoid:

* giant utility modules
* class hierarchies without a demonstrated need
* dynamic imports
* mutation of input arrays
* model-specific assumptions

---

# Documentation

Update `docs/migration.md` to mark feature construction as implemented only after parity passes.

If useful, add a short description of the feature schema to:

```text
docs/architecture.md
```

Do not duplicate large amounts of legacy documentation.

---

# Validation

Run:

```bash
uv run pytest tests/unit
uv run pytest tests/parity
uv run pytest
uv run ruff check src tests
uv run ruff format --check src tests
```

Expected outcome:

```text
all relevant tests pass
no legacy files modified
no graph implementation added
```

Inspect:

```bash
git status
git diff
```

before committing.

---

# Completion criteria

Task 4 is complete when:

1. New shared feature construction exists under `src/gnn4colliders/features/`.
2. Active node-feature behavior matches the legacy characterization tests.
3. Feature order is explicit.
4. Dtypes are explicit.
5. Derived energy behavior is covered.
6. Node-type behavior is covered.
7. Multiple-object ordering is covered where required.
8. Unit tests pass.
9. Parity tests pass.
10. No DGL graph construction has been implemented.
11. Legacy code remains untouched.

---

# Completion report

Report:

1. files created
2. files modified
3. public API added
4. feature schema implemented
5. parity results
6. tolerances used
7. intentional deviations, if any
8. validation commands and results
9. anything deferred

After validation succeeds, create a single commit:

```text
feat: implement shared collider feature construction
```

Do not include unrelated changes.

	# Task 4: Implement Shared Physics Feature Construction

	Implement the first production component of the GNN4Colliders rewrite: the shared collider-object feature construction layer.

	This task should implement new code under:

	```text
	src/gnn4colliders/features/
	```

	Do not implement graph construction, datasets, training, or models yet.

	Before making changes, read:

	```text
	AGENTS.md
	docs/architecture.md
	docs/migration.md
	```

	Then inspect:

	```text
	tests/parity/
	legacy/root_gnn_dgl/
	```

	Use the Task 3 characterization tests as the behavioral specification.

	---

	## Goal

	Implement architecture-independent collider-object feature construction that reproduces the active legacy behavior where compatibility is required.

	The output of this layer should be plain feature arrays/tensors.

	It should not depend on:

	```text
	DGL
	graph topology
	training code
	model code
	```

	The same feature layer should be reusable by future model families such as:

	```text
	root_gnn
	root_transformer
	```

	---

	# Scope

	Implement the active node/object feature behavior characterized in Task 3.

	At minimum, support the active feature schema:

	```text
	pt
	eta
	phi
	energy
	btag
	charge
	node_type
	```

	Verify all details against the parity tests and legacy implementation.

	Implement only behavior needed by the active standard workflow.

	Do not migrate experimental or unused feature paths unless required by tests.

	---

	# Package structure

	Prefer a small structure such as:

	```text
	src/gnn4colliders/features/
	__init__.py
	objects.py
	```

	If selection logic is genuinely needed for this task, a separate module such as:

	```text
	selections.py
	```

	is acceptable.

	Do not create many speculative modules.

	---

	# API design

	Design a small explicit API.

	For example, something conceptually similar to:

	```python
	def build_object_features(...):
	...
	```

	or:

	```python
	def build_node_features(...):
	...
	```

	The exact function name and arguments should be chosen based on the actual legacy inputs and characterization tests.

	Prefer:

	* explicit inputs
	* explicit outputs
	* typed signatures
	* no hidden global configuration
	* no mutation of caller-owned data
	* no dependency on repository-relative paths

	Do not expose legacy dynamic-import configuration behavior through this API.

	---

	# Feature behavior

	Implement and test:

	## pt

	Preserve active legacy behavior.

	## eta

	Preserve active legacy behavior.

	## phi

	Preserve active legacy behavior.

	Do not normalize or wrap phi unless the legacy node-feature path actually does so.

	## energy

	Support the active derived-energy behavior.

	If Task 3 confirmed:

	```text
	CALC_E = pt * cosh(eta)
	```

	implement that behavior directly and test it.

	Do not add alternative physics definitions unless required.

	## btag

	Preserve active legacy constants/branch behavior.

	## charge

	Preserve active legacy constants/branch behavior.

	## node_type

	Preserve active legacy node-type coding and object ordering.

	Do not redesign node-type identifiers in this task.

	---

	# Feature ordering

	The output column order must match the characterized active behavior.

	Expected order is likely:

	```text
	pt
	eta
	phi
	energy
	btag
	charge
	node_type
	```

	but verify against tests.

	Add a named constant or schema representation if that makes the ordering explicit and reduces magic indices.

	For example:

	```python
	NODE_FEATURE_NAMES = (
	"pt",
	"eta",
	"phi",
	"energy",
	"btag",
	"charge",
	"node_type",
	)
	```

	Keep it simple.

	---

	# Feature scaling

	If Task 3 characterized legacy feature scaling, implement that behavior.

	Scaling should be explicit.

	Avoid hidden mutable global scale state.

	Prefer passing scale information into the feature builder or representing it through a small typed configuration object.

	Do not introduce Hydra integration yet.

	---

	# Multiple object types

	If the active legacy path combines multiple collider-object collections:

	* preserve concatenation order
	* preserve node-type identifiers
	* preserve feature ordering
	* preserve dtype

	Characterize any additional behavior with unit tests.

	Do not generalize prematurely to every possible collider object type.

	Implement only what the active workflow requires.

	---

	# Data representation

	Prefer the natural representation used by the implementation and tests.

	Awkward Arrays may be accepted at the boundary where useful.

	The feature layer should produce a predictable dense representation suitable for later graph or sequence construction.

	Do not return a DGL graph.

	Do not make DGL a dependency of this module.

	---

	# Dtype behavior

	Preserve the characterized dtype where compatibility matters.

	Avoid relying on implicit dtype conversion.

	Use explicit dtype handling where appropriate.

	Add tests that assert dtype.

	---

	# Tests

	Add new unit tests under:

	```text
	tests/unit/features/
	```

	or an equivalent clear location.

	Unit tests should test the new implementation directly.

	Then adapt or extend parity testing so that the new implementation is compared against the legacy behavior.

	The important distinction is:

	```text
	tests/unit/
	tests the new implementation in isolation

	tests/parity/
	compares new behavior with legacy/reference behavior
	```

	Do not remove the legacy characterization tests.

	---

	# Parity

	For deterministic feature construction, parity should be as strict as practical.

	Verify:

	* shape
	* dtype
	* column ordering
	* numeric values
	* object ordering
	* node-type values
	* scaling
	* derived energy

	Use explicit tolerances only for floating-point calculations.

	Discrete values should use exact equality.

	---

	# No graph implementation

	Do not implement:

	```text
	fully connected graphs
	DGLGraph construction
	edge indices
	deta
	dphi
	dR
	graph padding
	```

	Those belong to the next task.

	Even if graph construction is simple, keep this task scoped to shared feature construction.

	---

	# No ROOT I/O rewrite

	Do not implement the new ROOT reader yet unless absolutely necessary for a small interface boundary.

	Tests may continue to construct small in-memory arrays/events.

	The new feature implementation should not need to open ROOT files itself.

	ROOT/Awkward I/O will be implemented separately.

	---

	# No configuration system

	Do not implement Hydra schemas or YAML loading in this task.

	The feature API should be usable directly from Python.

	Configuration will wrap it later.

	---

	# Code quality

	Use:

	* type annotations
	* docstrings for public functions
	* small pure functions where practical
	* explicit constants for schema/order
	* no hidden global state

	Avoid:

	* giant utility modules
	* class hierarchies without a demonstrated need
	* dynamic imports
	* mutation of input arrays
	* model-specific assumptions

	---

	# Documentation

	Update `docs/migration.md` to mark feature construction as implemented only after parity passes.

	If useful, add a short description of the feature schema to:

	```text
	docs/architecture.md
	```

	Do not duplicate large amounts of legacy documentation.

	---

	# Validation

	Run:

	```bash
	uv run pytest tests/unit
	uv run pytest tests/parity
	uv run pytest

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

	Expected outcome:

	```text
	all relevant tests pass
	no legacy files modified
	no graph implementation added
	```

	Inspect:

	```bash
	git status
	git diff
	```

	before committing.

	---

	# Completion criteria

	Task 4 is complete when:

	1. New shared feature construction exists under `src/gnn4colliders/features/`.
	2. Active node-feature behavior matches the legacy characterization tests.
	3. Feature order is explicit.
	4. Dtypes are explicit.
	5. Derived energy behavior is covered.
	6. Node-type behavior is covered.
	7. Multiple-object ordering is covered where required.
	8. Unit tests pass.
	9. Parity tests pass.
	10. No DGL graph construction has been implemented.
	11. Legacy code remains untouched.

	---

	# Completion report

	Report:

	1. files created
	2. files modified
	3. public API added
	4. feature schema implemented
	5. parity results
	6. tolerances used
	7. intentional deviations, if any
	8. validation commands and results
	9. anything deferred

	After validation succeeds, create a single commit:

	```text
	feat: implement shared collider feature construction
	```

	Do not include unrelated changes.