| # Task 6: Implement ROOT/Awkward Data Ingestion and Sample Representation |
|
|
| Implement the new shared data-ingestion layer for GNN4Colliders. |
|
|
| This task should build on the completed feature-construction and graph-construction work from Tasks 4 and 5. |
|
|
| Before making changes, read: |
|
|
| ```text |
| AGENTS.md |
| docs/architecture.md |
| docs/migration.md |
| ``` |
|
|
| Then inspect: |
|
|
| ```text |
| src/gnn4colliders/features/ |
| src/gnn4colliders/graphs/ |
| tests/unit/ |
| tests/parity/ |
| legacy/root_gnn_dgl/ |
| ``` |
|
|
| Use the legacy implementation and existing parity tests as behavioral references. |
|
|
| Do not implement caching, pre-batching, training, or models in this task. |
|
|
| --- |
|
|
| ## Goal |
|
|
| Implement the shared data layer that turns ROOT data into architecture-neutral event/sample representations. |
|
|
| The intended flow is: |
|
|
| ```text |
| ROOT file |
| -> |
| Uproot / Awkward |
| -> |
| event/sample representation |
| -> |
| gnn4colliders.features |
| -> |
| model-specific representation |
| ``` |
|
|
| For ROOT-GNN, the later model-specific representation is: |
|
|
| ```text |
| sample |
| -> |
| gnn4colliders.features |
| -> |
| gnn4colliders.graphs |
| ``` |
|
|
| For a future ROOT-Transformer, the same sample/data layer should be reusable without DGL. |
|
|
| The new data layer should therefore not return DGL graphs directly. |
|
|
| --- |
|
|
| # 1. Scope |
|
|
| Implement only the active shared ROOT/Awkward ingestion behavior required by the standard ROOT-GNN workflow. |
|
|
| At minimum support: |
|
|
| * opening ROOT files with Uproot |
| * selecting the configured tree |
| * reading required branches |
| * event indexing |
| * converting event data into a clean in-memory sample structure |
| * labels |
| * tracking information |
| * global features |
| * object collections needed by Task 4 feature construction |
|
|
| Do not migrate every experimental dataset class. |
|
|
| Focus on the active path only. |
|
|
| --- |
|
|
| # 2. Package structure |
|
|
| Prefer a small structure such as: |
|
|
| ```text |
| src/gnn4colliders/data/ |
| __init__.py |
| root_io.py |
| sample.py |
| dataset.py |
| ``` |
|
|
| Possible responsibilities: |
|
|
| ```text |
| root_io.py |
| low-level ROOT/Uproot reading |
| |
| sample.py |
| typed event/sample representation |
| |
| dataset.py |
| dataset abstraction over one or more ROOT files |
| ``` |
|
|
| If fewer modules are clearer, use fewer. |
|
|
| Do not create speculative abstractions or a deep class hierarchy. |
|
|
| --- |
|
|
| # 3. Architecture-neutral sample representation |
|
|
| Introduce a small explicit representation for one event/sample. |
|
|
| Prefer a typed dataclass or similarly simple structure. |
|
|
| Conceptually something like: |
|
|
| ```python |
| @dataclass(frozen=True) |
| class EventSample: |
| objects: ... |
| label: ... |
| tracking: ... |
| global_features: ... |
| ``` |
|
|
| The exact fields should reflect the actual active legacy behavior. |
|
|
| The sample representation should: |
|
|
| * be independent of DGL |
| * be independent of ROOT-GNN model classes |
| * be usable by future model families |
| * have explicit types |
| * avoid hidden global state |
| * avoid mutation where practical |
|
|
| Do not force raw ROOT branch names to leak throughout downstream code if a cleaner normalized sample structure is practical. |
|
|
| At the same time, do not over-generalize beyond the active workflow. |
|
|
| --- |
|
|
| # 4. ROOT I/O |
|
|
| Implement low-level ROOT reading with Uproot. |
|
|
| Responsibilities may include: |
|
|
| * opening a ROOT file |
| * selecting a tree by name |
| * validating requested branches |
| * reading selected branches |
| * returning Awkward arrays or an equivalent architecture-neutral structure |
|
|
| Prefer explicit functions. |
|
|
| For example, conceptually: |
|
|
| ```python |
| def read_tree( |
| path: Path, |
| tree_name: str, |
| branches: Sequence[str], |
| ) -> ak.Array: |
| ... |
| ``` |
|
|
| The exact API may differ. |
|
|
| Do not use PyROOT unless the active path requires behavior that Uproot cannot provide. |
|
|
| Do not add DGL to this module. |
|
|
| --- |
|
|
| # 5. Branch selection |
|
|
| Determine the branch set needed by the active standard configs. |
|
|
| Do not read every branch in the tree by default if only a subset is required. |
|
|
| The data layer should make required branch names explicit. |
|
|
| If branch sets depend on object type or task, represent that cleanly without reproducing the old dynamic-import configuration system. |
|
|
| Do not implement the full future Hydra configuration layer yet. |
|
|
| A plain typed argument/config object is acceptable. |
|
|
| --- |
|
|
| # 6. Object collections |
|
|
| Preserve the object-collection semantics needed by Task 4 feature construction. |
|
|
| For each active object collection, characterize and preserve: |
|
|
| * object ordering |
| * branch mapping |
| * jagged structure |
| * missing/empty collection behavior |
| * scalar vs vector branch behavior where relevant |
|
|
| Do not duplicate feature construction here. |
|
|
| This layer should expose raw or normalized object data. |
|
|
| Task 4 remains responsible for constructing: |
|
|
| ```text |
| pt |
| eta |
| phi |
| energy |
| btag |
| charge |
| node_type |
| ``` |
|
|
| Do not calculate graph edges in this layer. |
|
|
| --- |
|
|
| # 7. Labels |
|
|
| Characterize and implement active label extraction. |
|
|
| Verify: |
|
|
| * source branch or source rule |
| * scalar vs vector representation |
| * dtype |
| * shape |
| * binary-task behavior |
| * multiclass behavior where active configs require it |
|
|
| Do not expand into every historical multi-label/experimental path unless required by active configs or existing parity tests. |
|
|
| The sample representation should expose labels cleanly. |
|
|
| --- |
|
|
| # 8. Tracking information |
|
|
| Implement tracking extraction matching the active legacy behavior. |
|
|
| At minimum preserve the characterized semantics: |
|
|
| ```text |
| tracking column 0 -> fold identifier |
| tracking column 1 -> event weight |
| ``` |
|
|
| Verify actual source branches/rules in the legacy implementation. |
|
|
| Preserve: |
|
|
| * dtype |
| * shape |
| * ordering |
|
|
| Do not implement fold filtering yet. |
|
|
| Do not mutate tracking arrays in place. |
|
|
| If the legacy code mutates tracking information as part of dataset processing, reproduce the resulting external behavior without preserving unnecessary mutation internally. |
|
|
| --- |
|
|
| # 9. Global features |
|
|
| Implement active global-feature extraction. |
|
|
| Characterize: |
|
|
| * whether globals are present |
| * source branches |
| * shape |
| * dtype |
| * empty/default representation when not used |
|
|
| Do not invent model-specific global semantics. |
|
|
| The data layer should expose globals as generic sample data. |
|
|
| --- |
|
|
| # 10. Dataset abstraction |
|
|
| Implement a simple dataset abstraction over ROOT data. |
|
|
| Conceptually: |
|
|
| ```python |
| dataset = RootEventDataset(...) |
| sample = dataset[i] |
| ``` |
|
|
| The exact class/API may differ. |
|
|
| The dataset should preferably support: |
|
|
| * `__len__` |
| * `__getitem__` |
| * deterministic event indexing |
|
|
| Do not add DGL-specific behavior. |
|
|
| Do not make it inherit from a DGL dataset class. |
|
|
| It may inherit from `torch.utils.data.Dataset` only if that provides clear value and does not force unwanted coupling. |
|
|
| If plain Python indexing is sufficient, keep it simpler. |
|
|
| --- |
|
|
| # 11. One file vs multiple files |
|
|
| Support the smallest practical file abstraction required by the active workflow. |
|
|
| If the active standard path requires multiple ROOT files, support them explicitly. |
|
|
| Preserve deterministic event ordering across files. |
|
|
| Document the ordering rule. |
|
|
| Do not implement distributed file sharding yet. |
|
|
| Do not implement cache chunking yet. |
|
|
| --- |
|
|
| # 12. Event indexing |
|
|
| Characterize and preserve event indexing semantics. |
|
|
| Determine: |
|
|
| * whether indexing is global across files |
| * how event offsets are computed |
| * whether ordering follows input file order |
| * whether any implicit sorting occurs |
|
|
| Prefer a transparent deterministic rule. |
|
|
| Add tests for file boundaries if multiple files are supported. |
|
|
| --- |
|
|
| # 13. Error behavior |
|
|
| Characterize active behavior for: |
|
|
| * missing file |
| * missing tree |
| * missing required branch |
| * malformed branch shape |
|
|
| Do not implement an elaborate validation framework. |
|
|
| Raise clear Python exceptions. |
|
|
| Where legacy behavior is poor or accidental, preserve compatibility only if downstream behavior depends on it. |
|
|
| Document intentional improvements in error clarity. |
|
|
| --- |
|
|
| # 14. Small ROOT fixture |
|
|
| Create a tiny deterministic ROOT fixture if one does not already exist and if actual ROOT I/O needs direct testing. |
|
|
| Store it under: |
|
|
| ```text |
| tests/fixtures/ |
| ``` |
|
|
| Keep it very small. |
|
|
| It should include enough events to exercise: |
|
|
| * at least one object collection |
| * multiple objects in an event |
| * different event sizes |
| * labels |
| * fold information |
| * non-unit event weights |
| * global features if active |
| * empty collection behavior if supported and useful |
|
|
| Prefer generating the fixture in test setup if that is clearer and avoids committing binary data. |
|
|
| If a committed ROOT file is used, keep it tiny. |
|
|
| --- |
|
|
| # 15. Tests |
|
|
| Add unit tests under: |
|
|
| ```text |
| tests/unit/data/ |
| ``` |
|
|
| Suggested files: |
|
|
| ```text |
| test_root_io.py |
| test_sample.py |
| test_dataset.py |
| ``` |
|
|
| Use fewer files if clearer. |
|
|
| Unit tests should cover the new implementation directly. |
|
|
| --- |
|
|
| # 16. Parity tests |
|
|
| Extend parity tests to compare the new data layer against the active legacy behavior. |
|
|
| At minimum compare, where practical: |
|
|
| * event count |
| * per-event object arrays |
| * object ordering |
| * labels |
| * tracking |
| * global features |
| * dtypes |
| * shapes |
|
|
| Do not compare DGL graphs in this task unless an existing integration test naturally verifies downstream compatibility. |
|
|
| Graph parity belongs primarily to Task 5. |
|
|
| --- |
|
|
| # 17. Integration with Task 4 features |
|
|
| Add a small integration test showing that one sample from the new dataset can be passed into the Task 4 feature builder. |
|
|
| Conceptually: |
|
|
| ```text |
| ROOT fixture |
| -> |
| RootEventDataset |
| -> |
| EventSample |
| -> |
| build_object_features(...) |
| ``` |
|
|
| Verify the result matches the expected/legacy feature construction. |
|
|
| Do not continue into DGL graph construction unless a tiny integration check is already natural. |
|
|
| The main goal is to prove the shared data-to-feature boundary works. |
|
|
| --- |
|
|
| # 18. No DGL dependency in shared data layer |
|
|
| The following should not be imported by: |
|
|
| ```text |
| gnn4colliders.data |
| ``` |
|
|
| unless absolutely unavoidable: |
|
|
| ```text |
| dgl |
| gnn4colliders.graphs |
| gnn4colliders.models.root_gnn |
| ``` |
|
|
| The data layer must remain reusable by future architectures. |
|
|
| --- |
|
|
| # 19. No caching yet |
|
|
| Do not implement: |
|
|
| * DGL `.bin` cache files |
| * chunked graph caches |
| * lazy loading |
| * ring buffers |
| * pre-batching |
| * padding |
| * cache hashing |
| * cache invalidation |
|
|
| Those belong to Task 7. |
|
|
| This task should read ROOT data and expose samples directly. |
|
|
| --- |
|
|
| # 20. No fold selection yet |
|
|
| Although tracking should expose fold IDs, do not implement train/validation/test fold filtering in this task. |
|
|
| That belongs with dataset orchestration/batching in Task 7. |
|
|
| --- |
|
|
| # 21. No DataLoader/pre-batching yet |
|
|
| Do not implement: |
|
|
| * GraphDataLoader |
| * PyTorch DataLoader tuning |
| * pre-batching |
| * padding |
| * multiprocessing workers |
| * worker seeding |
| * distributed samplers |
|
|
| Those belong to later tasks. |
|
|
| Keep this task focused on deterministic sample access. |
|
|
| --- |
|
|
| # 22. No model or training changes |
|
|
| Do not modify: |
|
|
| ```text |
| src/gnn4colliders/models/ |
| src/gnn4colliders/training/ |
| src/gnn4colliders/inference/ |
| ``` |
|
|
| Do not migrate model code. |
|
|
| Do not implement losses, metrics, optimizers, checkpoints, CLI workflows, or distributed execution. |
|
|
| --- |
|
|
| # 23. No configuration-system implementation |
|
|
| Do not implement Hydra composition yet. |
|
|
| If the new data APIs require structured settings, use small typed Python objects or explicit function arguments. |
|
|
| The future config layer should wrap these APIs rather than define their internal behavior. |
|
|
| --- |
|
|
| # 24. Code quality |
|
|
| Prefer: |
|
|
| * explicit typed APIs |
| * `pathlib.Path` |
| * small pure functions |
| * immutable sample objects where practical |
| * deterministic ordering |
| * architecture-neutral data structures |
| * clear validation errors |
|
|
| Avoid: |
|
|
| * mutable global configuration |
| * repository-relative `sys.path` hacks |
| * dynamic imports |
| * giant dataset classes |
| * hidden file discovery |
| * implicit environment-specific paths |
| * model-specific behavior in shared data code |
|
|
| --- |
|
|
| # 25. Documentation |
|
|
| Update `docs/migration.md` once the new data-ingestion layer reaches parity. |
|
|
| Update `docs/architecture.md` only if needed to clarify the new boundary: |
|
|
| ```text |
| gnn4colliders.data |
| ROOT/Awkward I/O and architecture-neutral samples |
| |
| gnn4colliders.features |
| collider feature construction |
| |
| gnn4colliders.graphs |
| graph-specific representation |
| ``` |
|
|
| Keep documentation changes concise. |
|
|
| --- |
|
|
| # Validation |
|
|
| Run: |
|
|
| ```bash |
| uv run pytest tests/unit/data -v |
| ``` |
|
|
| Then run relevant feature integration tests: |
|
|
| ```bash |
| uv run pytest tests/unit/features -v |
| uv run pytest tests/parity -v |
| ``` |
|
|
| Then run the full suite: |
|
|
| ```bash |
| uv run pytest |
| ``` |
|
|
| Run lint/format validation: |
|
|
| ```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 DGL-specific data implementation was added |
| * no cache implementation was added |
| * no batching implementation was added |
| * no training/model implementation was added |
| * fixtures are small |
| * no generated ROOT or cache artifacts are accidentally tracked |
|
|
| --- |
|
|
| # Completion criteria |
|
|
| Task 6 is complete when: |
|
|
| 1. ROOT/Uproot reading exists under `src/gnn4colliders/data/`. |
| 2. A clean architecture-neutral sample representation exists. |
| 3. Deterministic event indexing works. |
| 4. Required object collections are exposed correctly. |
| 5. Labels match active legacy behavior. |
| 6. Tracking info matches active legacy behavior. |
| 7. Global-feature behavior is characterized and implemented. |
| 8. Shared data code does not depend on DGL. |
| 9. A sample can feed directly into the Task 4 feature layer. |
| 10. Unit tests pass. |
| 11. Relevant parity tests pass. |
| 12. The full current suite passes. |
| 13. Legacy code remains untouched. |
| 14. No caching, batching, model, or training implementation has been added. |
|
|
| --- |
|
|
| # Completion report |
|
|
| Report: |
|
|
| 1. files created |
| 2. files modified |
| 3. public APIs added |
| 4. sample representation introduced |
| 5. ROOT I/O behavior implemented |
| 6. object collections supported |
| 7. label behavior |
| 8. tracking behavior |
| 9. global-feature behavior |
| 10. event-ordering/indexing rule |
| 11. parity results |
| 12. integration-test results |
| 13. ambiguities discovered |
| 14. intentional deviations from legacy behavior |
| 15. validation commands and results |
| 16. anything deferred |
|
|
| After validation succeeds, create one Git commit containing only Task 6 changes. |
|
|
| Use: |
|
|
| ```text |
| feat: implement ROOT data ingestion and sample representation |
| ``` |
|
|
| Before committing, inspect the final diff and ensure no unrelated files are included. |
|
|
