radius.py 4.29 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import torch
rusty1s's avatar
rusty1s committed
2
import scipy.spatial
rusty1s's avatar
rusty1s committed
3
4

if torch.cuda.is_available():
5
    import torch_cluster.radius_cuda
rusty1s's avatar
rusty1s committed
6
7
8


def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
rusty1s's avatar
rusty1s committed
9
10
    r"""Finds for each element in :obj:`y` all points in :obj:`x` within
    distance :obj:`r`.
rusty1s's avatar
docs  
rusty1s committed
11
12

    Args:
rusty1s's avatar
rusty1s committed
13
14
15
        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
16
            :math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
rusty1s's avatar
docs  
rusty1s committed
17
        r (float): The radius.
rusty1s's avatar
rusty1s committed
18
19
20
21
22
23
        batch_x (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example. (default: :obj:`None`)
        batch_y (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each
            node to a specific example. (default: :obj:`None`)
rusty1s's avatar
docs  
rusty1s committed
24
        max_num_neighbors (int, optional): The maximum number of neighbors to
rusty1s's avatar
rusty1s committed
25
            return for each element in :obj:`y`. (default: :obj:`32`)
rusty1s's avatar
docs  
rusty1s committed
26
27
28

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
29
30
31
32
33
34
35
    .. testsetup::

        import torch
        from torch_cluster import radius

    .. testcode::

rusty1s's avatar
docs  
rusty1s committed
36
37

        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
rusty1s's avatar
rusty1s committed
38
        >>> batch_x = torch.tensor([0, 0, 0, 0])
rusty1s's avatar
docs  
rusty1s committed
39
        >>> y = torch.Tensor([[-1, 0], [1, 0]])
Vadim Bereznyuk's avatar
typos  
Vadim Bereznyuk committed
40
        >>> batch_y = torch.tensor([0, 0])
rusty1s's avatar
rusty1s committed
41
        >>> assign_index = radius(x, y, 1.5, batch_x, batch_y)
rusty1s's avatar
docs  
rusty1s committed
42
    """
rusty1s's avatar
rusty1s committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58

    if batch_x is None:
        batch_x = x.new_zeros(x.size(0), dtype=torch.long)

    if batch_y is None:
        batch_y = y.new_zeros(y.size(0), dtype=torch.long)

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

    assert x.dim() == 2 and batch_x.dim() == 1
    assert y.dim() == 2 and batch_y.dim() == 1
    assert x.size(1) == y.size(1)
    assert x.size(0) == batch_x.size(0)
    assert y.size(0) == batch_y.size(0)

rusty1s's avatar
rusty1s committed
59
    if x.is_cuda:
60
61
        return torch_cluster.radius_cuda.radius(x, y, r, batch_x, batch_y,
                                                max_num_neighbors)
rusty1s's avatar
rusty1s committed
62

rusty1s's avatar
rusty1s committed
63
64
65
    x = torch.cat([x, 2 * r * batch_x.view(-1, 1).to(x.dtype)], dim=-1)
    y = torch.cat([y, 2 * r * batch_y.view(-1, 1).to(y.dtype)], dim=-1)

66
67
68
69
    tree = scipy.spatial.cKDTree(x.detach().numpy())
    _, col = tree.query(
        y.detach().numpy(), k=max_num_neighbors, distance_upper_bound=r + 1e-8)
    col = [torch.from_numpy(c).to(torch.long) for c in col]
rusty1s's avatar
rusty1s committed
70
71
    row = [torch.full_like(c, i) for i, c in enumerate(col)]
    row, col = torch.cat(row, dim=0), torch.cat(col, dim=0)
72
    mask = col < int(tree.n)
rusty1s's avatar
mask  
rusty1s committed
73
    return torch.stack([row[mask], col[mask]], dim=0)
rusty1s's avatar
rusty1s committed
74
75


rusty1s's avatar
rusty1s committed
76
77
78
79
80
81
def radius_graph(x,
                 r,
                 batch=None,
                 loop=False,
                 max_num_neighbors=32,
                 flow='source_to_target'):
rusty1s's avatar
rusty1s committed
82
    r"""Computes graph edges to all points within a given distance.
rusty1s's avatar
docs  
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
docs  
rusty1s committed
87
        r (float): The radius.
rusty1s's avatar
rusty1s committed
88
89
90
        batch (LongTensor, optional): Batch vector
            :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each
            node to a specific example. (default: :obj:`None`)
rusty1s's avatar
rusty1s committed
91
92
        loop (bool, optional): If :obj:`True`, the graph will contain
            self-loops. (default: :obj:`False`)
rusty1s's avatar
docs  
rusty1s committed
93
        max_num_neighbors (int, optional): The maximum number of neighbors to
rusty1s's avatar
rusty1s committed
94
            return for each element in :obj:`y`. (default: :obj:`32`)
rusty1s's avatar
rusty1s committed
95
96
97
        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
docs  
rusty1s committed
98
99
100

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
101
102
103
104
105
106
    .. testsetup::

        import torch
        from torch_cluster import radius_graph

    .. testcode::
rusty1s's avatar
docs  
rusty1s committed
107
108

        >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
rusty1s's avatar
rusty1s committed
109
110
        >>> batch = torch.tensor([0, 0, 0, 0])
        >>> edge_index = radius_graph(x, r=1.5, batch=batch, loop=False)
rusty1s's avatar
docs  
rusty1s committed
111
112
    """

rusty1s's avatar
rusty1s committed
113
114
115
    assert flow in ['source_to_target', 'target_to_source']
    row, col = radius(x, x, r, batch, batch, max_num_neighbors + 1)
    row, col = (col, row) if flow == 'source_to_target' else (row, col)
rusty1s's avatar
rusty1s committed
116
117
118
    if not loop:
        mask = row != col
        row, col = row[mask], col[mask]
rusty1s's avatar
rusty1s committed
119
    return torch.stack([row, col], dim=0)