mindscience.e3nn.utils.radius_graph_full

mindscience.e3nn.utils.radius_graph_full(x, batch=None, loop=False, flow='source_to_target')[source]

Computes graph edges to all points within a given distance.

Parameters

  • x (Tensor) – node feature matrix.

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

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

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

Returns

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

  • batch (numpy.ndarray) - batch vector.

Raises

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

Examples

>>> from mindscience.e3nn.utils import radius_graph_full
>>> from mindspore import ops, Tensor
>>> x = Tensor(ops.ones((5, 12, 3)))
>>> edge_index, batch = radius_graph_full(x)
>>> print(edge_index.shape)
(2, 660)
>>> print(batch.shape)
(60,)