Commit a9ad9d55 authored by rusty1s's avatar rusty1s
Browse files

nearest doc

parent fd475022
...@@ -20,6 +20,7 @@ The package consists of the following clustering algorithms: ...@@ -20,6 +20,7 @@ The package consists of the following clustering algorithms:
* **[Voxel Grid Pooling](#voxelgrid)** from, *e.g.*, Simonovsky and Komodakis: [Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) (CVPR 2017) * **[Voxel Grid Pooling](#voxelgrid)** from, *e.g.*, Simonovsky and Komodakis: [Dynamic Edge-Conditioned Filters in Convolutional Neural Networks on Graphs](https://arxiv.org/abs/1704.02901) (CVPR 2017)
* **[Iterative Farthest Point Sampling](#farthestpointsampling)** from, *e.g.* Qi *et al.*: [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://arxiv.org/abs/1706.02413) (NIPS 2017) * **[Iterative Farthest Point Sampling](#farthestpointsampling)** from, *e.g.* Qi *et al.*: [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://arxiv.org/abs/1706.02413) (NIPS 2017)
* **[k-NN](#knn-graph)** and **[Radius](#radius-graph)** graph generation * **[k-NN](#knn-graph)** and **[Radius](#radius-graph)** graph generation
* Clustering based on **[Nearest](#nearest)** points
All included operations work on varying data types and are implemented both for CPU and GPU. All included operations work on varying data types and are implemented both for CPU and GPU.
...@@ -143,6 +144,26 @@ tensor([[0, 0, 1, 1, 2, 2, 3, 3], ...@@ -143,6 +144,26 @@ tensor([[0, 0, 1, 1, 2, 2, 3, 3],
[1, 2, 0, 2, 0, 3, 1, 2]]) [1, 2, 0, 2, 0, 3, 1, 2]])
``` ```
## Nearest
Clusters points which are nearest to a given query point in metric space.
```python
import torch
from torch_cluster import nearest
x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch_x = torch.Tensor([0, 0, 0, 0])
query_x = torch.Tensor([[-1, 0], [1, 0]])
query_batch = torch.Tensor([0, 0])
cluster = nearest(x, query_x, batch_x, query_batch)
```
```
print(cluster)
tensor([0, 0, 1, 1])
```
## Running tests ## Running tests
``` ```
......
...@@ -23,7 +23,7 @@ def nearest(x, y, batch_x=None, batch_y=None): ...@@ -23,7 +23,7 @@ def nearest(x, y, batch_x=None, batch_y=None):
>>> batch_x = torch.Tensor([0, 0, 0, 0]) >>> batch_x = torch.Tensor([0, 0, 0, 0])
>>> y = torch.Tensor([[-1, 0], [1, 0]]) >>> y = torch.Tensor([[-1, 0], [1, 0]])
>>> batch_x = torch.Tensor([0, 0]) >>> batch_x = torch.Tensor([0, 0])
>>> out = fps(x, y, batch_x, batch_y) >>> cluster = nearest(x, y, batch_x, batch_y)
""" """
if batch_x is None: if batch_x is None:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment