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
|
import torch
from torch_geometric.data import Data
from torch_geometric.transforms import VirtualNode
def test_virtual_node():
assert str(VirtualNode()) == 'VirtualNode()'
x = torch.randn(4, 16)
edge_index = torch.tensor([[2, 0, 2], [3, 1, 0]])
edge_weight = torch.rand(edge_index.size(1))
edge_attr = torch.randn(edge_index.size(1), 8)
data = Data(x=x, edge_index=edge_index, edge_weight=edge_weight,
edge_attr=edge_attr, num_nodes=x.size(0))
data = VirtualNode()(data)
assert len(data) == 6
assert data.x.size() == (5, 16)
assert torch.allclose(data.x[:4], x)
assert data.x[4:].abs().sum() == 0
assert data.edge_index.tolist() == [[2, 0, 2, 0, 1, 2, 3, 4, 4, 4, 4],
[3, 1, 0, 4, 4, 4, 4, 0, 1, 2, 3]]
assert data.edge_weight.size() == (11, )
assert torch.allclose(data.edge_weight[:3], edge_weight)
assert data.edge_weight[3:].abs().sum() == 8
assert data.edge_attr.size() == (11, 8)
assert torch.allclose(data.edge_attr[:3], edge_attr)
assert data.edge_attr[3:].abs().sum() == 0
assert data.num_nodes == 5
assert data.edge_type.tolist() == [0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]
data = Data(x=x, edge_index=torch.empty(2, 0, dtype=torch.long))
data = VirtualNode()(data)
assert len(data) == 3
assert data.x.size() == (5, 16)
assert torch.allclose(data.x[:4], x)
assert data.x[4:].abs().sum() == 0
assert data.edge_index.tolist() == [
[0, 1, 2, 3, 4, 4, 4, 4],
[4, 4, 4, 4, 0, 1, 2, 3],
]
assert data.edge_type.tolist() == [1, 1, 1, 1, 2, 2, 2, 2]
|
