Source code for torch_cluster.nearest

import torch
import scipy.cluster

if torch.cuda.is_available():
    import torch_cluster.nearest_cuda


[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`) .. testsetup:: import torch from torch_cluster import nearest .. testcode:: >>> 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 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) if x.is_cuda: return torch_cluster.nearest_cuda.nearest(x, y, batch_x, batch_y) # Rescale x and y. min_xy = min(x.min().item(), y.min().item()) x, y = x - min_xy, y - min_xy max_xy = max(x.max().item(), y.max().item()) x, y, = x / max_xy, y / max_xy # Concat batch/features to ensure no cross-links between examples exist. x = torch.cat([x, 2 * x.size(1) * batch_x.view(-1, 1).to(x.dtype)], dim=-1) y = torch.cat([y, 2 * y.size(1) * batch_y.view(-1, 1).to(y.dtype)], dim=-1) return torch.from_numpy( scipy.cluster.vq.vq(x.detach().cpu(), y.detach().cpu())[0]).to(torch.long)