Merge pull request #27 from DeepRank/align

Align

Author: NicoRenaud

pdb2sql/StructureSimilarity.py

@@ -2,6 +2,8 @@
 import numpy as np
 from .pdb2sqlcore import pdb2sql
 from .interface import interface
+from .superpose import get_trans_vect, get_rotation_matrix, superpose_selection
+
 from . import transform
 import os
 import pickle
@@ -134,27 +136,8 @@ def compute_lrmsd_fast(self, lzone=None, method='svd', check=True):
             xyz_ref_long, xyz_ref_short = self.get_xyz_zone_backbone(
                 self.ref, resData, return_not_in_zone=True)
 
-        # get the translation so that both A chains are centered
-        tr_decoy = self.get_trans_vect(xyz_decoy_long)
-        tr_ref = self.get_trans_vect(xyz_ref_long)
-
-        # translate everything for 1
-        xyz_decoy_short += tr_decoy
-        xyz_decoy_long += tr_decoy
-
-        # translate everuthing for 2
-        xyz_ref_short += tr_ref
-        xyz_ref_long += tr_ref
-
-        # get the ideal rotation matrix
-        # to superimpose the A chains
-        U = self.get_rotation_matrix(
-            xyz_decoy_long, xyz_ref_long, method=method)
-
-        # rotate the entire fragment
-
-        xyz_decoy_short = transform.rotate(
-            xyz_decoy_short, U, center=self.origin)
+        xyz_decoy_short = superpose_selection(
+                xyz_decoy_short, xyz_decoy_long, xyz_ref_long, method)
 
         # compute the RMSD
         return self.get_rmsd(xyz_decoy_short, xyz_ref_short)
@@ -281,24 +264,9 @@ def compute_irmsd_fast(
             xyz_contact_decoy = self.get_xyz_zone_backbone(self.decoy, resData)
             xyz_contact_ref = self.get_xyz_zone_backbone(self.ref, resData)
 
-        # get the translation so that both A chains are centered
-        tr_decoy = self.get_trans_vect(xyz_contact_decoy)
-        tr_ref = self.get_trans_vect(xyz_contact_ref)
-
-        # translate everything
-        xyz_contact_decoy += tr_decoy
-        xyz_contact_ref += tr_ref
-
-        # get the ideal rotation matrix
-        # to superimpose the A chains
-        U = self.get_rotation_matrix(
-            xyz_contact_decoy,
-            xyz_contact_ref,
-            method=method)
-
-        # rotate the entire fragment
-        xyz_contact_decoy = transform.rotate(
-            xyz_contact_decoy, U, center=self.origin)
+        # superpose the fragments
+        xyz_contact_decoy = superpose_selection(xyz_contact_decoy,
+            xyz_contact_decoy, xyz_contact_ref, method)
 
         # return the RMSD
         return self.get_rmsd(xyz_contact_decoy, xyz_contact_ref)
@@ -559,8 +527,8 @@ def compute_lrmsd_pdb2sql(self, exportpath=None, method='svd'):
             xyz_ref_short = xyz_ref_A
 
         # get the translation so that both A chains are centered
-        tr_decoy = self.get_trans_vect(xyz_decoy_long)
-        tr_ref = self.get_trans_vect(xyz_ref_long)
+        tr_decoy = get_trans_vect(xyz_decoy_long)
+        tr_ref = get_trans_vect(xyz_ref_long)
 
         # translate everything for 1
         xyz_decoy_short += tr_decoy
@@ -572,7 +540,7 @@ def compute_lrmsd_pdb2sql(self, exportpath=None, method='svd'):
 
         # get the ideal rotation matrix
         # to superimpose the A chains
-        U = self.get_rotation_matrix(
+        U = get_rotation_matrix(
             xyz_decoy_long, xyz_ref_long, method=method)
 
         # rotate the entire fragment
@@ -756,23 +724,23 @@ def compute_irmsd_pdb2sql(
                 'Error in i-rmsd: only one chain represented in one chain')
 
         # get the translation so that both A chains are centered
-        tr_decoy = self.get_trans_vect(xyz_contact_decoy)
-        tr_ref = self.get_trans_vect(xyz_contact_ref)
+        tr_decoy = get_trans_vect(xyz_contact_decoy)
+        tr_ref = get_trans_vect(xyz_contact_ref)
 
         # translate everything
         xyz_contact_decoy += tr_decoy
         xyz_contact_ref += tr_ref
 
         # get the ideql rotation matrix
         # to superimpose the A chains
-        U = self.get_rotation_matrix(
-            xyz_contact_decoy,
-            xyz_contact_ref,
-            method=method)
+        rot_mat = get_rotation_matrix(
+                    xyz_contact_decoy,
+                    xyz_contact_ref,
+                    method=method)
 
         # rotate the entire fragment
         xyz_contact_decoy = transform.rotate(
-            xyz_contact_decoy, U, center=self.origin)
+            xyz_contact_decoy, rot_mat, center=self.origin)
 
         # compute the RMSD
         irmsd = self.get_rmsd(xyz_contact_decoy, xyz_contact_ref)
@@ -1185,7 +1153,7 @@ def compute_clashes(pdb):
 
     ##########################################################################
     #
-    #   ROUTINES TO ACTUALY ALIGN THE MOLECULES
+    #   ROUTINES TO ACTUALY SUPERPOSE THE MOLECULES
     #
     ##########################################################################
 
@@ -1205,170 +1173,3 @@ def get_rmsd(P, Q):
         """
         n = len(P)
         return round(np.sqrt(1. / n * np.sum((P - Q)**2)), 3)
-
-    # compute the translation vector to center a set of points
-    @staticmethod
-    def get_trans_vect(P):
-        """Get the translationv vector to the origin.
-
-        Args:
-            P (np.array(nx3)): position of the points in the molecule
-
-        Returns:
-            float: minus mean value of the xyz columns
-        """
-        return -np.mean(P, 0)
-
-    # main switch for the rotation matrix
-    # add new methods here if necessary
-    def get_rotation_matrix(self, P, Q, method='svd'):
-
-        # get the matrix with Kabsh method
-        if method.lower() == 'svd':
-            mat = self.get_rotation_matrix_Kabsh(P, Q)
-
-        # or with the quaternion method
-        elif method.lower() == 'quaternion':
-            mat = self.get_rotation_matrix_quaternion(P, Q)
-
-        else:
-            raise ValueError(
-                f'{method} is not a valid method for rmsd alignement. '
-                f'Options are svd or quaternions')
-
-        return mat
-
-    # get the rotation matrix via a SVD
-    # decomposition of the correlation matrix
-    @staticmethod
-    def get_rotation_matrix_Kabsh(P, Q):
-        """Get the rotation matrix to aligh two point clouds.
-
-        The method is based on th Kabsh approach
-        https://cnx.org/contents/HV-RsdwL@23/Molecular-Distance-Measures
-
-        Args:
-            P (np.array): xyz of the first point cloud
-            Q (np.array): xyz of the second point cloud
-
-        Returns:
-            np.array: rotation matrix
-
-        Raises:
-            ValueError: matrix have different sizes
-        """
-        pshape = P.shape
-        qshape = Q.shape
-
-        if pshape[0] == qshape[0]:
-            npts = pshape[0]
-        else:
-            raise ValueError("Matrix don't have the same number of points",
-                             P.shape, Q.shape)
-
-        p0, q0 = np.abs(np.mean(P, 0)), np.abs(np.mean(Q, 0))
-        eps = 1E-6
-        if any(p0 > eps) or any(q0 > eps):
-            raise ValueError('You must center the fragment first', p0, q0)
-
-        # form the covariance matrix
-        A = np.dot(P.T, Q) / npts
-
-        # SVD the matrix
-        V, _, W = np.linalg.svd(A)
-
-        # the W matrix returned here is
-        # already its transpose
-        # https://docs.scipy.org/doc/numpy-1.13.0/reference/generated/numpy.linalg.svd.html
-        W = W.T
-
-        # determinant
-        d = np.linalg.det(np.dot(W, V.T))
-
-        # form the U matrix
-        Id = np.eye(3)
-        if d < 0:
-            Id[2, 2] = -1
-
-        U = np.dot(W, np.dot(Id, V.T))
-
-        return U
-
-    # get the rotation amtrix via the quaternion approach
-    # doesn't work great so far
-    @staticmethod
-    def get_rotation_matrix_quaternion(P, Q):
-        """Get the rotation matrix to aligh two point clouds.
-
-        The method is based on the quaternion approach
-        http://www.ams.stonybrook.edu/~coutsias/papers/rmsd17.pdf
-
-        Args:
-            P (np.array): xyz of the first point cloud
-            Q (np.array): xyz of the second point cloud
-
-        Returns:
-            np.array: rotation matrix
-
-        Raises:
-            ValueError: matrix have different sizes
-        """
-        pshape = P.shape
-        qshape = Q.shape
-
-        if pshape[0] != qshape[0]:
-            raise ValueError("Matrix don't have the same number of points",
-                             P.shape, Q.shape)
-
-        p0, q0 = np.abs(np.mean(P, 0)), np.abs(np.mean(Q, 0))
-        eps = 1E-6
-        if any(p0 > eps) or any(q0 > eps):
-            raise ValueError('You must center the fragment first', p0, q0)
-
-        # form the correlation matrix
-        R = np.dot(P.T, Q)
-
-        # form the F matrix (eq. 10 of ref[1])
-        F = np.zeros((4, 4))
-
-        F[0, 0] = np.trace(R)
-        F[0, 1] = R[1, 2] - R[2, 1]
-        F[0, 2] = R[2, 0] - R[0, 2]
-        F[0, 3] = R[0, 1] - R[1, 0]
-
-        F[1, 0] = R[1, 2] - R[2, 1]
-        F[1, 1] = R[0, 0] - R[1, 1] - R[2, 2]
-        F[1, 2] = R[0, 1] + R[1, 0]
-        F[1, 3] = R[0, 2] + R[2, 0]
-
-        F[2, 0] = R[2, 0] - R[0, 2]
-        F[2, 1] = R[0, 1] + R[1, 0]
-        F[2, 2] = -R[0, 0] + R[1, 1] - R[2, 2]
-        F[2, 3] = R[1, 2] + R[2, 1]
-
-        F[3, 0] = R[0, 1] - R[1, 0]
-        F[3, 1] = R[0, 2] + R[2, 0]
-        F[3, 2] = R[1, 2] + R[2, 1]
-        F[3, 3] = -R[0, 0] - R[1, 1] + R[2, 2]
-
-        # diagonalize it
-        l, U = np.linalg.eig(F)
-
-        # extract the eigenvect of the highest eigenvalues
-        indmax = np.argmax(l)
-        q0, q1, q2, q3 = U[:, indmax]
-
-        # form the rotation matrix (eq. 33 ref[1])
-        U = np.zeros((3, 3))
-
-        U[0, 0] = q0**2 + q1**2 - q2**2 - q3**2
-        U[0, 1] = 2 * (q1 * q2 - q0 * q3)
-        U[0, 2] = 2 * (q1 * q3 + q0 * q2)
-        U[1, 0] = 2 * (q1 * q2 + q0 * q3)
-        U[1, 1] = q0**2 - q1**2 + q2**2 - q3**2
-        U[1, 2] = 2 * (q2 * q3 - q0 * q1)
-        U[2, 0] = 2 * (q1 * q3 - q0 * q2)
-        U[2, 1] = 2 * (q2 * q3 + q0 * q1)
-        U[2, 2] = q0**2 - q1**2 - q2**2 + q3**2
-
-        return U