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
|
import torch
from torch_geometric.testing import is_full_test
from torch_geometric.utils import (
contains_isolated_nodes,
remove_isolated_nodes,
)
def test_contains_isolated_nodes():
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
assert not contains_isolated_nodes(edge_index)
assert contains_isolated_nodes(edge_index, num_nodes=3)
if is_full_test():
jit = torch.jit.script(contains_isolated_nodes)
assert not jit(edge_index)
assert jit(edge_index, num_nodes=3)
edge_index = torch.tensor([[0, 1, 2, 0], [1, 0, 2, 0]])
assert contains_isolated_nodes(edge_index)
def test_remove_isolated_nodes():
edge_index = torch.tensor([[0, 1, 0], [1, 0, 0]])
out, _, mask = remove_isolated_nodes(edge_index)
assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
assert mask.tolist() == [1, 1]
if is_full_test():
jit = torch.jit.script(remove_isolated_nodes)
out, _, mask = jit(edge_index)
assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
assert mask.tolist() == [1, 1]
out, _, mask = remove_isolated_nodes(edge_index, num_nodes=3)
assert out.tolist() == [[0, 1, 0], [1, 0, 0]]
assert mask.tolist() == [1, 1, 0]
edge_index = torch.tensor([[0, 2, 1, 0, 2], [2, 0, 1, 0, 2]])
edge_attr = torch.tensor([1, 2, 3, 4, 5])
out1, out2, mask = remove_isolated_nodes(edge_index, edge_attr)
assert out1.tolist() == [[0, 1, 0, 1], [1, 0, 0, 1]]
assert out2.tolist() == [1, 2, 4, 5]
assert mask.tolist() == [1, 0, 1]
|
