tfep.utils.geometry.reference_frame_rotation_matrix

tfep.utils.geometry.reference_frame_rotation_matrix(axis_atom_positions: Tensor, plane_atom_positions: Tensor, axis: Tensor, plane_axis: Tensor, plane_normal: Tensor | None = None, project_on_positive_axis: bool = False) → Tensor[source]

Return the rotation matrix required to rotate the frame of reference based on two atoms.

After the rotation matrix is applied to the coordinates, axis_atom_positions lie on the given axis vector while plane_atom_positions lie on the plane spanned by the axis and plane_axis vectors.

Parameters:

  • axis_atom_positions (torch.Tensor) – Shape (batch_size, 3). The position of the atom placed on axis.

  • plane_atom_positions (torch.Tensor) – Shape (batch_size, 3). The position of the atom placed on the axis-plane_axis plane.

  • axis (torch.Tensor) – Shape (3,). The axis on which to the axis atom is placed. Must be a unit vector.

  • plane_axis (torch.Tensor,) – Shape (3,). The second axis used to determine the plane where the plane atom is placed. Must be a unit vector and not parallel to axis.

  • plane_normal (Optional[torch.Tensor]) – The vector normal to axis and plane_axis. If not given, it is computed here.

  • project_on_positive_axis (bool) – If True, the axis atom is rotated so that it always lies on the positive axis. Otherwise, it is rotated on the positive or negative axis based on whichever is closest.

    Note that if this is True, a transformation that flips the sign of the coordinate of the axis atom might become impossible to invert in practice.

Returns:

rotation_matrices – Shape (batch_size, 3, 3). The rotation matrices.

Return type:

torch.Tensor

Examples

>>> # Initialize the coordinates.
>>> batch_size, n_atoms = 2, 4
>>> coordinates = torch.randn(batch_size, n_atoms, 3)

>>> # Fix the orientation of the coordiante frames based on the 2nd and 4th atoms.
>>> axis_atom_pos = coordinates[:, 1]
>>> plane_atom_pos = coordinates[:, 3]
>>> rotation_matrices = reference_frame_rotation_matrix(
...     axis_atom_pos,
...     plane_atom_pos,
...     axis=torch.tensor([1.0, 0, 0]),  # axis atom lies on x-axis
...     plane_axis=torch.tensor([0.0, 0, 1]),  # plane atomlies on x-z plane
... )
...

>>> # Rotate the coordinates.
>>> new_coordinates = batchwise_rotate(coordinates, rotation_matrices)
>>> # Reverse the change of reference frame.
>>> old_coordinates = batchwise_rotate(new_coordinates, rotation_matrices, inverse=True)