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