knn.py 4.79 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
from typing import Optional

rusty1s's avatar
rusty1s committed
3
4
5
import torch


rusty1s's avatar
rusty1s committed
6
@torch.jit.script
rusty1s's avatar
rusty1s committed
7
8
def knn(x: torch.Tensor, y: torch.Tensor, k: int,
        batch_x: Optional[torch.Tensor] = None,
rusty1s's avatar
rusty1s committed
9
10
        batch_y: Optional[torch.Tensor] = None, cosine: bool = False,
        num_workers: int = 1) -> torch.Tensor:
rusty1s's avatar
rusty1s committed
11
12
    r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
    :obj:`x`.
rusty1s's avatar
rusty1s committed
13
14

    Args:
rusty1s's avatar
rusty1s committed
15
16
17
18
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
        y (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
rusty1s's avatar
rusty1s committed
19
        k (int): The number of neighbors.
rusty1s's avatar
rusty1s committed
20
21
        batch_x (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
rusty1s's avatar
rusty1s committed
22
23
            node to a specific example. :obj:`batch_x` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
24
25
        batch_y (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
rusty1s's avatar
rusty1s committed
26
27
28
29
30
31
32
33
            node to a specific example. :obj:`batch_y` needs to be sorted.
            (default: :obj:`None`)
        cosine (boolean, optional): If :obj:`True`, will use the Cosine
            distance instead of the Euclidean distance to find nearest
            neighbors. (default: :obj:`False`)
        num_workers (int): Number of workers to use for computation. Has no
            effect in case :obj:`batch_x` or :obj:`batch_y` is not
            :obj:`None`, or the input lies on the GPU. (default: :obj:`1`)
rusty1s's avatar
rusty1s committed
34
35
36

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
37
    .. code-block:: python
rusty1s's avatar
rusty1s committed
38
39
40
41

        import torch
        from torch_cluster import knn

rusty1s's avatar
rusty1s committed
42
43
44
45
46
        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
        batch_x = torch.tensor([0, 0, 0, 0])
        y = torch.Tensor([[-1, 0], [1, 0]])
        batch_x = torch.tensor([0, 0])
        assign_index = knn(x, y, 2, batch_x, batch_y)
rusty1s's avatar
rusty1s committed
47
48
49
50
51
    """

    x = x.view(-1, 1) if x.dim() == 1 else x
    y = y.view(-1, 1) if y.dim() == 1 else y

rusty1s's avatar
rusty1s committed
52
53
54
    if batch_x is not None:
        assert x.size(0) == batch_x.numel()
        batch_size = int(batch_x.max()) + 1
rusty1s's avatar
rusty1s committed
55

rusty1s's avatar
rusty1s committed
56
57
        deg = x.new_zeros(batch_size, dtype=torch.long)
        deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))
rusty1s's avatar
rusty1s committed
58

rusty1s's avatar
rusty1s committed
59
60
        ptr_x = deg.new_zeros(batch_size + 1)
        torch.cumsum(deg, 0, out=ptr_x[1:])
rusty1s's avatar
rusty1s committed
61

rusty1s's avatar
rusty1s committed
62
63
64
    if batch_y is not None:
        assert y.size(0) == batch_y.numel()
        batch_size = int(batch_y.max()) + 1
rusty1s's avatar
rusty1s committed
65

rusty1s's avatar
rusty1s committed
66
67
        deg = y.new_zeros(batch_size, dtype=torch.long)
        deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))
rusty1s's avatar
rusty1s committed
68

rusty1s's avatar
rusty1s committed
69
70
        ptr_y = deg.new_zeros(batch_size + 1)
        torch.cumsum(deg, 0, out=ptr_y[1:])
rusty1s's avatar
rusty1s committed
71
    else:
rusty1s's avatar
rusty1s committed
72
        ptr_y = torch.tensor([0, y.size(0)], device=y.device)
rusty1s's avatar
rusty1s committed
73

rusty1s's avatar
rusty1s committed
74
75
    return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine,
                                       num_workers)
rusty1s's avatar
rusty1s committed
76
77


rusty1s's avatar
rusty1s committed
78
@torch.jit.script
rusty1s's avatar
rusty1s committed
79
80
def knn_graph(x: torch.Tensor, k: int, batch: Optional[torch.Tensor] = None,
              loop: bool = False, flow: str = 'source_to_target',
rusty1s's avatar
rusty1s committed
81
              cosine: bool = False, num_workers: int = 1) -> torch.Tensor:
rusty1s's avatar
rusty1s committed
82
    r"""Computes graph edges to the nearest :obj:`k` points.
rusty1s's avatar
rusty1s committed
83
84

    Args:
rusty1s's avatar
rusty1s committed
85
86
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
rusty1s's avatar
rusty1s committed
87
        k (int): The number of neighbors.
rusty1s's avatar
rusty1s committed
88
89
        batch (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
rusty1s's avatar
rusty1s committed
90
91
            node to a specific example. :obj:`batch` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
92
93
        loop (bool, optional): If :obj:`True`, the graph will contain
            self-loops. (default: :obj:`False`)
rusty1s's avatar
rusty1s committed
94
95
96
        flow (string, optional): The flow direction when using in combination
            with message passing (:obj:`"source_to_target"` or
            :obj:`"target_to_source"`). (default: :obj:`"source_to_target"`)
rusty1s's avatar
rusty1s committed
97
98
99
        cosine (boolean, optional): If :obj:`True`, will use the cosine
            distance instead of euclidean distance to find nearest neighbors.
            (default: :obj:`False`)
rusty1s's avatar
rusty1s committed
100
101
102
        num_workers (int): Number of workers to use for computation. Has no
            effect in case :obj:`batch` is not :obj:`None`, or the input lies
            on the GPU. (default: :obj:`1`)
rusty1s's avatar
rusty1s committed
103
104
105

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
106
    .. code-block:: python
rusty1s's avatar
rusty1s committed
107
108
109
110

        import torch
        from torch_cluster import knn_graph

rusty1s's avatar
rusty1s committed
111
112
113
        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
        batch = torch.tensor([0, 0, 0, 0])
        edge_index = knn_graph(x, k=2, batch=batch, loop=False)
rusty1s's avatar
rusty1s committed
114
115
    """

rusty1s's avatar
rusty1s committed
116
    assert flow in ['source_to_target', 'target_to_source']
rusty1s's avatar
rusty1s committed
117
118
    row, col = knn(x, x, k if loop else k + 1, batch, batch, cosine,
                   num_workers)
rusty1s's avatar
rusty1s committed
119
    row, col = (col, row) if flow == 'source_to_target' else (row, col)
rusty1s's avatar
rusty1s committed
120
121
122
    if not loop:
        mask = row != col
        row, col = row[mask], col[mask]
rusty1s's avatar
rusty1s committed
123
    return torch.stack([row, col], dim=0)