from .max_pool import max_pool, max_pool_x
from .avg_pool import avg_pool, avg_pool_x
from .graclus import graclus
from .voxel_grid import voxel_grid
from .topk_pool import TopKPooling
from .sag_pool import SAGPooling
from .edge_pool import EdgePooling
try:
import torch_cluster
except ImportError:
torch_cluster = None
[docs]def fps(x, batch=None, ratio=0.5, random_start=True):
r""""A sampling algorithm from the `"PointNet++: Deep Hierarchical Feature
Learning on Point Sets in a Metric Space"
<https://arxiv.org/abs/1706.02413>`_ paper, which iteratively samples the
most distant point with regard to the rest points.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
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`)
ratio (float, optional): Sampling ratio. (default: :obj:`0.5`)
random_start (bool, optional): If set to :obj:`False`, use the first
node in :math:`\mathbf{X}` as starting node. (default: obj:`True`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import fps
x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch = torch.tensor([0, 0, 0, 0])
index = fps(x, batch, ratio=0.5)
"""
if torch_cluster is None:
raise ImportError('`fps` requires `torch-cluster`.')
return torch_cluster.fps(x, batch, ratio, random_start)
[docs]def knn(x, y, k, batch_x=None, batch_y=None, cosine=False):
r"""Finds for each element in :obj:`y` the :obj:`k` nearest points in
:obj:`x`.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
y (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{M \times F}`.
k (int): The number of neighbors.
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`)
cosine (boolean, optional): If :obj:`True`, will use the cosine
distance instead of euclidean distance to find nearest neighbors.
(default: :obj:`False`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import knn
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_x = torch.tensor([0, 0])
assign_index = knn(x, y, 2, batch_x, batch_y)
"""
if torch_cluster is None:
raise ImportError('`knn` requires `torch-cluster`.')
return torch_cluster.knn(x, y, k, batch_x, batch_y, cosine)
[docs]def knn_graph(x, k, batch=None, loop=False, flow='source_to_target',
cosine=False):
r"""Computes graph edges to the nearest :obj:`k` points.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
k (int): The number of neighbors.
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`)
loop (bool, optional): If :obj:`True`, the graph will contain
self-loops. (default: :obj:`False`)
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"`)
cosine (boolean, optional): If :obj:`True`, will use the cosine
distance instead of euclidean distance to find nearest neighbors.
(default: :obj:`False`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import knn_graph
x = torch.Tensor([[-1, -1], [-1, 1], [1, -1], [1, 1]])
batch = torch.tensor([0, 0, 0, 0])
edge_index = knn_graph(x, k=2, batch=batch, loop=False)
"""
if torch_cluster is None:
raise ImportError('`knn_graph` requires `torch-cluster`.')
return torch_cluster.knn_graph(x, k, batch, loop, flow, cosine)
[docs]def radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32):
r"""Finds for each element in :obj:`y` all points in :obj:`x` within
distance :obj:`r`.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
y (Tensor): Node feature matrix
:math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
r (float): The radius.
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`)
max_num_neighbors (int, optional): The maximum number of neighbors to
return for each element in :obj:`y`. (default: :obj:`32`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import radius
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)
"""
if torch_cluster is None:
raise ImportError('`radius` requires `torch-cluster`.')
return torch_cluster.radius(x, y, r, batch_x, batch_y, max_num_neighbors)
[docs]def radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32,
flow='source_to_target'):
r"""Computes graph edges to all points within a given distance.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
r (float): The radius.
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`)
loop (bool, optional): If :obj:`True`, the graph will contain
self-loops. (default: :obj:`False`)
max_num_neighbors (int, optional): The maximum number of neighbors to
return for each element in :obj:`y`. (default: :obj:`32`)
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"`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import radius_graph
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)
"""
if torch_cluster is None:
raise ImportError('`radius_graph` requires `torch-cluster`.')
return torch_cluster.radius_graph(x, r, batch, loop, max_num_neighbors,
flow)
[docs]def nearest(x, y, batch_x=None, batch_y=None):
r"""Clusters points in :obj:`x` together which are nearest to a given query
point in :obj:`y`.
Args:
x (Tensor): Node feature matrix
:math:`\mathbf{X} \in \mathbb{R}^{N \times F}`.
y (Tensor): Node feature matrix
:math:`\mathbf{Y} \in \mathbb{R}^{M \times F}`.
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`)
:rtype: :class:`LongTensor`
.. code-block:: python
import torch
from torch_geometric.nn import nearest
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)
"""
if torch_cluster is None:
raise ImportError('`radius` requires `torch-cluster`.')
return torch_cluster.nearest(x, y, batch_x, batch_y)
__all__ = [
'TopKPooling',
'SAGPooling',
'EdgePooling',
'max_pool',
'avg_pool',
'max_pool_x',
'avg_pool_x',
'graclus',
'voxel_grid',
'fps',
'knn',
'knn_graph',
'radius',
'radius_graph',
'nearest',
]