mindchemistry.e3.utils.radius_full

View Source On Gitee
mindchemistry.e3.utils.radius_full(x, y, batch_x=None, batch_y=None)[source]

Find all points in x for each element in y.

Parameters

  • x (Tensor) – node feature matrix.

  • y (Tensor) – node feature matrix.

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

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

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.

Supported Platforms:

Ascend

Examples

>>> from mindchemistry.e3.utils import radius_full
>>> from mindspore import ops, Tensor
>>> x = Tensor(ops.ones((5, 12, 3)))
>>> edge_index, batch_x, batch_y = radius_full(x, x)
>>> print(edge_index.shape)
(2, 720)
>>> print(batch_x.shape)
(60,)
>>> print(batch_y.shape)
(60,)