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

Bases: torch.nn.modules.module.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

>>> 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]])
>>> 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]])
>>> 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.