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
|
import torch
from torch_geometric.data import Data
from torch_geometric.transforms import LocalDegreeProfile
def test_target_indegree():
assert str(LocalDegreeProfile()) == 'LocalDegreeProfile()'
edge_index = torch.tensor([[0, 1, 1, 2], [1, 0, 2, 1]])
x = torch.tensor([[1.0], [1.0], [1.0], [1.0]]) # One isolated node.
expected = torch.tensor([
[1, 2, 2, 2, 0],
[2, 1, 1, 1, 0],
[1, 2, 2, 2, 0],
[0, 0, 0, 0, 0],
], dtype=torch.float)
data = Data(edge_index=edge_index, num_nodes=x.size(0))
data = LocalDegreeProfile()(data)
assert torch.allclose(data.x, expected, atol=1e-2)
data = Data(edge_index=edge_index, x=x)
data = LocalDegreeProfile()(data)
assert torch.allclose(data.x[:, :1], x)
assert torch.allclose(data.x[:, 1:], expected, atol=1e-2)
|
