tfep.utils.geometry.batchwise_rotate

tfep.utils.geometry.batchwise_rotate(x, rotation_matrices, inverse=False)

Rotate a batch of configurations with their respective rotation matrix.

Parameters:

• x (torch.Tensor) – A tensor of shape (batch_size, n_vectors, 3).

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

• inverse (bool, optional) – If True the inverse rotation is performed.

Returns:

y – A tensor of shape (batch_size, n_vectors, 3)) where y[b][i] is the result of rotating vector x[b][i] with the rotation matrix rotation_matrices[b].

Return type:

torch.Tensor