RMSD including initial alignment by principal moments of inertia (#242)

* Adding in RMSD to do rigorous_rmsd logic after aligning about the principal moments of inertia

* Cleaning up comments

* [pre-commit.ci] auto fixes from pre-commit hooks

* Adding test cases

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

molSimplify/Classes/mol3D.py

@@ -2533,6 +2533,32 @@ def molsize(self):
                 maxd = distance(cm, atom.coords())
         return maxd
 
+    def moments_of_inertia(self):
+        """
+        Determines the moments of inertia for the object, in the specified coordinates
+        (after centering about the center of mass)
+
+        Returns
+        -------
+            I : np.array
+                Moments of inertia tensor
+        """
+        I = np.zeros((3, 3))
+        #center about the center of mass
+        cm = self.centermass()
+        for atom in self.atoms:
+            atom.setcoords(np.array(atom.coords()) - cm)
+            I[0, 0] += atom.mass * (atom.coords()[1]**2 + atom.coords()[2]**2) #xx
+            I[1, 1] += atom.mass * (atom.coords()[0]**2 + atom.coords()[2]**2) #yy
+            I[2, 2] += atom.mass * (atom.coords()[0]**2 + atom.coords()[1]**2) #zz
+            I[0, 1] -= atom.mass * (atom.coords()[0] * atom.coords()[1]) #xy
+            I[1, 0] -= atom.mass * (atom.coords()[0] * atom.coords()[1]) #yx
+            I[0, 2] -= atom.mass * (atom.coords()[2] * atom.coords()[0]) #xz
+            I[2, 0] -= atom.mass * (atom.coords()[2] * atom.coords()[0]) #zx
+            I[1, 2] -= atom.mass * (atom.coords()[1] * atom.coords()[2]) #yz
+            I[2, 1] -= atom.mass * (atom.coords()[1] * atom.coords()[2]) #zy
+        return I
+
     def overlapcheck(self, mol, silence=False):
         """
         Measure the smallest distance between an atom and a point.
@@ -2631,6 +2657,39 @@ def populateBOMatrixAug(self):
                 molBOMat[error_idx[i].tolist()[0], error_idx[i].tolist()[1]] = 1
         return (molBOMat)
 
+    def principal_moments_of_inertia(self, return_transform=False):
+        """
+        Returns the diagonalized moments of inertia tensor, and optionally the
+        matrices required to diagonalize this tensor.
+
+        Parameters
+        ----------
+            return_transform : bool
+                Flag for if the matrices used to diagonalize I should be returned.
+                Default is False.
+        Returns
+        -------
+            pmom : np.array
+                3x1 array of the principal moments of inertia, in the provided Cartesian frame.
+            P : np.array
+                Rotation matrix to rotate original molecule
+                in order to have the axes correspond to the principal moments of inertia
+                Use by running atom.setcoords(P.dot(atom.coords())) for each atom
+
+        """
+        I = self.moments_of_inertia()
+        #diagonalize the moments of inertia
+        eigvals, eigvecs = np.linalg.eig(I)
+        D = np.linalg.inv(eigvecs) @ I @ eigvecs
+        pmom = np.diag(D)
+        #transformation for the original coordinates defined as inverse of eigenvectors
+        P = np.linalg.inv(eigvecs)
+
+        if return_transform:
+            return pmom, P
+        else:
+            return pmom
+
     def printxyz(self):
         """
         Print XYZ info of mol3D class instance to stdout. To write to file

molSimplify/Scripts/rmsd.py

@@ -372,7 +372,8 @@ def reorder_distance(p_atoms, q_atoms, p_coord, q_coord):
 
 
 def rmsd_reorder_rotate(p_atoms, q_atoms, p_coord, q_coord,
-                        rotation="kabsch", reorder="hungarian"):
+                        rotation="kabsch", reorder="hungarian",
+                        translate=True):
     """Reorder and rotate for RMSD.
 
     Parameters
@@ -389,6 +390,10 @@ def rmsd_reorder_rotate(p_atoms, q_atoms, p_coord, q_coord,
             Rotation method. Default is kabsch.
         reorder : str, optional
             Reorder method. Default is hungarian.
+        translate : bool, optional
+            Whether or not the molecules should be translated
+            such that their centroid is at the origin.
+            Default is True.
 
     Returns
     -------
@@ -404,11 +409,11 @@ def rmsd_reorder_rotate(p_atoms, q_atoms, p_coord, q_coord,
         print(("Warning: Atom types do not match!",
                np.unique(p_atoms), np.unique(q_atoms)))
         return 1000
-
-    p_cent = centroid(p_coord)
-    q_cent = centroid(q_coord)
-    p_coord -= p_cent
-    q_coord -= q_cent
+    if translate:
+        p_cent = centroid(p_coord)
+        q_cent = centroid(q_coord)
+        p_coord -= p_cent
+        q_coord -= q_cent
 
     # set rotation method
     if rotation.lower() == "kabsch":
@@ -472,6 +477,81 @@ def rigorous_rmsd(mol1, mol2, rotation: str = "kabsch",
                                       rotation=rotation, reorder=reorder)
     return result_rmsd
 
+def align_rmsd(mol1, mol2, rotation: str = "kabsch",
+               reorder: str = "hungarian") -> float:
+    """
+    Computes the RMSD between 2 mol objects after:
+    - translating them both such that the center of mass is at the origin
+    - projecting the coordinates onto the principal axes
+    - reordering x, y, z such that Ixx < Iyy < Izz
+    (will allow for 180degree rotations about x, y, z, as well as
+    reflections about the xy, xz, yz, and all three of those planes)
+
+    Parameters
+    ----------
+        mol1 : mol3D
+            mol3D instance of initial molecule.
+        mol2 : np.mol3D
+            mol3D instance of final molecule.
+        rotation : str, optional
+            Rotation method. Default is kabsch.
+        reorder : str, optional
+            Reorder method. Default is hungarian.
+
+    Returns
+    -------
+        rmsd : float
+            Resulting RMSD from aligning and rotating.
+    """
+
+    mol1_atoms = mol1.symvect()
+    cm1 = mol1.centermass()
+    pmom1, P1 = mol1.principal_moments_of_inertia(return_transform=True)
+    for atom in mol1.atoms:
+        atom.setcoords(np.array(atom.coords()) - cm1) #center
+        atom.setcoords(P1.dot(atom.coords())) #project onto principal moments
+        #sort the coordinates so that the largest princiipal moment is first
+        atom.setcoords(np.array(atom.coords())[np.argsort(pmom1)].flatten())
+    mol1_coords = mol1.coordsvect()
+
+    #note: the above aligns the largest moment of inertia with z, the smallest with x
+    #does not account for whether it is aligned with + or - part of an axis
+    #so, have to allow for 180deg rotations and reflections as well
+
+    rmsd = np.inf
+    x_rot = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]]) #180 about x
+    y_rot = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]]) #180 about y
+    z_rot = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, 1]]) #180 about z
+    x_ref = np.array([[-1, 0, 0], [0, 1, 0], [0, 0, 1]]) #reflect about yz
+    y_ref = np.array([[1, 0, 0], [0, -1, 0], [0, 0, 1]]) #reflect about xz
+    z_ref = np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]]) #reflect about xy
+    a_ref = np.array([[-1, 0, 0], [0, -1, 0], [0, 0, -1]]) #reflect all 3
+    transformations = [
+        np.eye(3), #no change
+        x_rot, y_rot, z_rot,
+        x_ref, y_ref, z_ref, a_ref
+    ]
+    mol2_atoms = mol2.symvect()
+    cm2 = mol2.centermass()
+    pmom2, P2 = mol2.principal_moments_of_inertia(return_transform=True)
+    for atom in mol2.atoms:
+        atom.setcoords(np.array(atom.coords()) - cm2) #center
+        atom.setcoords(P2.dot(atom.coords())) #project onto principal moments
+        #sort the coordinates so that the largest princiipal moment is first
+        atom.setcoords(np.array(atom.coords())[np.argsort(pmom2)].flatten())
+    for transformation in transformations:
+        for atom in mol2.atoms:
+            atom.setcoords(transformation @ np.array(atom.coords()))
+        mol2_coords = mol2.coordsvect()
+        #revert transformations
+        for atom in mol2.atoms:
+            atom.setcoords(np.linalg.inv(transformation) @ np.array(atom.coords()))
+
+        result_rmsd = rmsd_reorder_rotate(mol1_atoms, mol2_atoms, mol1_coords, mol2_coords,
+                                          rotation=rotation, reorder=reorder, translate=True)
+        if result_rmsd < rmsd:
+            rmsd = result_rmsd
+    return rmsd
 
 def test_case():
     p_atoms = np.array(["N", "H", "H", "H"])

tests/test_rmsd.py

@@ -2,7 +2,7 @@
 import numpy as np
 from molSimplify.Classes.mol3D import mol3D
 from molSimplify.Scripts.geometry import rotate_around_axis
-from molSimplify.Scripts.rmsd import rigorous_rmsd, quaternion_rotate
+from molSimplify.Scripts.rmsd import rigorous_rmsd, quaternion_rotate, align_rmsd
 from scipy.spatial.transform import Rotation
 
 
@@ -31,8 +31,6 @@ def test_rigorous_rmsd(resource_path_root, path1, path2, ref_hungarian, ref_none
     r = rigorous_rmsd(mol1, mol2, reorder='none')
     assert abs(r - ref_none) < atol
 
-
-@pytest.mark.skip
 def test_methane_rotation(atol=1e-3):
     """This test case is intended to show the problem with our current RMSD implementation"""
     # XYZ data copied from
@@ -53,9 +51,45 @@ def test_methane_rotation(atol=1e-3):
     mol2.readfromstring(xyz_string)
     # rotate 180 degrees around the z axis
     mol2 = rotate_around_axis(mol2, [0.0, 0.0, 0.0], [0.0, 0.0, 1.0], 180)
-    assert rigorous_rmsd(mol1, mol2, reorder='none') < atol
+    #assert rigorous_rmsd(mol1, mol2, reorder='none') < atol
+    assert align_rmsd(mol1, mol2, reorder='none') < atol
 
-    assert rigorous_rmsd(mol1, mol2, reorder='hungarian') < atol
+    #assert rigorous_rmsd(mol1, mol2, reorder='hungarian') < atol
+    assert align_rmsd(mol1, mol2, reorder='hungarian') < atol
+
+def test_tbpd_rotation(atol=1e-3):
+    """Designed to test align_rmsd and show how rigorous_rmsd can fail"""
+    xyz_string = (
+        """
+        Fe 0 0 0
+        O  -0.5  0.866 0
+        O  -0.5 -0.866 0
+        O  1 0 0
+        O  0 0 0.8
+        O  0 0 -1
+        """
+    )
+    mol1 = mol3D()
+    mol1.readfromstring(xyz_string)
+
+    mol2 = mol3D()
+    mol2.copymol3D(mol1)
+
+    #if testing permutations as well, can use the following instead of mol2.copymol3D
+    """
+    atoms = [atom for atom in mol1.atoms]
+    idxs = np.arange(len(atoms))
+    np.random.shuffle(idxs)
+    for idx in idxs:
+    mol2.addAtom(atoms[idx])
+    """
+
+    mol2 = rotate_around_axis(mol2, [0.0, 0.0, 0.0], [1, 0, 0], 180)
+    #to test arbitrary rotations, use the following
+    #mol2 = rotate_around_axis(mol2, [0.0, 0.0, 0.0], 2*np.pi*np.random.rand(3), 180)
+
+    assert rigorous_rmsd(mol1, mol2, reorder='none') < atol
+    assert align_rmsd(mol1, mol2, reorder='hungarian') < atol
 
 
 def test_quaternion_rotate(atol=1e-3):