From a9eeff8a5ce437a2d2eb486fbafc494b8ce5f71b Mon Sep 17 00:00:00 2001 From: aarongarrison <84917438+aarongarrison@users.noreply.github.com> Date: Tue, 18 Jun 2024 17:22:44 -0400 Subject: [PATCH] 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 | 59 +++++++++++++++++++++++ molSimplify/Scripts/rmsd.py | 92 +++++++++++++++++++++++++++++++++--- tests/test_rmsd.py | 44 +++++++++++++++-- 3 files changed, 184 insertions(+), 11 deletions(-) diff --git a/molSimplify/Classes/mol3D.py b/molSimplify/Classes/mol3D.py index 64c6e6a3..f7580119 100644 --- a/molSimplify/Classes/mol3D.py +++ b/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 diff --git a/molSimplify/Scripts/rmsd.py b/molSimplify/Scripts/rmsd.py index 9c7831cb..a57f0747 100644 --- a/molSimplify/Scripts/rmsd.py +++ b/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"]) diff --git a/tests/test_rmsd.py b/tests/test_rmsd.py index 5c417ca6..233ba5f4 100644 --- a/tests/test_rmsd.py +++ b/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):