tfep.utils.geometry.rotation_matrix_3d

tfep.utils.geometry.rotation_matrix_3d(angles, directions)

Return the matrix rotating vectors for the given angle about a direction.

The rotation matrix is built using Rodrigues’ rotation formula.

Parameters:

• angles (torch.Tensor) – A tensor of shape (batch_size,).

• directions (torch.Tensor) – A tensor of shape (batch_size, 3) or (3,).

Returns:

R – R[i] is the 3 by 3 matrix rotating by the angle angles[i] about the vector directions[i].

Return type:

torch.Tensor