mindscience.e3nn.utils.radius

mindscience.e3nn.utils.radius(x, y, r, batch_x=None, batch_y=None, max_num_neighbors=32)[source]

Find all points in x for each element in y within distance r.

Parameters

  • x (ndarray) – node feature matrix of x.

  • y (ndarray) – node feature matrix of y.

  • r (Union[ndarray, float]) – the radius.

  • batch_x (ndarray, optional) – batch vector of x. If it is none, then calculate based on x and return. Default: None.

  • batch_y (ndarray, optional) – batch vector of y. If it is none, then calculate based on y and return. Default: None.

  • max_num_neighbors (int, optional) – The maximum number of neighbors to return for each element in y. Default: 32.

Returns

  • edge_index (numpy.ndarray) - including edges of source and destination.

  • batch_x (numpy.ndarray) - batch vector of x.

  • batch_y (numpy.ndarray) - batch vector of y.

Raises

ValueError – If the last dimension of x and y do not match.

Examples

>>> from mindscience.e3nn.utils import radius
>>> import numpy as np
>>> np.random.seed(1)
>>> x = np.random.random((5, 12, 3))
>>> r = 0.5
>>> edge_index, batch_x, batch_y = radius(x, x, r)
>>> print(edge_index.shape)
(2, 222)
>>> print(batch_x.shape)
(60,)
>>> print(batch_y.shape)
(60,)