| """Smoke tests against a small fixture derived from the real ttH ROOT sample.""" |
|
|
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
|
|
| import awkward as ak |
| import pytest |
| import torch |
| import uproot |
|
|
| from gnn4colliders.features import build_node_features |
| from gnn4colliders.graphs import build_dgl_graph |
|
|
| pytest.importorskip("dgl") |
|
|
| pytestmark = [pytest.mark.integration, pytest.mark.real_data] |
|
|
| FEATURE_BRANCHES = [ |
| ["jet_pt", "ele_pt", "mu_pt", "ph_pt", "MET_met"], |
| ["jet_eta", "ele_eta", "mu_eta", "ph_eta", 0], |
| ["jet_phi", "ele_phi", "mu_phi", "ph_phi", "MET_phi"], |
| "CALC_E", |
| ["jet_btag", 0, 0, 0, 0], |
| [0, "ele_charge", "mu_charge", 0, 0], |
| "NODE_TYPE", |
| ] |
| OBJECT_TYPES = ["vector", "vector", "vector", "vector", "single"] |
| FEATURE_SCALES = [0.1, 1, 1, 0.1, 1, 1, 1] |
|
|
|
|
| def _fixture_path() -> Path: |
| configured = os.environ.get("GNN4COLLIDERS_ROOT_FIXTURE") |
| if configured: |
| return Path(configured) |
| return Path("data/processed/ttH_NLO_64.root") |
|
|
|
|
| @pytest.fixture(scope="module") |
| def root_tree(): |
| path = _fixture_path() |
| if not path.exists(): |
| pytest.skip( |
| f"ROOT sample fixture is absent: {path}; download ttH_NLO.root and " |
| "create the reduced fixture first" |
| ) |
| root_file = uproot.open(path) |
| tree = root_file["output"] |
| try: |
| yield tree |
| finally: |
| root_file.close() |
|
|
|
|
| def test_real_root_sample_has_expected_fixture_shape(root_tree): |
| assert root_tree.num_entries == 64 |
| assert {"jet_pt", "MET_met", "weight", "Number"}.issubset(root_tree.keys()) |
|
|
|
|
| def test_real_root_events_build_features_and_graphs(root_tree): |
| arrays = root_tree.arrays(entry_start=0, entry_stop=8, library="ak") |
| for event in ak.Array(arrays): |
| event_features, lengths = build_node_features( |
| event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES |
| ) |
| graph = build_dgl_graph(event_features) |
|
|
| assert event_features.dtype == torch.float32 |
| assert event_features.shape == (sum(lengths), 7) |
| assert graph.number_of_nodes() == event_features.shape[0] |
| expected_edges = event_features.shape[0] * max(event_features.shape[0] - 1, 1) |
| assert graph.number_of_edges() == expected_edges |
| assert graph.edata["features"].shape == (expected_edges, 3) |
| assert torch.isfinite(event_features).all() |
| assert torch.isfinite(graph.edata["features"]).all() |
|
|
