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