mindscience.e3nn.so2_conv.init_edge_rot_mat

mindscience.e3nn.so2_conv.init_edge_rot_mat(edge_distance_vec)

Initialize the rotation matrix from the edge distance vector.

Parameters

edge_distance_vec (Tensor) – Edge distance vector with shape (batch_size, 3).

Returns

Tensor, Rotation matrix with shape (batch_size, 3, 3).

Examples

>>> import mindspore as ms
>>> from mindspore import Tensor
>>> from mindscience.e3nn.so2_conv import init_edge_rot_mat
>>> edge_vec = ms.Tensor([[1.0, 0.0, 0.0],
...                       [0.0, 1.0, 0.0],
...                       [0.0, 0.0, 1.0]])
>>> rot_mat = init_edge_rot_mat(edge_vec)
>>> print(rot_mat.shape)
(3, 3, 3)