[Feature] Radius Graph (#3829)

* radius graph

* remove trailing whitespaces from docs

* disable invalid name for transform func

* disable radius graph module invalid name

* move pylint disable before init

* fix missing nodes from point set

* update docs indexing

* add compute mode as optional param

* radius graph test

* remove trailing whitespaces

* fix precision when comparing tensors
Co-authored-by: Mufei Li <mufeili1996@gmail.com>

docs/source/api/python/dgl.rst

@@ -26,6 +26,7 @@ Operators for constructing :class:`DGLGraph` from raw data formats.
     rand_bipartite
     knn_graph
     segmented_knn_graph
+    radius_graph
     create_block
     block_to_graph
     merge

docs/source/api/python/nn-pytorch.rst

@@ -104,6 +104,7 @@ Utility Modules
     ~dgl.nn.pytorch.utils.WeightBasis
     ~dgl.nn.pytorch.factory.KNNGraph
     ~dgl.nn.pytorch.factory.SegmentedKNNGraph
+    ~dgl.nn.pytorch.factory.RadiusGraph
     ~dgl.nn.pytorch.utils.JumpingKnowledge
     ~dgl.nn.pytorch.sparse_emb.NodeEmbedding
     ~dgl.nn.pytorch.explain.GNNExplainer

python/dgl/nn/pytorch/factory.py

 """Modules that transforms between graphs and between graph and tensors."""
 import torch.nn as nn
-from ...transforms import knn_graph, segmented_knn_graph
+from ...transforms import knn_graph, segmented_knn_graph, radius_graph
 
 def pairwise_squared_distance(x):
     '''
@@ -230,3 +230,114 @@ class SegmentedKNNGraph(nn.Module):
         """
 
         return segmented_knn_graph(x, self.k, segs, algorithm=algorithm, dist=dist)
+
+
+class RadiusGraph(nn.Module):
+    r"""Layer that transforms one point set into a bidirected graph with
+    neighbors within given distance.
+
+    The RadiusGraph is implemented in the following steps:
+
+    1. Compute an NxN matrix of pairwise distance for all points.
+    2. Pick the points within distance to each point as their neighbors.
+    3. Construct a graph with edges to each point as a node from its neighbors.
+
+    The nodes of the returned graph correspond to the points, where the neighbors
+    of each point are within given distance.
+
+    Parameters
+    ----------
+    r : float
+        Radius of the neighbors.
+    p : float, optional
+        Power parameter for the Minkowski metric. When :attr:`p = 1` it is the
+        equivalent of Manhattan distance (L1 norm) and Euclidean distance
+        (L2 norm) for :attr:`p = 2`.
+
+        (default: 2)
+    self_loop : bool, optional
+        Whether the radius graph will contain self-loops.
+
+        (default: False)
+    compute_mode : str, optional
+        ``use_mm_for_euclid_dist_if_necessary`` - will use matrix multiplication
+        approach to calculate euclidean distance (p = 2) if P > 25 or R > 25
+        ``use_mm_for_euclid_dist`` - will always use matrix multiplication
+        approach to calculate euclidean distance (p = 2)
+        ``donot_use_mm_for_euclid_dist`` - will never use matrix multiplication
+        approach to calculate euclidean distance (p = 2).
+
+        (default: donot_use_mm_for_euclid_dist)
+
+    Examples
+    --------
+    The following examples uses PyTorch backend.
+
+    >>> import dgl
+    >>> from dgl.nn.pytorch.factory import RadiusGraph
+
+    >>> x = torch.tensor([[0.0, 0.0, 1.0],
+    ...                   [1.0, 0.5, 0.5],
+    ...                   [0.5, 0.2, 0.2],
+    ...                   [0.3, 0.2, 0.4]])
+    >>> rg = RadiusGraph(0.75)
+    >>> g = rg(x)  # Each node has neighbors within 0.75 distance
+    >>> g.edges()
+    (tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
+
+    When :attr:`get_distances` is True, forward pass returns the radius graph and
+    distances for the corresponding edges.
+
+    >>> x = torch.tensor([[0.0, 0.0, 1.0],
+    ...                   [1.0, 0.5, 0.5],
+    ...                   [0.5, 0.2, 0.2],
+    ...                   [0.3, 0.2, 0.4]])
+    >>> rg = RadiusGraph(0.75)
+    >>> g, dist = rg(x, get_distances=True)
+    >>> g.edges()
+    (tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
+    >>> dist
+    tensor([[0.7000],
+            [0.6557],
+            [0.6557],
+            [0.2828],
+            [0.7000],
+            [0.2828]])
+    """
+    #pylint: disable=invalid-name
+    def __init__(self, r, p=2, self_loop=False,
+                 compute_mode='donot_use_mm_for_euclid_dist'):
+        super(RadiusGraph, self).__init__()
+        self.r = r
+        self.p = p
+        self.self_loop = self_loop
+        self.compute_mode = compute_mode
+
+    #pylint: disable=invalid-name
+    def forward(self, x, get_distances=False):
+        r"""
+        Forward computation.
+
+        Parameters
+        ----------
+        x : Tensor
+            The point coordinates. :math:`(N, D)` where :math:`N` means the
+            number of points in the point set, and :math:`D` means the size of
+            the features. It can be either on CPU or GPU. Device of the point
+            coordinates specifies device of the radius graph.
+        get_distances : bool, optional
+            Whether to return the distances for the corresponding edges in the
+            radius graph.
+
+            (default: False)
+
+        Returns
+        -------
+        DGLGraph
+            The constructed graph. The node IDs are in the same order as :attr:`x`.
+        torch.Tensor, optional
+            The distances for the edges in the constructed graph. The distances
+            are in the same order as edge IDs.
+        """
+        return radius_graph(x, self.r, self.p, self.self_loop,
+                            self.compute_mode, get_distances)

python/dgl/transforms/functional.py

@@ -20,6 +20,7 @@ from collections import defaultdict
 import numpy as np
 import scipy.sparse as sparse
 import scipy.sparse.linalg
+import torch as th
 
 from .._ffi.function import _init_api
 from ..base import dgl_warning, DGLError, NID, EID
@@ -71,6 +72,7 @@ __all__ = [
     'adj_sum_graph',
     'reorder_graph',
     'norm_by_dst',
+    'radius_graph',
     'random_walk_pe',
     'laplacian_pe'
     ]
@@ -3302,6 +3304,119 @@ def norm_by_dst(g, etype=None):
 
     return norm
 
+def radius_graph(x, r, p=2, self_loop=False,
+                 compute_mode='donot_use_mm_for_euclid_dist', get_distances=False):
+    r"""Construct a graph from a set of points with neighbors within given distance.
+
+    The function transforms the coordinates/features of a point set
+    into a bidirected homogeneous graph. The coordinates of the point
+    set is specified as a matrix whose rows correspond to points and
+    columns correspond to coordinate/feature dimensions.
+
+    The nodes of the returned graph correspond to the points, where the neighbors
+    of each point are within given distance.
+
+    The function requires the PyTorch backend.
+
+    Parameters
+    ----------
+    x : Tensor
+        The point coordinates. It can be either on CPU or GPU.
+        Device of the point coordinates specifies device of the radius graph and
+        ``x[i]`` corresponds to the i-th node in the radius graph.
+    r : float
+        Radius of the neighbors.
+    p : float, optional
+        Power parameter for the Minkowski metric. When :attr:`p = 1` it is the
+        equivalent of Manhattan distance (L1 norm) and Euclidean distance
+        (L2 norm) for :attr:`p = 2`.
+
+        (default: 2)
+    self_loop : bool, optional
+        Whether the radius graph will contain self-loops.
+
+        (default: False)
+    compute_mode : str, optional
+        ``use_mm_for_euclid_dist_if_necessary`` - will use matrix multiplication
+        approach to calculate euclidean distance (p = 2) if P > 25 or R > 25
+        ``use_mm_for_euclid_dist`` - will always use matrix multiplication
+        approach to calculate euclidean distance (p = 2)
+        ``donot_use_mm_for_euclid_dist`` - will never use matrix multiplication
+        approach to calculate euclidean distance (p = 2).
+
+        (default: donot_use_mm_for_euclid_dist)
+    get_distances : bool, optional
+        Whether to return the distances for the corresponding edges in the
+        radius graph.
+
+        (default: False)
+
+    Returns
+    -------
+    DGLGraph
+        The constructed graph. The node IDs are in the same order as :attr:`x`.
+    torch.Tensor, optional
+        The distances for the edges in the constructed graph. The distances are
+        in the same order as edge IDs.
+
+    Examples
+    --------
+
+    The following examples use PyTorch backend.
+
+    >>> import dgl
+    >>> import torch
+
+    >>> x = torch.tensor([[0.0, 0.0, 1.0],
+    ...                   [1.0, 0.5, 0.5],
+    ...                   [0.5, 0.2, 0.2],
+    ...                   [0.3, 0.2, 0.4]])
+    >>> r_g = dgl.radius_graph(x, 0.75)  # Each node has neighbors within 0.75 distance
+    >>> r_g.edges()
+    (tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
+
+    When :attr:`get_distances` is True, function returns the radius graph and
+    distances for the corresponding edges.
+
+    >>> x = torch.tensor([[0.0, 0.0, 1.0],
+    ...                   [1.0, 0.5, 0.5],
+    ...                   [0.5, 0.2, 0.2],
+    ...                   [0.3, 0.2, 0.4]])
+    >>> r_g, dist = dgl.radius_graph(x, 0.75, get_distances=True)
+    >>> r_g.edges()
+    (tensor([0, 1, 2, 2, 3, 3]), tensor([3, 2, 1, 3, 0, 2]))
+    >>> dist
+    tensor([[0.7000],
+            [0.6557],
+            [0.6557],
+            [0.2828],
+            [0.7000],
+            [0.2828]])
+    """
+    # check invalid r
+    if r <= 0:
+        raise DGLError("Invalid r value. expect r > 0, got r = {}".format(r))
+
+    # check empty point set
+    if F.shape(x)[0] == 0:
+        raise DGLError("Find empty point set")
+
+    distances = th.cdist(x, x, p=p, compute_mode=compute_mode)
+
+    if not self_loop:
+        distances.fill_diagonal_(r + 1e-4)
+
+    edges = th.nonzero(distances <= r, as_tuple=True)
+
+    g = convert.graph(edges, num_nodes=x.shape[0], device=x.device)
+
+    if get_distances:
+        distances = distances[edges].unsqueeze(-1)
+
+        return g, distances
+
+    return g
+
 def random_walk_pe(g, k, eweight_name=None):
     r"""Random Walk Positional Encoding, as introduced in
     `Graph Neural Networks with Learnable Structural and Positional Representations

tests/pytorch/test_nn.py

@@ -1330,6 +1330,86 @@ def test_hgt(idtype, in_size, num_heads):
     # TODO(minjie): enable the following check
     #assert th.allclose(y, sorted_y[rev_idx], atol=1e-4, rtol=1e-4)
 
+@pytest.mark.parametrize('self_loop', [True, False])
+@pytest.mark.parametrize('get_distances', [True, False])
+def test_radius_graph(self_loop, get_distances):
+    pos = th.tensor([[0.1, 0.3, 0.4],
+                     [0.5, 0.2, 0.1],
+                     [0.7, 0.9, 0.5],
+                     [0.3, 0.2, 0.5],
+                     [0.2, 0.8, 0.2],
+                     [0.9, 0.2, 0.1],
+                     [0.7, 0.4, 0.4],
+                     [0.2, 0.1, 0.6],
+                     [0.5, 0.3, 0.5],
+                     [0.4, 0.2, 0.6]])
+
+    rg = nn.RadiusGraph(0.3, self_loop=self_loop)
+
+    if get_distances:
+        g, dists = rg(pos, get_distances=get_distances)
+    else:
+        g = rg(pos)
+
+    if self_loop:
+        src_target = th.tensor([0, 0, 1, 2, 3, 3, 3, 3, 3, 4, 5, 6, 6, 7, 7, 7,
+                                8, 8, 8, 8, 9, 9, 9, 9])
+        dst_target = th.tensor([0, 3, 1, 2, 0, 3, 7, 8, 9, 4, 5, 6, 8, 3, 7, 9,
+                                3, 6, 8, 9, 3, 7, 8, 9])
+
+        if get_distances:
+            dists_target = th.tensor([[0.0000],
+                                      [0.2449],
+                                      [0.0000],
+                                      [0.0000],
+                                      [0.2449],
+                                      [0.0000],
+                                      [0.1732],
+                                      [0.2236],
+                                      [0.1414],
+                                      [0.0000],
+                                      [0.0000],
+                                      [0.0000],
+                                      [0.2449],
+                                      [0.1732],
+                                      [0.0000],
+                                      [0.2236],
+                                      [0.2236],
+                                      [0.2449],
+                                      [0.0000],
+                                      [0.1732],
+                                      [0.1414],
+                                      [0.2236],
+                                      [0.1732],
+                                      [0.0000]])
+    else:
+        src_target = th.tensor([0, 3, 3, 3, 3, 6, 7, 7, 8, 8, 8, 9, 9, 9])
+        dst_target = th.tensor([3, 0, 7, 8, 9, 8, 3, 9, 3, 6, 9, 3, 7, 8])
+
+        if get_distances:
+            dists_target = th.tensor([[0.2449],
+                                      [0.2449],
+                                      [0.1732],
+                                      [0.2236],
+                                      [0.1414],
+                                      [0.2449],
+                                      [0.1732],
+                                      [0.2236],
+                                      [0.2236],
+                                      [0.2449],
+                                      [0.1732],
+                                      [0.1414],
+                                      [0.2236],
+                                      [0.1732]])
+
+    src, dst = g.edges()
+
+    assert th.equal(src, src_target)
+    assert th.equal(dst, dst_target)
+
+    if get_distances:
+        assert th.allclose(dists, dists_target, rtol=1e-03)
+
 @parametrize_dtype
 def test_group_rev_res(idtype):
     dev = F.ctx()