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:
src/gnn4colliders/features/
Do not implement graph construction, datasets, training, or models yet.
Before making changes, read:
AGENTS.md
docs/architecture.md
docs/migration.md
Then inspect:
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:
DGL
graph topology
training code
model code
The same feature layer should be reusable by future model families such as:
root_gnn
root_transformer
Scope
Implement the active node/object feature behavior characterized in Task 3.
At minimum, support the active feature schema:
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:
src/gnn4colliders/features/
__init__.py
objects.py
If selection logic is genuinely needed for this task, a separate module such as:
selections.py
is acceptable.
Do not create many speculative modules.
API design
Design a small explicit API.
For example, something conceptually similar to:
def build_object_features(...):
...
or:
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:
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:
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:
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:
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:
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:
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:
docs/architecture.md
Do not duplicate large amounts of legacy documentation.
Validation
Run:
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:
all relevant tests pass
no legacy files modified
no graph implementation added
Inspect:
git status
git diff
before committing.
Completion criteria
Task 4 is complete when:
- New shared feature construction exists under
src/gnn4colliders/features/. - Active node-feature behavior matches the legacy characterization tests.
- Feature order is explicit.
- Dtypes are explicit.
- Derived energy behavior is covered.
- Node-type behavior is covered.
- Multiple-object ordering is covered where required.
- Unit tests pass.
- Parity tests pass.
- No DGL graph construction has been implemented.
- Legacy code remains untouched.
Completion report
Report:
- files created
- files modified
- public API added
- feature schema implemented
- parity results
- tolerances used
- intentional deviations, if any
- validation commands and results
- anything deferred
After validation succeeds, create a single commit:
feat: implement shared collider feature construction
Do not include unrelated changes.