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 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117
|
import torch
from torch_geometric.data import Data
from torch_geometric.testing import withPackage
from torch_geometric.transforms import (
AddLaplacianEigenvectorPE,
AddRandomWalkPE,
)
@withPackage('scipy')
def test_add_laplacian_eigenvector_pe():
x = torch.randn(6, 4)
edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5],
[1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
data = Data(x=x, edge_index=edge_index)
transform = AddLaplacianEigenvectorPE(k=3)
assert str(transform) == 'AddLaplacianEigenvectorPE()'
out = transform(data)
assert out.laplacian_eigenvector_pe.size() == (6, 3)
transform = AddLaplacianEigenvectorPE(k=3, attr_name=None)
out = transform(data)
assert out.x.size() == (6, 4 + 3)
transform = AddLaplacianEigenvectorPE(k=3, attr_name='x')
out = transform(data)
assert out.x.size() == (6, 3)
# Output tests:
edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5, 2, 5],
[1, 0, 4, 0, 4, 1, 3, 2, 5, 3, 5, 2]])
data = Data(x=x, edge_index=edge_index)
transform1 = AddLaplacianEigenvectorPE(k=1, is_undirected=True)
transform2 = AddLaplacianEigenvectorPE(k=1, is_undirected=False)
# Clustering test with first non-trivial eigenvector (Fiedler vector)
pe = transform1(data).laplacian_eigenvector_pe
pe_cluster_1 = pe[[0, 1, 4]]
pe_cluster_2 = pe[[2, 3, 5]]
assert not torch.allclose(pe_cluster_1, pe_cluster_2)
assert torch.allclose(pe_cluster_1, pe_cluster_1.mean())
assert torch.allclose(pe_cluster_2, pe_cluster_2.mean())
pe = transform2(data).laplacian_eigenvector_pe
pe_cluster_1 = pe[[0, 1, 4]]
pe_cluster_2 = pe[[2, 3, 5]]
assert not torch.allclose(pe_cluster_1, pe_cluster_2)
assert torch.allclose(pe_cluster_1, pe_cluster_1.mean())
assert torch.allclose(pe_cluster_2, pe_cluster_2.mean())
@withPackage('scipy')
def test_eigenvector_permutation_invariance():
edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5],
[1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
data = Data(edge_index=edge_index, num_nodes=6)
perm = torch.randperm(data.num_nodes)
transform = AddLaplacianEigenvectorPE(
k=2,
is_undirected=True,
attr_name='x',
)
out1 = transform(data)
transform = AddLaplacianEigenvectorPE(
k=2,
is_undirected=True,
attr_name='x',
)
out2 = transform(data.subgraph(perm))
assert torch.allclose(out1.x[perm].abs(), out2.x.abs(), atol=1e-6)
def test_add_random_walk_pe():
x = torch.randn(6, 4)
edge_index = torch.tensor([[0, 1, 0, 4, 1, 4, 2, 3, 3, 5],
[1, 0, 4, 0, 4, 1, 3, 2, 5, 3]])
data = Data(x=x, edge_index=edge_index)
transform = AddRandomWalkPE(walk_length=3)
assert str(transform) == 'AddRandomWalkPE()'
out = transform(data)
assert out.random_walk_pe.size() == (6, 3)
transform = AddRandomWalkPE(walk_length=3, attr_name=None)
out = transform(data)
assert out.x.size() == (6, 4 + 3)
transform = AddRandomWalkPE(walk_length=3, attr_name='x')
out = transform(data)
assert out.x.size() == (6, 3)
# Output tests:
assert out.x.tolist() == [
[0.0, 0.5, 0.25],
[0.0, 0.5, 0.25],
[0.0, 0.5, 0.00],
[0.0, 1.0, 0.00],
[0.0, 0.5, 0.25],
[0.0, 0.5, 0.00],
]
edge_index = torch.tensor([[0, 1, 2], [0, 1, 2]])
data = Data(edge_index=edge_index, num_nodes=4)
out = transform(data)
assert out.x.tolist() == [
[1.0, 1.0, 1.0],
[1.0, 1.0, 1.0],
[1.0, 1.0, 1.0],
[0.0, 0.0, 0.0],
]
|