feat: implement shared collider feature construction

src/gnn4colliders/features/__init__.py

@@ -0,0 +1,5 @@
+"""Collider-domain feature transformations shared across model families."""
+
+from .objects import NODE_FEATURE_NAMES, build_node_features, build_object_features
+
+__all__ = ["NODE_FEATURE_NAMES", "build_node_features", "build_object_features"]

src/gnn4colliders/features/objects.py

@@ -0,0 +1,120 @@
+"""Architecture-independent collider-object feature construction."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping, Sequence
+from numbers import Number
+from typing import Any
+
+import numpy as np
+import torch
+
+NODE_FEATURE_NAMES = (
+    "pt",
+    "eta",
+    "phi",
+    "energy",
+    "btag",
+    "charge",
+    "node_type",
+)
+
+_CALCULATED_ENERGY = "CALC_E"
+_NODE_TYPE = "NODE_TYPE"
+
+
+def _as_feature_tensor(values: list[Any], dtype: torch.dtype) -> torch.Tensor:
+    return torch.as_tensor(np.asarray(values, dtype=np.float32), dtype=dtype)
+
+
+def build_node_features(
+    event: Mapping[str, Any],
+    feature_branches: Sequence[Any],
+    object_types: Sequence[str],
+    scales: Sequence[Number] | torch.Tensor,
+    *,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, list[int]]:
+    """Build dense node features from one collider event.
+
+    ``feature_branches`` follows the active legacy schema: each item describes
+    one output column and contains one branch name or constant per object
+    type. ``object_types`` contains ``"vector"`` or ``"single"`` for each
+    object type. Rows are concatenated in object-type order.
+
+    The returned tensor has columns in :data:`NODE_FEATURE_NAMES` order for
+    the active schema, and the second result records rows contributed by each
+    object type. Inputs are read but never modified.
+    """
+    if not feature_branches:
+        return torch.empty((0, 0), dtype=dtype), []
+    first_column = feature_branches[0]
+    if len(first_column) != len(object_types):
+        raise ValueError("one branch specification is required per object type")
+
+    lengths: list[int] = []
+    for branch, object_type in zip(first_column, object_types):
+        if object_type == "single":
+            lengths.append(1)
+        elif object_type == "vector":
+            lengths.append(len(event[branch]))
+        else:
+            raise ValueError(f"unknown object type: {object_type!r}")
+
+    columns: list[torch.Tensor] = []
+    for column_index, specification in enumerate(feature_branches):
+        if specification == _CALCULATED_ENERGY:
+            columns.append(columns[0] * torch.cosh(columns[1]))
+            continue
+        if specification == _NODE_TYPE:
+            values = [
+                object_type_index
+                for object_type_index, length in enumerate(lengths)
+                for _ in range(length)
+            ]
+            columns.append(torch.tensor(values, dtype=dtype))
+            continue
+        if len(specification) != len(object_types):
+            raise ValueError(
+                f"feature column {column_index} has the wrong number of branches"
+            )
+
+        values: list[Any] = []
+        for object_type_index, (length, branch, object_type) in enumerate(
+            zip(lengths, specification, object_types)
+        ):
+            if isinstance(branch, Number):
+                values.extend([branch] * length)
+            elif branch == _CALCULATED_ENERGY:
+                start = sum(lengths[:object_type_index])
+                stop = start + length
+                values.extend(
+                    (
+                        columns[0][start:stop] * torch.cosh(columns[1][start:stop])
+                    ).tolist()
+                )
+            elif object_type == "single":
+                values.append(event[branch])
+            else:
+                values.extend(event[branch])
+        columns.append(_as_feature_tensor(values, dtype))
+
+    features = torch.stack(columns, dim=1)
+    scale_tensor = torch.as_tensor(scales, dtype=dtype)
+    if scale_tensor.ndim != 1 or scale_tensor.numel() != features.shape[1]:
+        raise ValueError("scales must contain one value per feature column")
+    return features * scale_tensor, lengths
+
+
+def build_object_features(
+    event: Mapping[str, Any],
+    feature_branches: Sequence[Any],
+    object_types: Sequence[str],
+    scales: Sequence[Number] | torch.Tensor,
+    *,
+    dtype: torch.dtype = torch.float32,
+) -> tuple[torch.Tensor, list[int]]:
+    """Alias for :func:`build_node_features` using domain-neutral wording."""
+    return build_node_features(
+        event, feature_branches, object_types, scales, dtype=dtype
+    )

tests/parity/test_new_node_features.py

@@ -0,0 +1,45 @@
+import numpy as np
+import torch
+
+from gnn4colliders.features import build_node_features
+
+
+def test_new_builder_matches_legacy_builder(legacy_dataset_module_without_dgl):
+    event = {
+        "jet_pt": np.array([100.0, 50.0], dtype=np.float32),
+        "ele_pt": np.array([20.0], dtype=np.float32),
+        "mu_pt": np.array([30.0], dtype=np.float32),
+        "ph_pt": np.array([40.0], dtype=np.float32),
+        "MET_met": np.float32(25.0),
+        "jet_eta": np.array([1.0, -0.5], dtype=np.float32),
+        "ele_eta": np.array([0.25], dtype=np.float32),
+        "mu_eta": np.array([-0.75], dtype=np.float32),
+        "ph_eta": np.array([0.5], dtype=np.float32),
+        "jet_phi": np.array([3.0, -3.0], dtype=np.float32),
+        "ele_phi": np.array([0.2], dtype=np.float32),
+        "mu_phi": np.array([-0.4], dtype=np.float32),
+        "ph_phi": np.array([1.0], dtype=np.float32),
+        "MET_phi": np.float32(-1.2),
+        "jet_btag": np.array([0.8, 0.1], dtype=np.float32),
+        "ele_charge": np.array([-1.0], dtype=np.float32),
+        "mu_charge": np.array([1.0], dtype=np.float32),
+    }
+    names = [
+        ["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 = torch.tensor([0.1, 1, 1, 0.1, 1, 1, 1])
+    expected, expected_lengths = (
+        legacy_dataset_module_without_dgl.node_features_from_tree(
+            event, names, object_types, scales
+        )
+    )
+    actual, actual_lengths = build_node_features(event, names, object_types, scales)
+    assert actual_lengths == expected_lengths
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)

tests/unit/features/__init__.py

@@ -0,0 +1 @@
+"""Unit tests for shared collider feature construction."""

tests/unit/features/test_objects.py

@@ -0,0 +1,89 @@
+import numpy as np
+import pytest
+import torch
+
+from gnn4colliders.features import NODE_FEATURE_NAMES, build_node_features
+
+
+@pytest.fixture
+def schema():
+    return (
+        [
+            ["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",
+        ],
+        ["vector", "vector", "vector", "vector", "single"],
+        [0.1, 1, 1, 0.1, 1, 1, 1],
+    )
+
+
+@pytest.fixture
+def event():
+    return {
+        "jet_pt": np.array([100.0, 50.0], dtype=np.float32),
+        "ele_pt": np.array([20.0], dtype=np.float32),
+        "mu_pt": np.array([30.0], dtype=np.float32),
+        "ph_pt": np.array([40.0], dtype=np.float32),
+        "MET_met": np.float32(25.0),
+        "jet_eta": np.array([1.0, -0.5], dtype=np.float32),
+        "ele_eta": np.array([0.25], dtype=np.float32),
+        "mu_eta": np.array([-0.75], dtype=np.float32),
+        "ph_eta": np.array([0.5], dtype=np.float32),
+        "jet_phi": np.array([3.0, -3.0], dtype=np.float32),
+        "ele_phi": np.array([0.2], dtype=np.float32),
+        "mu_phi": np.array([-0.4], dtype=np.float32),
+        "ph_phi": np.array([1.0], dtype=np.float32),
+        "MET_phi": np.float32(-1.2),
+        "jet_btag": np.array([0.8, 0.1], dtype=np.float32),
+        "ele_charge": np.array([-1.0], dtype=np.float32),
+        "mu_charge": np.array([1.0], dtype=np.float32),
+    }
+
+
+def test_schema_and_values(event, schema):
+    names, object_types, scales = schema
+    features, lengths = build_node_features(event, names, object_types, scales)
+    assert NODE_FEATURE_NAMES == (
+        "pt",
+        "eta",
+        "phi",
+        "energy",
+        "btag",
+        "charge",
+        "node_type",
+    )
+    assert lengths == [2, 1, 1, 1, 1]
+    assert features.shape == (6, 7)
+    assert features.dtype == torch.float32
+    np.testing.assert_allclose(
+        features.numpy(),
+        [
+            [10.0, 1.0, 3.0, 15.431, 0.8, 0.0, 0.0],
+            [5.0, -0.5, -3.0, 5.638, 0.1, 0.0, 0.0],
+            [2.0, 0.25, 0.2, 2.063, 0.0, -1.0, 1.0],
+            [3.0, -0.75, -0.4, 3.884, 0.0, 1.0, 2.0],
+            [4.0, 0.5, 1.0, 4.511, 0.0, 0.0, 3.0],
+            [2.5, 0.0, -1.2, 2.500, 0.0, 0.0, 4.0],
+        ],
+        rtol=0,
+        atol=2e-3,
+    )
+
+
+def test_empty_vectors_and_input_immutability(event, schema):
+    names, object_types, scales = schema
+    event = dict(event)
+    event["jet_pt"] = np.array([], dtype=np.float32)
+    event["jet_eta"] = np.array([], dtype=np.float32)
+    event["jet_phi"] = np.array([], dtype=np.float32)
+    event["jet_btag"] = np.array([], dtype=np.float32)
+    before = event["ele_pt"].copy()
+    features, lengths = build_node_features(event, names, object_types, scales)
+    assert lengths == [0, 1, 1, 1, 1]
+    assert features.shape == (4, 7)
+    np.testing.assert_array_equal(event["ele_pt"], before)