rewriting codebase (#7)

916755e 2 months ago

preview code

Raw History Blame

14.2 kB

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:

AGENTS.md
docs/architecture.md
docs/migration.md

Then inspect:

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:

ROOT file
    ->
Uproot / Awkward
    ->
event/sample representation
    ->
gnn4colliders.features
    ->
model-specific representation

For ROOT-GNN, the later model-specific representation is:

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:

src/gnn4colliders/data/
    __init__.py
    root_io.py
    sample.py
    dataset.py

Possible responsibilities:

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:

@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:

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:

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:

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:

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:

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:

tests/unit/data/

Suggested files:

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:

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:

gnn4colliders.data

unless absolutely unavoidable:

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:

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:

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:

uv run pytest tests/unit/data -v

Then run relevant feature integration tests:

uv run pytest tests/unit/features -v
uv run pytest tests/parity -v

Then run the full suite:

uv run pytest

Run lint/format validation:

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

Inspect:

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:

ROOT/Uproot reading exists under src/gnn4colliders/data/.
A clean architecture-neutral sample representation exists.
Deterministic event indexing works.
Required object collections are exposed correctly.
Labels match active legacy behavior.
Tracking info matches active legacy behavior.
Global-feature behavior is characterized and implemented.
Shared data code does not depend on DGL.
A sample can feed directly into the Task 4 feature layer.
Unit tests pass.
Relevant parity tests pass.
The full current suite passes.
Legacy code remains untouched.
No caching, batching, model, or training implementation has been added.

Completion report

Report:

files created
files modified
public APIs added
sample representation introduced
ROOT I/O behavior implemented
object collections supported
label behavior
tracking behavior
global-feature behavior
event-ordering/indexing rule
parity results
integration-test results
ambiguities discovered
intentional deviations from legacy behavior
validation commands and results
anything deferred

After validation succeeds, create one Git commit containing only Task 6 changes.

Use:

feat: implement ROOT data ingestion and sample representation

Before committing, inspect the final diff and ensure no unrelated files are included.