import torch.nn.functional
import autogl
from ._basic import BaseFeatureGenerator
from ._pyg_impl import degree, scatter_min, scatter_max, scatter_mean, scatter_std
from .._feature_engineer_registry import FeatureEngineerUniversalRegistry
[docs]@FeatureEngineerUniversalRegistry.register_feature_engineer("LocalDegreeProfile")
class LocalDegreeProfileGenerator(BaseFeatureGenerator):
r"""Appends the Local Degree Profile (LDP) from the `"A Simple yet
Effective Baseline for Non-attribute Graph Classification"
<https://arxiv.org/abs/1811.03508>`_ paper
.. math::
\mathbf{x}_i = \mathbf{x}_i \, \Vert \, (\deg(i), \min(DN(i)),
\max(DN(i)), \textrm{mean}(DN(i)), \textrm{std}(DN(i)))
to the node features, where :math:`DN(i) = \{ \deg(j) \mid j \in
\mathcal{N}(i) \}`.
"""
def _extract_nodes_feature(self, data: autogl.data.Data) -> torch.Tensor:
row, col = data.edge_index
if data.x is not None and isinstance(data.x, torch.Tensor):
N = data.x.size(0)
else:
N = (torch.max(data.edge_index).item() + 1)
deg = degree(row, N, dtype=torch.float)
deg_col = deg[col]
min_deg, _ = scatter_min(deg_col, row, dim_size=N)
min_deg[min_deg > 10000] = 0
max_deg, _ = scatter_max(deg_col, row, dim_size=N)
max_deg[max_deg < -10000] = 0
mean_deg = scatter_mean(deg_col, row, dim_size=N)
std_deg = scatter_std(deg_col, row, dim_size=N)
x = torch.stack([deg, min_deg, max_deg, mean_deg, std_deg], dim=1)
return x
@FeatureEngineerUniversalRegistry.register_feature_engineer("NormalizeFeatures")
class NormalizeFeatures(BaseFeatureGenerator):
def __init__(self):
super(NormalizeFeatures, self).__init__(override_features=True)
def _extract_nodes_feature(self, data: autogl.data.Data) -> torch.Tensor:
if data.x is not None and isinstance(data.x, torch.Tensor):
data.x.div_(data.x.sum(dim=-1, keepdim=True).clamp_(min=1.))
return data.x
[docs]@FeatureEngineerUniversalRegistry.register_feature_engineer("OneHotDegree")
class OneHotDegreeGenerator(BaseFeatureGenerator):
r"""Adds the node degree as one hot encodings to the node features.
Args:
max_degree (int): Maximum degree.
in_degree (bool, optional): If set to :obj:`True`, will compute the
in-degree of nodes instead of the out-degree.
(default: :obj:`False`)
cat (bool, optional): Concat node degrees to node features instead
of replacing them. (default: :obj:`True`)
"""
def __init__(
self, max_degree: int = 1000,
in_degree: bool = False, cat: bool = True
):
self.__max_degree: int = max_degree
self.__in_degree: bool = in_degree
self.__cat: bool = cat
super(OneHotDegreeGenerator, self).__init__()
def _extract_nodes_feature(self, data: autogl.data.Data) -> torch.Tensor:
idx, x = data.edge_index[1 if self.__in_degree else 0], data.x
deg = degree(idx, data.num_nodes, dtype=torch.long)
deg = torch.nn.functional.one_hot(
deg, num_classes=self.__max_degree + 1
).to(torch.float)
return deg