mindchemistry.e3.utils.radius_graph

View Source On Gitee
mindchemistry.e3.utils.radius_graph(x, r, batch=None, loop=False, max_num_neighbors=32, flow='source_to_target')[source]

Computes graph edges to all points within a given distance.

Parameters

  • x (ndarray) – node feature matrix.

  • r (ndarray, float) – the radius.

  • batch (Tensor) – batch vector. If it is none, then calculate and return. Default: None.

  • loop (bool) – whether contain self-loops in the graph. Dufault: False.

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

  • flow (str) – {'source_to_target', 'target_to_source'}, the flow direction when using in combination with message passing. Dufault: 'source_to_target'.

Returns

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

batch (ndarray) - batch vector.

Raises

ValueError – If flow is not in {'source_to_target', 'target_to_source'}.

Supported Platforms:

Ascend

Examples

>>> from mindchemistry.e3.utils import radius_graph
>>> import numpy as np
>>> np.random.seed(1)
>>> x = np.random.random((5, 12, 3))
>>> r = 0.5
>>> edge_index, batch = radius_graph(x, r)
>>> print(edge_index.shape)
(2, 162)
>>> print(batch.shape)
(60,)