|
6 | 6 |
|
7 | 7 |
|
8 | 8 | def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32): |
9 | | - """Finds for each element in `y` all points in `x` within distance `r`. |
| 9 | + r"""Finds for each element in :obj:`y` all points in :obj:`x` within |
| 10 | + distance :obj:`r`. |
10 | 11 |
|
11 | 12 | Args: |
12 | | - x (Tensor): D-dimensional point features. |
13 | | - y (Tensor): D-dimensional point features. |
| 13 | + x (Tensor): Node feature matrix |
| 14 | + :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`. |
| 15 | + y (Tensor): Node feature matrix |
| 16 | + :math:`\mathbf{X} \in \mathbb{R}^{M \times F}`. |
14 | 17 | r (float): The radius. |
15 | | - batch_x (LongTensor, optional): Vector that maps each point to its |
16 | | - example identifier. If :obj:`None`, all points belong to the same |
17 | | - example. If not :obj:`None`, points in the same example need to |
18 | | - have contiguous memory layout and :obj:`batch` needs to be |
19 | | - ascending. (default: :obj:`None`) |
20 | | - batch_y (LongTensor, optional): See `batch_x` (default: :obj:`None`) |
| 18 | + batch_x (LongTensor, optional): Batch vector |
| 19 | + :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each |
| 20 | + node to a specific example. (default: :obj:`None`) |
| 21 | + batch_y (LongTensor, optional): Batch vector |
| 22 | + :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^M`, which assigns each |
| 23 | + node to a specific example. (default: :obj:`None`) |
21 | 24 | max_num_neighbors (int, optional): The maximum number of neighbors to |
22 | | - return for each element in `y`. (default: :obj:`32`) |
| 25 | + return for each element in :obj:`y`. (default: :obj:`32`) |
23 | 26 |
|
24 | 27 | :rtype: :class:`LongTensor` |
25 | 28 |
|
26 | | - Examples:: |
| 29 | + .. testsetup:: |
| 30 | +
|
| 31 | + import torch |
| 32 | + from torch_cluster import radius |
| 33 | +
|
| 34 | + .. testcode:: |
| 35 | +
|
27 | 36 |
|
28 | 37 | >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) |
29 | 38 | >>> batch_x = torch.tensor([0, 0, 0, 0]) |
@@ -63,24 +72,28 @@ def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32): |
63 | 72 |
|
64 | 73 |
|
65 | 74 | def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32): |
66 | | - """Finds for each element in `x` all points in `x` within distance `r`. |
| 75 | + r"""Computes graph edges to all points within a given distance. |
67 | 76 |
|
68 | 77 | Args: |
69 | | - x (Tensor): D-dimensional point features. |
| 78 | + x (Tensor): Node feature matrix |
| 79 | + :math:`\mathbf{X} \in \mathbb{R}^{N \times F}`. |
70 | 80 | r (float): The radius. |
71 | | - batch (LongTensor, optional): Vector that maps each point to its |
72 | | - example identifier. If :obj:`None`, all points belong to the same |
73 | | - example. If not :obj:`None`, points in the same example need to |
74 | | - have contiguous memory layout and :obj:`batch` needs to be |
75 | | - ascending. (default: :obj:`None`) |
| 81 | + batch (LongTensor, optional): Batch vector |
| 82 | + :math:`\mathbf{b} \in {\{ 0, \ldots, B-1\}}^N`, which assigns each |
| 83 | + node to a specific example. (default: :obj:`None`) |
76 | 84 | loop (bool, optional): If :obj:`True`, the graph will contain |
77 | 85 | self-loops. (default: :obj:`False`) |
78 | 86 | max_num_neighbors (int, optional): The maximum number of neighbors to |
79 | | - return for each element in `y`. (default: :obj:`32`) |
| 87 | + return for each element in :obj:`y`. (default: :obj:`32`) |
80 | 88 |
|
81 | 89 | :rtype: :class:`LongTensor` |
82 | 90 |
|
83 | | - Examples:: |
| 91 | + .. testsetup:: |
| 92 | +
|
| 93 | + import torch |
| 94 | + from torch_cluster import radius_graph |
| 95 | +
|
| 96 | + .. testcode:: |
84 | 97 |
|
85 | 98 | >>> x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]]) |
86 | 99 | >>> batch = torch.tensor([0, 0, 0, 0]) |
|
0 commit comments