tests/integration/test_real_root_sample.py

"""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, build_object_feature_table
from gnn4colliders.graphs import build_tensor_graph

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/fixtures/testing/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}; run the public validation "
            "command to download it"
        )
    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_tensor_graph(
            event_features, torch.tensor([event_features.shape[0]])
        )

        assert event_features.dtype == torch.float32
        assert event_features.shape == (sum(lengths), 7)
        assert graph.node_features.shape[0] == event_features.shape[0]
        expected_edges = event_features.shape[0] * max(event_features.shape[0] - 1, 1)
        assert graph.edge_features.shape == (expected_edges, 3)
        assert torch.isfinite(event_features).all()
        assert torch.isfinite(graph.edge_features).all()


def test_vectorized_feature_table_matches_per_event_features(root_tree):
    arrays = root_tree.arrays(library="ak")
    per_event = [
        build_node_features(event, FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES)[0]
        for event in ak.Array(arrays)
    ]
    table, counts = build_object_feature_table(
        arrays, len(arrays), FEATURE_BRANCHES, OBJECT_TYPES, FEATURE_SCALES
    )
    assert torch.equal(table, torch.cat(per_event))
    assert counts.tolist() == [len(features) for features in per_event]