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

rusty1s's avatar
rusty1s committed
3
4
5
import torch


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

    Args:
rusty1s's avatar
rusty1s committed
20
21
22
23
        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
24
        k (int): The number of neighbors.
rusty1s's avatar
rusty1s committed
25
26
        batch_x (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
rusty1s's avatar
rusty1s committed
27
28
            node to a specific example. :obj:`batch_x` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
29
30
        batch_y (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
rusty1s's avatar
rusty1s committed
31
32
33
34
35
36
37
38
            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`)
39
40
        batch_size (int, optional): The number of examples :math:`B`.
            Automatically calculated if not given. (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
41
42
43

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
44
    .. code-block:: python
rusty1s's avatar
rusty1s committed
45
46
47
48

        import torch
        from torch_cluster import knn

rusty1s's avatar
rusty1s committed
49
50
51
        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]])
Ahmed Ibrahim's avatar
Ahmed Ibrahim committed
52
        batch_y = torch.tensor([0, 0])
rusty1s's avatar
rusty1s committed
53
        assign_index = knn(x, y, 2, batch_x, batch_y)
rusty1s's avatar
rusty1s committed
54
    """
rusty1s's avatar
rusty1s committed
55
56
    if x.numel() == 0 or y.numel() == 0:
        return torch.empty(2, 0, dtype=torch.long, device=x.device)
rusty1s's avatar
rusty1s committed
57
58
59

    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
60
    x, y = x.contiguous(), y.contiguous()
rusty1s's avatar
rusty1s committed
61

62
63
64
65
66
67
68
69
70
    if batch_size is None:
        batch_size = 1
        if batch_x is not None:
            assert x.size(0) == batch_x.numel()
            batch_size = int(batch_x.max()) + 1
        if batch_y is not None:
            assert y.size(0) == batch_y.numel()
            batch_size = max(batch_size, int(batch_y.max()) + 1)
    assert batch_size > 0
rusty1s's avatar
rusty1s committed
71

72
73
74
75
76
77
78
79
    ptr_x: Optional[torch.Tensor] = None
    ptr_y: Optional[torch.Tensor] = None
    if batch_size > 1:
        assert batch_x is not None
        assert batch_y is not None
        arange = torch.arange(batch_size + 1, device=x.device)
        ptr_x = torch.bucketize(arange, batch_x)
        ptr_y = torch.bucketize(arange, batch_y)
rusty1s's avatar
rusty1s committed
80

rusty1s's avatar
rusty1s committed
81
82
    return torch.ops.torch_cluster.knn(x, y, ptr_x, ptr_y, k, cosine,
                                       num_workers)
rusty1s's avatar
rusty1s committed
83
84


85
86
87
88
89
90
91
92
93
94
def knn_graph(
    x: torch.Tensor,
    k: int,
    batch: Optional[torch.Tensor] = None,
    loop: bool = False,
    flow: str = 'source_to_target',
    cosine: bool = False,
    num_workers: int = 1,
    batch_size: Optional[int] = None,
) -> torch.Tensor:
rusty1s's avatar
rusty1s committed
95
    r"""Computes graph edges to the nearest :obj:`k` points.
rusty1s's avatar
rusty1s committed
96
97

    Args:
rusty1s's avatar
rusty1s committed
98
99
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
rusty1s's avatar
rusty1s committed
100
        k (int): The number of neighbors.
rusty1s's avatar
rusty1s committed
101
102
        batch (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
rusty1s's avatar
rusty1s committed
103
104
            node to a specific example. :obj:`batch` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
105
106
        loop (bool, optional): If :obj:`True`, the graph will contain
            self-loops. (default: :obj:`False`)
107
        flow (string, optional): The flow direction when used in combination
rusty1s's avatar
rusty1s committed
108
109
            with message passing (:obj:`"source_to_target"` or
            :obj:`"target_to_source"`). (default: :obj:`"source_to_target"`)
110
111
        cosine (boolean, optional): If :obj:`True`, will use the Cosine
            distance instead of Euclidean distance to find nearest neighbors.
rusty1s's avatar
rusty1s committed
112
            (default: :obj:`False`)
rusty1s's avatar
rusty1s committed
113
114
115
        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`)
116
117
        batch_size (int, optional): The number of examples :math:`B`.
            Automatically calculated if not given. (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
118
119
120

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
121
    .. code-block:: python
rusty1s's avatar
rusty1s committed
122
123
124
125

        import torch
        from torch_cluster import knn_graph

rusty1s's avatar
rusty1s committed
126
127
128
        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
129
130
    """

rusty1s's avatar
rusty1s committed
131
    assert flow in ['source_to_target', 'target_to_source']
rusty1s's avatar
rusty1s committed
132
    edge_index = knn(x, x, k if loop else k + 1, batch, batch, cosine,
133
                     num_workers, batch_size)
rusty1s's avatar
rusty1s committed
134
135
136
137
138
139

    if flow == 'source_to_target':
        row, col = edge_index[1], edge_index[0]
    else:
        row, col = edge_index[0], edge_index[1]

rusty1s's avatar
rusty1s committed
140
141
142
    if not loop:
        mask = row != col
        row, col = row[mask], col[mask]
rusty1s's avatar
rusty1s committed
143

rusty1s's avatar
rusty1s committed
144
    return torch.stack([row, col], dim=0)