RadiusGraph

class dgl.nn.pytorch.factory.RadiusGraph(r, p=2, self_loop=False, compute_mode='donot_use_mm_for_euclid_dist')[source]

Bases: Module

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 p = 1 it is the equivalent of Manhattan distance (L1 norm) and Euclidean distance (L2 norm) for 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 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]])
forward(x, get_distances=False)[source]

Forward computation.

Parameters:

  • x (Tensor) – The point coordinates. \((N, D)\) where \(N\) means the number of points in the point set, and \(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 x.

  • torch.Tensor, optional – The distances for the edges in the constructed graph. The distances are in the same order as edge IDs.