test_knn.py 2.04 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
11


@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
rusty1s's avatar
rusty1s committed
12
def test_knn(dtype, device):
rusty1s's avatar
rusty1s committed
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
    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
31
32
33
    row, col = knn(x, y, 2)
    assert row.tolist() == [0, 0, 1, 1]
    assert col.tolist() == [2, 3, 0, 1]
rusty1s's avatar
rusty1s committed
34

rusty1s's avatar
rusty1s committed
35
    row, col = knn(x, y, 2, batch_x, batch_y)
rusty1s's avatar
rusty1s committed
36
37
38
    assert row.tolist() == [0, 0, 1, 1]
    assert col.tolist() == [2, 3, 4, 5]

rusty1s's avatar
rusty1s committed
39
40
41
42
43
    if x.is_cuda:
        row, col = knn(x, y, 2, batch_x, batch_y, cosine=True)
        assert row.tolist() == [0, 0, 1, 1]
        assert col.tolist() == [0, 1, 4, 5]

rusty1s's avatar
rusty1s committed
44
45
46
47
48
49
50
51
52
53

@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
54
    row, col = knn_graph(x, k=2, flow='target_to_source')
rusty1s's avatar
rusty1s committed
55
56
    assert row.tolist() == [0, 0, 1, 1, 2, 2, 3, 3]
    assert col.tolist() == [1, 3, 0, 2, 1, 3, 0, 2]
rusty1s's avatar
rusty1s committed
57
58
59
60

    row, col = knn_graph(x, k=2, flow='source_to_target')
    assert row.tolist() == [1, 3, 0, 2, 1, 3, 0, 2]
    assert col.tolist() == [0, 0, 1, 1, 2, 2, 3, 3]
61
62
63
64


@pytest.mark.parametrize('dtype,device', product(grad_dtypes, devices))
def test_knn_graph_large(dtype, device):
rusty1s's avatar
rusty1s committed
65
66
67
68
69
70
71
72
73
74
75
    x = torch.randn(1000, 3)

    row, col = knn_graph(x, k=5, flow='target_to_source', loop=True,
                         num_workers=6)
    pred = set([(i, j) for i, j in zip(row.tolist(), col.tolist())])

    tree = scipy.spatial.cKDTree(x.numpy())
    _, col = tree.query(x.cpu(), k=5)
    truth = set([(i, j) for i, ns in enumerate(col) for j in ns])

    assert pred == truth