| 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) |
|
|
|
|
| def test_reordering_vector_objects_reorders_feature_rows(event, schema): |
| names, object_types, scales = schema |
| reordered = dict(event) |
| for name in ("jet_pt", "jet_eta", "jet_phi", "jet_btag"): |
| reordered[name] = event[name][::-1].copy() |
| first, _ = build_node_features(event, names, object_types, scales) |
| second, _ = build_node_features(reordered, names, object_types, scales) |
| torch.testing.assert_close(second[:2], first[[1, 0]]) |
| torch.testing.assert_close(second[2:], first[2:]) |
|
|
|
|
| def test_feature_scaling_changes_only_the_requested_columns(event, schema): |
| names, object_types, scales = schema |
| base, _ = build_node_features(event, names, object_types, scales) |
| changed_scales = list(scales) |
| changed_scales[0] *= 2 |
| scaled, _ = build_node_features(event, names, object_types, changed_scales) |
| torch.testing.assert_close(scaled[:, 0], base[:, 0] * 2) |
| torch.testing.assert_close(scaled[:, 1:], base[:, 1:]) |
|
|
