validation/extract_common.py

"""Shared portable extraction helpers; no legacy or DGL classes cross the boundary."""

from __future__ import annotations

from pathlib import Path
from typing import Any, Callable

import awkward as ak
import numpy as np
import torch
import uproot

from .artifacts import ValidationArtifact, file_sha256, save_artifact

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"]
SCALES = [1e-1, 1, 1, 1e-1, 1, 1, 1]
BRANCHES = (
    "jet_pt",
    "ele_pt",
    "mu_pt",
    "ph_pt",
    "MET_met",
    "jet_eta",
    "ele_eta",
    "mu_eta",
    "ph_eta",
    "jet_phi",
    "ele_phi",
    "mu_phi",
    "ph_phi",
    "MET_phi",
    "jet_btag",
    "ele_charge",
    "mu_charge",
    "Number",
    "weight",
)


def _row(value: Any) -> Any:
    """Convert an Awkward scalar while retaining a jagged vector."""
    converted = ak.to_numpy(value) if isinstance(value, ak.Array) else value
    array = np.asarray(converted)
    return array.item() if array.ndim == 0 else converted


def extract_root(
    path: str | Path,
    output: str | Path,
    *,
    build_features: Callable[[dict[str, Any]], torch.Tensor],
    sample_name: str = "testing/ttH_NLO_64.root",
    label: int = 0,
    filter_zero_pt: bool = False,
    source_sha256: str | None = None,
    label_branch: str | None = None,
) -> Path:
    """Extract one ROOT tree into the implementation-neutral artifact format."""
    path = Path(path)
    requested_branches = list(BRANCHES)
    if label_branch and label_branch not in requested_branches:
        requested_branches.append(label_branch)
    with uproot.open(path) as root_file:
        tree = root_file["output"]
        arrays = tree.arrays(requested_branches, library="ak")
        event_count = int(tree.num_entries)

    sample_ids: list[str] = []
    labels: list[int] = []
    folds: list[int] = []
    weights: list[float] = []
    node_parts: list[np.ndarray] = []
    node_offsets = [0]
    edge_src_parts: list[np.ndarray] = []
    edge_dst_parts: list[np.ndarray] = []
    edge_feature_parts: list[np.ndarray] = []
    edge_offsets = [0]

    for index in range(event_count):
        event = {name: _row(arrays[name][index]) for name in requested_branches}
        nodes = build_features(event)
        if filter_zero_pt:
            nodes = nodes[nodes[:, 0] != 0]
        node_count = int(nodes.shape[0])
        source = torch.arange(node_count).repeat_interleave(node_count)
        destination = torch.arange(node_count).repeat(node_count)
        if node_count > 1:
            keep = source != destination
            source, destination = source[keep], destination[keep]
        eta = nodes[:, 1]
        phi = nodes[:, 2]
        deta = eta[source] - eta[destination]
        dphi = phi[source] - phi[destination]
        dphi = torch.where(dphi > np.pi, dphi - 2 * np.pi, dphi)
        dphi = torch.where(dphi < -np.pi, dphi + 2 * np.pi, dphi)
        edge_features = torch.stack(
            (deta, dphi, torch.sqrt(deta.square() + dphi.square())), dim=1
        )

        sample_ids.append(f"{sample_name}:output:{index}")
        labels.append(int(event[label_branch]) if label_branch else label)
        folds.append(int(event["Number"]))
        weights.append(float(event["weight"]))
        node_parts.append(nodes.detach().cpu().numpy())
        edge_src_parts.append(source.numpy())
        edge_dst_parts.append(destination.numpy())
        edge_feature_parts.append(edge_features.detach().cpu().numpy())
        node_offsets.append(node_offsets[-1] + node_count)
        edge_offsets.append(edge_offsets[-1] + int(source.numel()))

    artifact = ValidationArtifact(
        sample_id=np.asarray(sample_ids),
        labels=np.asarray(labels),
        folds=np.asarray(folds),
        weights=np.asarray(weights, dtype=np.float32),
        globals=np.empty((event_count, 0), dtype=np.float32),
        node_features_flat=np.concatenate(node_parts, axis=0),
        node_offsets=np.asarray(node_offsets, dtype=np.int64),
        edge_src_flat=np.concatenate(edge_src_parts),
        edge_dst_flat=np.concatenate(edge_dst_parts),
        edge_features_flat=np.concatenate(edge_feature_parts, axis=0),
        edge_offsets=np.asarray(edge_offsets, dtype=np.int64),
        manifest={
            "source": "huggingface://datasets/HWresearch/Delphes",
            "repository_path": sample_name,
            "tree": "output",
            "events": event_count,
            "sha256": file_sha256(path),
            "source_sha256": source_sha256 or file_sha256(path),
            "input_path": path.name,
            "feature_schema": [
                "pt",
                "eta",
                "phi",
                "energy",
                "btag",
                "charge",
                "node_type",
            ],
        },
    )
    return save_artifact(artifact, output)