| """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 |
| ) |
|
|
