File size: 4,216 Bytes
916755e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | """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
)
|