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

rusty1s's avatar
rusty1s committed
3
import scipy.cluster
Matthias Fey's avatar
Matthias Fey committed
4
import torch
rusty1s's avatar
rusty1s committed
5
6


Matthias Fey's avatar
Matthias Fey committed
7
8
9
10
11
12
def nearest(
    x: torch.Tensor,
    y: torch.Tensor,
    batch_x: Optional[torch.Tensor] = None,
    batch_y: Optional[torch.Tensor] = None,
) -> torch.Tensor:
rusty1s's avatar
typo  
rusty1s committed
13
    r"""Clusters points in :obj:`x` together which are nearest to a given query
rusty1s's avatar
rusty1s committed
14
    point in :obj:`y`.
rusty1s's avatar
docs  
rusty1s committed
15
16

    Args:
rusty1s's avatar
rusty1s committed
17
18
19
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
        y (Tensor): Node feature matrix
Vadim Bereznyuk's avatar
typos  
Vadim Bereznyuk committed
20
            :math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
rusty1s's avatar
rusty1s committed
21
22
        batch_x (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
rusty1s's avatar
rusty1s committed
23
24
            node to a specific example. :obj:`batch_x` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
25
26
        batch_y (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
rusty1s's avatar
rusty1s committed
27
28
            node to a specific example. :obj:`batch_y` needs to be sorted.
            (default: :obj:`None`)
rusty1s's avatar
docs  
rusty1s committed
29

rusty1s's avatar
rusty1s committed
30
31
32
    :rtype: :class:`LongTensor`

    .. code-block:: python
rusty1s's avatar
rusty1s committed
33
34
35
36

        import torch
        from torch_cluster import nearest

rusty1s's avatar
rusty1s committed
37
38
39
40
41
        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_y = torch.tensor([0, 0])
        cluster = nearest(x, y, batch_x, batch_y)
rusty1s's avatar
docs  
rusty1s committed
42
43
    """

rusty1s's avatar
rusty1s committed
44
45
    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
update  
rusty1s committed
46
    assert x.size(1) == y.size(1)
rusty1s's avatar
rusty1s committed
47

48
    if batch_x is not None and (batch_x[1:] - batch_x[:-1] < 0).any():
Matthias Fey's avatar
Matthias Fey committed
49
        raise ValueError("'batch_x' is not sorted")
50
    if batch_y is not None and (batch_y[1:] - batch_y[:-1] < 0).any():
Matthias Fey's avatar
Matthias Fey committed
51
        raise ValueError("'batch_y' is not sorted")
52

rusty1s's avatar
rusty1s committed
53
    if x.is_cuda:
rusty1s's avatar
rusty1s committed
54
55
56
57
58
59
60
61
        if batch_x is not None:
            assert x.size(0) == batch_x.numel()
            batch_size = int(batch_x.max()) + 1

            deg = x.new_zeros(batch_size, dtype=torch.long)
            deg.scatter_add_(0, batch_x, torch.ones_like(batch_x))

            ptr_x = deg.new_zeros(batch_size + 1)
rusty1s's avatar
fix  
rusty1s committed
62
            torch.cumsum(deg, 0, out=ptr_x[1:])
rusty1s's avatar
rusty1s committed
63
64
65
66
67
        else:
            ptr_x = torch.tensor([0, x.size(0)], device=x.device)

        if batch_y is not None:
            assert y.size(0) == batch_y.numel()
rusty1s's avatar
fix  
rusty1s committed
68
            batch_size = int(batch_y.max()) + 1
rusty1s's avatar
rusty1s committed
69
70
71
72
73

            deg = y.new_zeros(batch_size, dtype=torch.long)
            deg.scatter_add_(0, batch_y, torch.ones_like(batch_y))

            ptr_y = deg.new_zeros(batch_size + 1)
rusty1s's avatar
fix  
rusty1s committed
74
            torch.cumsum(deg, 0, out=ptr_y[1:])
rusty1s's avatar
rusty1s committed
75
76
77
        else:
            ptr_y = torch.tensor([0, y.size(0)], device=y.device)

Matthias Fey's avatar
Matthias Fey committed
78
79
80
81
82
83
84
        # If an instance in `batch_x` is non-empty, it must be non-empty in
        # `batch_y `as well:
        nonempty_ptr_x = (ptr_x[1:] - ptr_x[:-1]) > 0
        nonempty_ptr_y = (ptr_y[1:] - ptr_y[:-1]) > 0
        if not torch.equal(nonempty_ptr_x, nonempty_ptr_y):
            raise ValueError("Some batch indices occur in 'batch_x' "
                             "that do not occur in 'batch_y'")
85

rusty1s's avatar
rusty1s committed
86
        return torch.ops.torch_cluster.nearest(x, y, ptr_x, ptr_y)
Matthias Fey's avatar
Matthias Fey committed
87

rusty1s's avatar
rusty1s committed
88
    else:
Matthias Fey's avatar
Matthias Fey committed
89
90
91
92
93

        if batch_x is None and batch_y is not None:
            batch_x = x.new_zeros(x.size(0), dtype=torch.long)
        if batch_y is None and batch_x is not None:
            batch_y = y.new_zeros(y.size(0), dtype=torch.long)
94

rusty1s's avatar
rusty1s committed
95
        # Translate and rescale x and y to [0, 1].
rusty1s's avatar
rusty1s committed
96
        if batch_x is not None and batch_y is not None:
Matthias Fey's avatar
Matthias Fey committed
97
98
99
100
101
102
103
            # If an instance in `batch_x` is non-empty, it must be non-empty in
            # `batch_y `as well:
            unique_batch_x = batch_x.unique_consecutive()
            unique_batch_y = batch_y.unique_consecutive()
            if not torch.equal(unique_batch_x, unique_batch_y):
                raise ValueError("Some batch indices occur in 'batch_x' "
                                 "that do not occur in 'batch_y'")
104

rusty1s's avatar
typo  
rusty1s committed
105
106
107
108
109
            assert x.dim() == 2 and batch_x.dim() == 1
            assert y.dim() == 2 and batch_y.dim() == 1
            assert x.size(0) == batch_x.size(0)
            assert y.size(0) == batch_y.size(0)

rusty1s's avatar
rusty1s committed
110
111
112
113
114
115
116
117
118
119
120
            min_xy = min(x.min().item(), y.min().item())
            x, y = x - min_xy, y - min_xy

            max_xy = max(x.max().item(), y.max().item())
            x.div_(max_xy)
            y.div_(max_xy)

            # Concat batch/features to ensure no cross-links between examples.
            D = x.size(-1)
            x = torch.cat([x, 2 * D * batch_x.view(-1, 1).to(x.dtype)], -1)
            y = torch.cat([y, 2 * D * batch_y.view(-1, 1).to(y.dtype)], -1)
rusty1s's avatar
rusty1s committed
121
122
123
124

        return torch.from_numpy(
            scipy.cluster.vq.vq(x.detach().cpu(),
                                y.detach().cpu())[0]).to(torch.long)