test_knn.py 2.38 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
3
4
from itertools import product

import pytest
import torch
rusty1s's avatar
rusty1s committed
5
import scipy.spatial
rusty1s's avatar
rusty1s committed
6
from torch_cluster import knn, knn_graph
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
from .utils import grad_dtypes, devices, tensor
rusty1s's avatar
rusty1s committed
9
10


rusty1s's avatar
rusty1s committed
11
12
13
14
def to_set(edge_index):
    return set([(i, j) for i, j in edge_index.t().tolist()])


rusty1s's avatar
rusty1s committed
15
@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
rusty1s's avatar
rusty1s committed
16
def test_knn(dtype, device):
rusty1s's avatar
rusty1s committed
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
    x = tensor([
        [-1, -1],
        [-1, +1],
        [+1, +1],
        [+1, -1],
        [-1, -1],
        [-1, +1],
        [+1, +1],
        [+1, -1],
    ], dtype, device)
    y = tensor([
        [1, 0],
        [-1, 0],
    ], dtype, device)

    batch_x = tensor([0, 0, 0, 0, 1, 1, 1, 1], torch.long, device)
    batch_y = tensor([0, 1], torch.long, device)

rusty1s's avatar
rusty1s committed
35
36
    edge_index = knn(x, y, 2)
    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 0), (1, 1)])
rusty1s's avatar
rusty1s committed
37

rusty1s's avatar
rusty1s committed
38
39
    edge_index = knn(x, y, 2, batch_x, batch_y)
    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
rusty1s's avatar
rusty1s committed
40

rusty1s's avatar
rusty1s committed
41
    if x.is_cuda:
rusty1s's avatar
rusty1s committed
42
        edge_index = knn(x, y, 2, batch_x, batch_y, cosine=True)
rusty1s's avatar
rusty1s committed
43
        assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])
rusty1s's avatar
rusty1s committed
44

45
46
47
48
49
50
    # Skipping a batch
    batch_x = tensor([0, 0, 0, 0, 2, 2, 2, 2], torch.long, device)
    batch_y = tensor([0, 2], torch.long, device)
    edge_index = knn(x, y, 2, batch_x, batch_y)
    assert to_set(edge_index) == set([(0, 2), (0, 3), (1, 4), (1, 5)])

rusty1s's avatar
rusty1s committed
51
52
53
54
55
56
57
58
59
60

@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
def test_knn_graph(dtype, device):
    x = tensor([
        [-1, -1],
        [-1, +1],
        [+1, +1],
        [+1, -1],
    ], dtype, device)

rusty1s's avatar
rusty1s committed
61
62
63
    edge_index = knn_graph(x, k=2, flow='target_to_source')
    assert to_set(edge_index) == set([(0, 1), (0, 3), (1, 0), (1, 2), (2, 1),
                                      (2, 3), (3, 0), (3, 2)])
rusty1s's avatar
rusty1s committed
64

rusty1s's avatar
rusty1s committed
65
66
67
    edge_index = knn_graph(x, k=2, flow='source_to_target')
    assert to_set(edge_index) == set([(1, 0), (3, 0), (0, 1), (2, 1), (1, 2),
                                      (3, 2), (0, 3), (2, 3)])
68
69


Matthias Fey's avatar
Matthias Fey committed
70
@pytest.mark.parametrize('dtype,device', product([torch.float], devices))
71
def test_knn_graph_large(dtype, device):
rusty1s's avatar
rusty1s committed
72
    x = torch.randn(1000, 3, dtype=dtype, device=device)
rusty1s's avatar
rusty1s committed
73

rusty1s's avatar
rusty1s committed
74
    edge_index = knn_graph(x, k=5, flow='target_to_source', loop=True)
rusty1s's avatar
rusty1s committed
75

rusty1s's avatar
rusty1s committed
76
    tree = scipy.spatial.cKDTree(x.cpu().numpy())
rusty1s's avatar
rusty1s committed
77
78
79
    _, col = tree.query(x.cpu(), k=5)
    truth = set([(i, j) for i, ns in enumerate(col) for j in ns])

rusty1s's avatar
rusty1s committed
80
    assert to_set(edge_index.cpu()) == truth