| """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 _grouped_energy( |
| pt: torch.Tensor, eta: torch.Tensor, sizes: np.ndarray |
| ) -> torch.Tensor: |
| """``pt * cosh(eta)`` evaluated one group (event) at a time. |
| |
| PyTorch's CPU ``cosh`` can differ by 1-2 ULP between SIMD lanes and the |
| scalar tail, so results depend on an element's position in the tensor. |
| The per-event builder evaluates each event separately; grouping the same |
| way keeps both builders bit-identical. This position (and CPU vector |
| width) dependence is a known legacy artifact. |
| """ |
| groups = [int(size) for size in sizes if size] |
| if not groups: |
| return pt.new_empty(0) |
| return pt * torch.cat([torch.cosh(part) for part in torch.split(eta, groups)]) |
|
|
|
|
| def build_object_feature_table( |
| arrays: Mapping[str, Any], |
| num_events: int, |
| feature_branches: Sequence[Any], |
| object_types: Sequence[str], |
| scales: Sequence[Number] | torch.Tensor, |
| *, |
| dtype: torch.dtype = torch.float32, |
| ) -> tuple[torch.Tensor, torch.Tensor]: |
| """Vectorized :func:`build_node_features` over many events at once. |
| |
| ``arrays`` maps branch names to per-event values for ``num_events`` |
| events: Awkward/NumPy jagged arrays for ``vector`` object types and flat |
| arrays for ``single`` ones. Returns the object rows of all events in event |
| order (identical to concatenating per-event results) and the number of |
| objects per event. |
| """ |
| import awkward as ak |
|
|
| if not feature_branches: |
| return torch.empty((0, 0), dtype=dtype), torch.zeros( |
| num_events, dtype=torch.long |
| ) |
| first_column = feature_branches[0] |
| if len(first_column) != len(object_types): |
| raise ValueError("one branch specification is required per object type") |
|
|
| lengths: list[np.ndarray] = [] |
| for branch, object_type in zip(first_column, object_types): |
| if object_type == "single": |
| lengths.append(np.ones(num_events, dtype=np.int64)) |
| elif object_type == "vector": |
| lengths.append(np.asarray(ak.num(arrays[branch], axis=1), dtype=np.int64)) |
| else: |
| raise ValueError(f"unknown object type: {object_type!r}") |
| per_type = np.stack(lengths, axis=1) if lengths else np.zeros((num_events, 0)) |
| counts = per_type.sum(axis=1) |
| event_start = np.cumsum(counts) - counts |
| total = int(counts.sum()) |
|
|
| |
| positions: list[np.ndarray] = [] |
| for type_index in range(len(object_types)): |
| type_lengths = per_type[:, type_index] |
| base = event_start + per_type[:, :type_index].sum(axis=1) |
| type_total = int(type_lengths.sum()) |
| local = np.arange(type_total, dtype=np.int64) - np.repeat( |
| np.cumsum(type_lengths) - type_lengths, type_lengths |
| ) |
| positions.append(np.repeat(base, type_lengths) + local) |
|
|
| def flat_values(branch: Any, type_index: int) -> np.ndarray: |
| size = positions[type_index].size |
| if isinstance(branch, Number): |
| return np.full(size, branch, dtype=np.float32) |
| values = arrays[branch] |
| if object_types[type_index] == "vector": |
| values = ak.flatten(values, axis=1) |
| return np.asarray( |
| ak.to_numpy(values) if isinstance(values, ak.Array) else values |
| ).astype(np.float32) |
|
|
| columns: list[torch.Tensor] = [] |
| for column_index, specification in enumerate(feature_branches): |
| if specification == _CALCULATED_ENERGY: |
| columns.append(_grouped_energy(columns[0], columns[1], counts)) |
| continue |
| if specification == _NODE_TYPE: |
| node_type = np.empty(total, dtype=np.float32) |
| for type_index, rows in enumerate(positions): |
| node_type[rows] = type_index |
| columns.append(torch.from_numpy(node_type).to(dtype)) |
| continue |
| if len(specification) != len(object_types): |
| raise ValueError( |
| f"feature column {column_index} has the wrong number of branches" |
| ) |
| column = np.empty(total, dtype=np.float32) |
| for type_index, branch in enumerate(specification): |
| rows = positions[type_index] |
| if branch == _CALCULATED_ENERGY: |
| index = torch.from_numpy(rows) |
| column[rows] = _grouped_energy( |
| columns[0][index], columns[1][index], per_type[:, type_index] |
| ).numpy() |
| else: |
| column[rows] = flat_values(branch, type_index) |
| columns.append(torch.from_numpy(column).to(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, torch.from_numpy(counts.astype(np.int64)) |
|
|
|
|
| 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 |
| ) |
|
|
