+import numpy as np
+from math import sin, cos, atan2, sqrt
+import scipy
+from typing import Tuple, Union, Sequence
+from scm import plams
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
[docs]
+
def get_rotmat(x: float = None, y: float = None, z: float = None) -> np.ndarray:
+
"""
+
Create a rotation matrix based on the Tait-Bryant sytem.
+
In this system, x, y, and z are angles of rotation around
+
the corresponding axes. This function uses the right-handed
+
convention
+
+
Args:
+
x: Rotation around the x-axis in radians.
+
y: Rotation around the y-axis in radians.
+
z: Rotation around the z-axis in radians.
+
+
Returns:
+
the rotation matrix :math:`\\textbf{R} \\in \\mathbb{R}^{3 \\times 3}` with the specified axis rotations.
+
+
+
.. seealso::
+
+
:func:`apply_rotmat`
+
For applying the rotation matrix to coordinates.
+
+
:func:`rotate`
+
For rotating coordinates directly, given Tait-Bryant angles.
+
+
:meth:`Transform.rotate`
+
The :class:`Transform` class allows you to also rotate.
+
"""
+
# start with identity matrix
+
R = np.eye(3)
+
+
# apply rotation around each axis separately
+
if x is not None:
+
c = cos(x)
+
s = sin(x)
+
R = R @ np.array(([1, 0, 0], [0, c, -s], [0, s, c]))
+
+
if y is not None:
+
c = cos(y)
+
s = sin(y)
+
R = R @ np.array(([c, 0, s], [0, 1, 0], [-s, 0, c]))
+
+
if z is not None:
+
c = cos(z)
+
s = sin(z)
+
R = R @ np.array(([c, -s, 0], [s, c, 0], [0, 0, 1]))
+
+
return R
+
+
+
+
+
[docs]
+
def rotmat_to_angles(R: np.ndarray) -> Tuple[float]:
+
thetax = atan2(R[2, 1], R[2, 2])
+
thetay = atan2(-R[2, 0], sqrt(R[2, 1]**2 + R[2, 2]**2))
+
thetaz = atan2(R[1, 0], R[0, 0])
+
return thetax, thetay, thetaz
+
+
+
+
+
[docs]
+
def apply_rotmat(coords: np.ndarray, R: np.ndarray) -> np.ndarray:
+
"""
+
Apply a rotation matrix to a set of coordinates.
+
+
Args:
+
coords: the coordinates :math`\\in \\mathbb{R}^{n \\times 3}` to rotate.
+
R: the rotation matrix to apply.
+
+
Returns:
+
New coordinates :math`\\in \\mathbb{R}^{n \\times 3}` rotated using the given rotation matrix.
+
+
.. seealso::
+
+
:func:`get_rotmat`
+
For creating a rotation matrix.
+
+
:func:`rotate`
+
For rotating coordinates directly, given Tait-Bryant angles.
+
"""
+
coords = np.atleast_2d(coords)
+
return (R @ coords.T).T.squeeze()
+
+
+
+
+
[docs]
+
def rotate(coords: np.ndarray, x: float = None, y: float = None, z: float = None) -> np.ndarray:
+
"""
+
Build and apply a rotation matrix to a set of coordinates.
+
+
Args:
+
coords: the coordinates :math`\\in \\mathbb{R}^{n \\times 3}` to rotate.
+
x: Rotation around the x-axis in radians.
+
y: Rotation around the y-axis in radians.
+
z: Rotation around the z-axis in radians.
+
+
.. seealso::
+
:func:`get_rotmat`
+
For creating a rotation matrix.
+
"""
+
return apply_rotmat(coords, get_rotmat(x, y, z))
+
+
+
+
+
[docs]
+
def vector_align_rotmat(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+
"""
+
Calculate a rotation matrix that aligns vector a onto vector b.
+
+
Args:
+
a: vector that is to be aligned.
+
b: vector that is the target of the alignment.
+
+
Returns:
+
Rotation matrix R, such that ``geometry.apply_rotmat(a, R) == b``.
+
"""
+
# normalize the vectors first
+
a = np.array(a) / np.linalg.norm(a)
+
b = np.array(b) / np.linalg.norm(b)
+
+
c = a @ b
+
if c == 1:
+
# if a == b we simply return the identity matrix
+
return np.eye(3)
+
if c == -1:
+
# when a == -b we cannot directly calculate R, as 1/(1+c) is undefined
+
# instead, we first create a random rotation matrix and apply it to a
+
# to get a new vector aprime. We then align aprime to b, which is possible since aprime != -b
+
# to get the final rotation matrix we simply multiply the random and alignment rotation matrices
+
Rrand = get_rotmat(np.pi / 3, np.pi / 3, np.pi / 3)
+
return vector_align_rotmat(apply_rotmat(a, Rrand), b) @ Rrand
+
+
v = np.cross(a, b)
+
skew = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
+
R = np.eye(3) + skew + skew @ skew / (1 + c)
+
return R
+
+
+
+
+
[docs]
+
def RMSD(X: np.ndarray, Y: np.ndarray, axis: Union[int, None] = None, use_kabsch: bool = True, include_mirror: bool = False) -> float:
+
r"""
+
Calculate Root Mean Squared Deviations between two sets of points ``X`` and ``Y``.
+
By default Kabsch' algorithm is used to align the sets of points prior to calculating the RMSD.
+
Optionally the axis can be given to calculate the RMSD along different axes.
+
+
RMSD is given as
+
+
:math:`\text{RMSD}(X, Y) = \frac{1}{N}\sqrt{\sum_i^N (X_i - Y_i)^2}`
+
+
when using the Kabsch algorithm to align the two sets of coordinates we first obtain the :class:`KabschTransform` :math:`T_{Kabsch}` and then
+
+
:math:`\text{RMSD}(X, Y) = \frac{1}{N}\sqrt{\sum_i^N (T_{Kabsch}(X_i) - Y_i)^2}`
+
+
Args:
+
X: the first set of coordinates to compare. It must have the same dimensions as ``Y``.
+
Y: the second set of coordinates to compare. It must have the same dimensions as ``X``.
+
axis: axis to compare. Defaults to ``None``.
+
use_kabsch: whether to use Kabsch' algorithm to align ``X`` and ``Y`` before calculating the RMSD. Defaults to ``True``.
+
include_mirror: return the lowest value between the RMSD of the supplied coordinates and also the RMSD of mirrored X with Y.
+
This will only be done if ``use_kabsch == True``.
+
+
Returns:
+
RMSD in the units of X and Y. If ``axis`` is set to an integer this function will return a vector of RMSD's along that axis.
+
+
.. note::
+
It is generally recommended to enable the use of the Kabsch-Umeyama algorithm prior to calculating the RMSD.
+
This will ensure you get the lowest possible RMSD for you sets of coordinates.
+
+
.. seealso::
+
:class:`KabschTransform`
+
"""
+
+
X = np.array(X)
+
Y = np.array(Y)
+
+
assert X.shape == Y.shape
+
+
# apply Kabsch transform
+
if use_kabsch:
+
Tkabsch = KabschTransform(X, Y)
+
Xprime = Tkabsch(X)
+
+
rmsd = np.sqrt(np.sum((Xprime - Y) ** 2, axis=axis) / X.shape[0])
+
+
# if we include the mirror image we have to apply a reflection to the coordinates
+
# and then recalculate the kabsch transform in the mirror coordinates
+
# then we calculate the new RMSD and take the smaller of the new and old RMSD
+
if include_mirror and use_kabsch:
+
Tmirror = Transform()
+
Tmirror.reflect()
+
Tkabsch_mirror = KabschTransform(Tmirror(X), Y)
+
Xprime = Tkabsch_mirror(Tmirror(X))
+
rmsd_mirrored = np.sqrt(np.sum((Xprime - Y) ** 2, axis=axis) / X.shape[0])
+
rmsd = min(rmsd, rmsd_mirrored)
+
+
return rmsd
+
+
+
+
+
[docs]
+
def random_points_on_sphere(shape: Tuple[int], radius: float = 1) -> np.ndarray:
+
"""
+
Generate random points on a sphere with a specified radius.
+
+
Args:
+
shape: The shape of the resulting points, generally shape[0] coordinates with shape[1] dimensions
+
radius: The radius of the sphere to generate the points on.
+
+
Returns:
+
Array of coordinates on a sphere.
+
"""
+
x = np.random.randn(*shape)
+
x = x / np.linalg.norm(x, axis=1, keepdims=True) * radius
+
return x
+
+
+
+
+
[docs]
+
def random_points_in_anular_sphere(shape: Tuple[int], min_radius: float = 0, max_radius: float = 1):
+
"""
+
Generate random points in an sphere or anular sphere with specified radii.
+
An anular sphere is a hollow sphere of a certain thickness.
+
+
Args:
+
shape: The shape of the resulting points, generally shape[0] coordinates with shape[1] dimensions
+
min_radius: The lowest radius of the sphere to generate the points in.
+
max_radius: The largest radius of the sphere to generate the points in.
+
+
Returns:
+
Array of coordinates on a sphere.
+
"""
+
random_radii = np.random.rand(shape[0]) * (max_radius - min_radius) + min_radius
+
return random_points_on_sphere(shape, random_radii)
+
+
+
+
+
[docs]
+
def random_points_on_spheroid(coordinates: np.ndarray, Nsamples: int = 1, margin: float = 0):
+
"""
+
Generate random points on a spheroid generated by a set of coordinates.
+
+
Args:
+
coordinates: The (n x dim) set of coordinates that is used to generate the minimum-volume spheroid.
+
Nsamples: The number of samples to return.
+
margin: the spacing between the sampling spheroid and the minimum-volume spheroid.
+
+
Returns:
+
Array of coordinates on a spheroid.
+
"""
+
# for this to work we should first get the centroid of our molecule
+
centroid = np.mean(coordinates, axis=0)
+
# and get the centered coordiantes
+
Xc = coordinates - centroid
+
+
# we then do a singular-value decomposition to obtain
+
# the three principle components (Vh) with their eigenvalues (s)
+
_, s, Vh = scipy.linalg.svd(Xc)
+
+
# then compute a transformation matrix for generating the correct spheroid
+
transform = Transform()
+
transform.translate(centroid)
+
transform.rotate((np.diag(s/2 + margin) @ Vh).T)
+
+
# to sample the spheroid we generate points on a
+
# sphere and transform them to our spheroid
+
p = random_points_on_sphere((Nsamples, Xc.shape[1]))
+
return transform(p)
+
+
+
+
+
[docs]
+
def parameter(coordinates, *indices, pyramidal=False):
+
'''
+
Return geometry information about a set of coordinates given indices.
+
'''
+
assert 1 <= len(indices) <= 4, "Number of indices must be between 1, 2, 3 or 4"
+
+
coordinates = np.array(coordinates)
+
selected_coords = [coordinates[i] for i in indices]
+
+
if len(indices) == 1:
+
return selected_coords[0]
+
+
if len(indices) == 2:
+
return np.linalg.norm(selected_coords[0] - selected_coords[1])
+
+
if len(indices) == 3:
+
a = selected_coords[0] - selected_coords[1]
+
b = selected_coords[2] - selected_coords[1]
+
a = a / np.linalg.norm(a)
+
b = b / np.linalg.norm(b)
+
+
return np.arccos(a @ b) / np.pi * 180
+
+
if len(indices) == 4 and not pyramidal:
+
a = selected_coords[0] - selected_coords[1]
+
b = selected_coords[2] - selected_coords[1]
+
+
u = selected_coords[1] - selected_coords[2]
+
v = selected_coords[3] - selected_coords[2]
+
+
n1, n2 = np.cross(a, b), np.cross(u, v)
+
+
n1 = n1 / np.linalg.norm(n1)
+
n2 = n2 / np.linalg.norm(n2)
+
+
return np.arccos(n1 @ n2) / np.pi * 180
+
+
+
if len(indices) == 4 and pyramidal:
+
ang1 = parameter(coordinates, indices[1], indices[0], indices[2])
+
ang2 = parameter(coordinates, indices[2], indices[0], indices[3])
+
ang3 = parameter(coordinates, indices[3], indices[0], indices[1])
+
+
return 360 - ang1 - ang2 - ang3
+
+
+
+