radius.py 5.06 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
from typing import Optional
rusty1s's avatar
rusty1s committed
2
3
import torch

4

rusty1s's avatar
rusty1s committed
5
6
7
8
def radius(x: torch.Tensor, y: torch.Tensor, r: float,
           batch_x: Optional[torch.Tensor] = None,
           batch_y: Optional[torch.Tensor] = None,
           max_num_neighbors: int = 32) -> torch.Tensor:
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

rusty1s's avatar
rusty1s committed
27
    .. code-block:: python
rusty1s's avatar
rusty1s committed
28
29
30
31

        import torch
        from torch_cluster import radius

rusty1s's avatar
rusty1s committed
32
33
34
35
36
        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])
        assign_index = radius(x, y, 1.5, batch_x, batch_y)
rusty1s's avatar
docs  
rusty1s committed
37
    """
rusty1s's avatar
rusty1s committed
38
39
40
41

    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
42
    if x.is_cuda:
rusty1s's avatar
rusty1s committed
43
44
45
46
47
48
49
50
        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
51
            torch.cumsum(deg, 0, out=ptr_x[1:])
rusty1s's avatar
rusty1s committed
52
53
54
55
56
        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
57
            batch_size = int(batch_y.max()) + 1
rusty1s's avatar
rusty1s committed
58
59
60
61

            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
62
            torch.cumsum(deg, 0, out=ptr_y[1:])
rusty1s's avatar
rusty1s committed
63
64
65
        else:
            ptr_y = torch.tensor([0, y.size(0)], device=y.device)

66
67
        result = torch.ops.torch_cluster.radius(x, y, ptr_x, ptr_y, r,
                                                max_num_neighbors)
rusty1s's avatar
rusty1s committed
68
69
70
71
72
73
74
    else:
        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)

75
76
77
        batch_x = batch_x.to(x.dtype)
        batch_y = batch_y.to(y.dtype)

rusty1s's avatar
rusty1s committed
78
79
80
81
82
83
        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)

84
85
        result = torch.ops.torch_cluster.radius(x, y, batch_x, batch_y, r,
                                                max_num_neighbors)
rusty1s's avatar
rusty1s committed
86

87
    return result
rusty1s's avatar
rusty1s committed
88

Alexander Liao's avatar
Alexander Liao committed
89

rusty1s's avatar
rusty1s committed
90
91
92
93
def radius_graph(x: torch.Tensor, r: float,
                 batch: Optional[torch.Tensor] = None, loop: bool = False,
                 max_num_neighbors: int = 32,
                 flow: str = 'source_to_target') -> torch.Tensor:
rusty1s's avatar
rusty1s committed
94
    r"""Computes graph edges to all points within a given distance.
rusty1s's avatar
docs  
rusty1s committed
95
96

    Args:
rusty1s's avatar
rusty1s committed
97
98
        x (Tensor): Node feature matrix
            :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
rusty1s's avatar
docs  
rusty1s committed
99
        r (float): The radius.
rusty1s's avatar
rusty1s committed
100
101
102
        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
103
104
        loop (bool, optional): If :obj:`True`, the graph will contain
            self-loops. (default: :obj:`False`)
rusty1s's avatar
docs  
rusty1s committed
105
        max_num_neighbors (int, optional): The maximum number of neighbors to
rusty1s's avatar
rusty1s committed
106
            return for each element in :obj:`y`. (default: :obj:`32`)
rusty1s's avatar
rusty1s committed
107
108
109
        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
110
111
112

    :rtype: :class:`LongTensor`

rusty1s's avatar
rusty1s committed
113
    .. code-block:: python
rusty1s's avatar
rusty1s committed
114
115
116
117

        import torch
        from torch_cluster import radius_graph

rusty1s's avatar
rusty1s committed
118
119
120
        x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
        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
121
122
    """

rusty1s's avatar
rusty1s committed
123
    assert flow in ['source_to_target', 'target_to_source']
124
125
    row, col = radius(x, x, r, batch, batch,
                      max_num_neighbors if loop else max_num_neighbors + 1)
126

Alexander Liao's avatar
Alexander Liao committed
127
128
129
130
    if x.is_cuda:
        row, col = (col, row) if flow == 'source_to_target' else (row, col)
    else:
        row, col = (col, row) if flow == 'target_to_source' else (row, col)
Alexander Liao's avatar
Alexander Liao committed
131

rusty1s's avatar
rusty1s committed
132
133
134
    if not loop:
        mask = row != col
        row, col = row[mask], col[mask]
rusty1s's avatar
rusty1s committed
135
    return torch.stack([row, col], dim=0)