| from pathlib import Path |
|
|
| import awkward as ak |
| import numpy as np |
| import pytest |
| import uproot |
|
|
| from gnn4colliders.data import EventSample, RootEventDataset |
| from gnn4colliders.data.root_io import branch_names_from_specs, read_tree |
| from gnn4colliders.features import build_object_features |
|
|
|
|
| def _write_fixture(path: Path) -> None: |
| with uproot.recreate(path) as root_file: |
| root_file["events"] = { |
| "eventNumber": np.array([10, 11, 12], dtype=np.int64), |
| "weight": np.array([1.5, -2.0, 0.25], dtype=np.float32), |
| "label_branch": np.array([4, 5, 6], dtype=np.int32), |
| "jet_pt": ak.Array([[10.0, 20.0], [], [30.0]]), |
| "global_x": np.array([0.5, 0.6, 0.7], dtype=np.float32), |
| } |
|
|
|
|
| def test_branch_selection_and_root_reading(tmp_path: Path): |
| path = tmp_path / "events.root" |
| _write_fixture(path) |
| assert branch_names_from_specs( |
| [["jet_pt"], "CALC_E", "NODE_TYPE"], |
| tracking_info=["eventNumber"], |
| global_features=["global_x"], |
| ) == ("jet_pt", "eventNumber", "global_x") |
| arrays = read_tree(path, "events", ["weight"]) |
| assert arrays.fields == ["weight"] |
| assert arrays["weight"].to_list() == [1.5, -2.0, 0.25] |
|
|
|
|
| def test_dataset_returns_architecture_neutral_samples(tmp_path: Path): |
| path = tmp_path / "events.root" |
| _write_fixture(path) |
| dataset = RootEventDataset( |
| path, |
| tree_name="events", |
| label="label_branch", |
| feature_branches=[["jet_pt"]], |
| global_features=["global_x"], |
| fold_var="eventNumber", |
| weight_var="weight", |
| ) |
|
|
| assert len(dataset) == 3 |
| assert dataset.branches == ( |
| "jet_pt", |
| "eventNumber", |
| "weight", |
| "global_x", |
| "label_branch", |
| ) |
| sample = dataset[1] |
| assert isinstance(sample, EventSample) |
| assert sample.event_index == 1 |
| assert sample.objects["jet_pt"].to_list() == [] |
| assert sample.label == 5 |
| assert sample.tracking.tolist() == [11.0, -2.0] |
| assert sample.global_features.tolist() == pytest.approx([0.6]) |
| assert dataset[-1].event_index == 2 |
|
|
|
|
| def test_multiple_files_keep_input_order_and_feed_feature_builder(tmp_path: Path): |
| first = tmp_path / "first.root" |
| second = tmp_path / "second.root" |
| _write_fixture(first) |
| _write_fixture(second) |
| dataset = RootEventDataset( |
| [first, second], |
| tree_name="events", |
| label=[0, 1], |
| feature_branches=[["jet_pt"], [0], [0], ["CALC_E"]], |
| fold_var="eventNumber", |
| ) |
|
|
| assert len(dataset) == 6 |
| assert dataset[3].event_index == 3 |
| assert dataset[3].label == 1 |
| features, lengths = build_object_features( |
| dataset[0].objects, |
| [["jet_pt"], [0], [0], ["CALC_E"]], |
| ["vector"], |
| [1, 1, 1, 1], |
| ) |
| assert features.shape == (2, 4) |
| assert lengths == [2] |
|
|
|
|
| def test_root_boundary_reports_missing_tree_and_branch(tmp_path: Path): |
| path = tmp_path / "events.root" |
| _write_fixture(path) |
| with pytest.raises(KeyError, match="tree"): |
| read_tree(path, "missing", ["weight"]) |
| with pytest.raises(KeyError, match="missing requested branches"): |
| read_tree(path, "events", ["missing_branch"]) |
|
|
|
|
| def test_root_boundary_reports_missing_file(tmp_path: Path): |
| with pytest.raises((OSError, FileNotFoundError)): |
| read_tree(tmp_path / "does-not-exist.root", "events", ["weight"]) |
|
|
|
|
| def test_sample_identity_is_stable_and_distinguishes_files(tmp_path: Path): |
| first = tmp_path / "first.root" |
| second = tmp_path / "second.root" |
| _write_fixture(first) |
| _write_fixture(second) |
| kwargs = dict( |
| tree_name="events", |
| label="label_branch", |
| feature_branches=[["jet_pt"]], |
| fold_var="eventNumber", |
| weight_var="weight", |
| ) |
| left = RootEventDataset(first, **kwargs) |
| repeated = RootEventDataset(first, **kwargs) |
| both = RootEventDataset([first, second], **kwargs) |
| assert left[0].metadata.sample_id == repeated[0].metadata.sample_id |
| assert left[0].metadata.sample_id != left[1].metadata.sample_id |
| assert both[0].metadata.sample_id != both[3].metadata.sample_id |
|
|
