-
Notifications
You must be signed in to change notification settings - Fork 1
/
test_functional.py
80 lines (62 loc) · 2.35 KB
/
test_functional.py
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
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
import torch
from torch_geometric.testing import withPackage
from mooon import drop_edge, drop_node, drop_path
def test_drop_node():
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])
edge_weight = torch.tensor([1., 2., 3., 4., 5., 6])
out = drop_node(edge_index, training=False)
assert edge_index.tolist() == out[0].tolist()
torch.manual_seed(5)
out = drop_node(edge_index)
print(out)
assert out[0].tolist() == [[2, 3], [3, 2]]
assert out[1] is None
torch.manual_seed(5)
out = drop_node(edge_index, edge_weight)
print(out)
assert out[0].tolist() == [[2, 3], [3, 2]]
assert out[1].tolist() == [
5.,
6.,
]
def test_drop_edge():
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])
edge_weight = torch.tensor([1., 2., 3., 4., 5., 6])
out = drop_edge(edge_index, training=False)
assert edge_index.tolist() == out[0].tolist()
assert out[1] is None
torch.manual_seed(5)
out = drop_edge(edge_index)
assert out[0].tolist() == [[0, 1, 2, 2], [1, 2, 1, 3]]
assert out[1] is None
torch.manual_seed(5)
out = drop_edge(edge_index, edge_weight)
assert out[0].tolist() == [[0, 1, 2, 2], [1, 2, 1, 3]]
assert out[1].tolist() == [1., 3., 4., 5.]
@withPackage('torch_cluster')
def test_drop_path():
edge_index = torch.tensor([[0, 1, 1, 2, 2, 3], [1, 0, 2, 1, 3, 2]])
edge_weight = torch.tensor([1., 2., 3., 4., 5., 6])
out = drop_path(edge_index, training=False)
assert edge_index.tolist() == out[0].tolist()
assert out[1] is None
torch.manual_seed(4)
out = drop_path(edge_index, p=0.2)
assert out[0].tolist() == [[1, 2, 3], [2, 3, 2]]
assert out[1] is None
torch.manual_seed(4)
out = drop_path(edge_index, edge_weight, p=0.2)
assert out[0].tolist() == [[1, 2, 3], [2, 3, 2]]
assert out[1].tolist() == [3., 5., 6.]
# test with unsorted edges
torch.manual_seed(6)
edge_index = torch.tensor([[3, 5, 2, 2, 2, 1], [1, 0, 0, 1, 3, 2]])
out = drop_path(edge_index, p=0.2)
assert out[0].tolist() == [[2, 3, 5], [0, 1, 0]]
assert out[1] is None
# test with isolated nodes
torch.manual_seed(7)
edge_index = torch.tensor([[0, 1, 2, 3], [1, 0, 2, 4]])
out = drop_path(edge_index, p=0.2)
assert out[0].tolist() == [[2, 3], [2, 4]]
assert out[1] is None