coords.py 3.62 KB
Newer Older
cmaffeo2's avatar
cmaffeo2 committed
1
2
3
import numpy as np
from scipy.optimize import newton

cmaffeo2's avatar
cmaffeo2 committed
4
5
6
7
8
9
10
11
12
13

def minimizeRmsd(coordsB, coordsA, weights=None, maxIter=100):
    ## Going through many iterations wasn't really needed
    tol = 1
    count = 0

    R = np.eye(3)
    comB = np.zeros([3,])
    cNext = coordsB

14
    while tol > 1e-6:
cmaffeo2's avatar
cmaffeo2 committed
15
16
        q,cB,comA = _minimizeRmsd(cNext,coordsA, weights)
        R = R.dot(quaternionToMatrix3(q))
17
        assert( np.all(np.isreal( R )) )
cmaffeo2's avatar
cmaffeo2 committed
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33

        comB += cB
        cLast = cNext
        cNext = (coordsB-comB).dot(R)

        tol = np.sum(((cNext-cLast)**2)[:]) / np.max(np.shape(coordsB))
        if count > maxIter:
            Exception("Exceeded maxIter (%d)" % maxIter)
        count += 1

    print("%d iterations",count)
    return R, comB, comA


def minimizeRmsd(coordsB, coordsA, weights=None):
    q,comA,comB = _minimizeRmsd(coordsB, coordsA, weights)
34
    assert( np.all(np.isreal( q )) )
cmaffeo2's avatar
cmaffeo2 committed
35
36
37
    return quaternionToMatrix3(q),comA,comB


cmaffeo2's avatar
cmaffeo2 committed
38
## http://onlinelibrary.wiley.com/doi/10.1002/jcc.21439/full
cmaffeo2's avatar
cmaffeo2 committed
39
def _minimizeRmsd(coordsB, coordsA, weights=None):
cmaffeo2's avatar
cmaffeo2 committed
40
41
    A = coordsA
    B = coordsB
cmaffeo2's avatar
cmaffeo2 committed
42
43
44
45
46
47
48
49
50

    shapeA,shapeB = [np.shape(X) for X in (A,B)]
    for s in (shapeA,shapeB):  assert( len(s) == 2 )

    A,B = [X.T if s[1] > s[0] else X for X,s in zip([A,B],(shapeA,shapeB))] # TODO: print warning

    shapeA,shapeB = [np.shape(X) for X in (A,B)]
    assert( shapeA == shapeB )
    for X,s in zip((A,B),(shapeA,shapeB)):
51
        assert( s[1] == 3 and s[0] >= s[1] )
cmaffeo2's avatar
cmaffeo2 committed
52
    
53
54
55
56
57
58
59
    # if weights is None: weights = np.ones(len(A))
    if weights is None:
        comA,comB = [np.mean( X, axis=0 ) for X in (A,B)]
    else:
        assert( len(weights[:]) == len(B) )
        W = np.diag(weights)
        comA,comB = [np.sum( W.dot(X), axis=0 ) / np.sum(W) for X in (A,B)]
cmaffeo2's avatar
cmaffeo2 committed
60

cmaffeo2's avatar
cmaffeo2 committed
61
62
    A = np.array( A-comA )
    B = np.array( B-comB )
63
64
65
66
67

    if weights is None:
        s = A.T.dot(B)
    else:
        s = A.T.dot(W.dot(B))
cmaffeo2's avatar
cmaffeo2 committed
68
69
70
71
    
    sxx,sxy,sxz = s[0,:]
    syx,syy,syz = s[1,:]
    szx,szy,szz = s[2,:]
cmaffeo2's avatar
cmaffeo2 committed
72
73
74
    
    K = [[ sxx+syy+szz, syz-szy, szx-sxz, sxy-syx],
         [syz-szy,  sxx-syy-szz, sxy+syx, sxz+szx],
75
         [szx-sxz, sxy+syx, -sxx+syy-szz, syz+szy],
cmaffeo2's avatar
cmaffeo2 committed
76
77
78
         [sxy-syx, sxz+szx, syz+szy, -sxx-syy+szz]]
    K = np.array(K)

cmaffeo2's avatar
cmaffeo2 committed
79
80
81
82
83
84
85
    # GA = np.trace( A.T.dot(W.dot(A)) )
    # GB = np.trace( B.T.dot(W.dot(B)) )
        
    ## Finding GA/GB can be done more quickly
    # I = np.eye(4)
    # x0 = (GA+GB)*0.5
    # vals = newtoon(lambda x: np.det(K-x*I), x0 = x0)
cmaffeo2's avatar
cmaffeo2 committed
86

cmaffeo2's avatar
cmaffeo2 committed
87
88
89
    vals, vecs = np.linalg.eig(K)
    i = np.argmax(vals)
    q = vecs[:,i]
cmaffeo2's avatar
cmaffeo2 committed
90

cmaffeo2's avatar
cmaffeo2 committed
91
92
93
    # RMSD = np.sqrt( (GA+GB-2*vals[i]) / len(A) )
    # print("CHECK:", K.dot(q)-vals[i]*q )
    return q, comB, comA
cmaffeo2's avatar
cmaffeo2 committed
94

cmaffeo2's avatar
cmaffeo2 committed
95
96
def quaternionToMatrix3(q):
    assert(len(q) == 4)
cmaffeo2's avatar
cmaffeo2 committed
97

cmaffeo2's avatar
cmaffeo2 committed
98
99
100
101
102
103
104
105
106
107
108
109
110
    ## It looks like the wikipedia article I used employed a less common convention for q (see below
    ## http://en.wikipedia.org/wiki/Rotation_formalisms_in_three_dimensions#Rotation_matrix_.E2.86.94_quaternion
    # q1,q2,q3,q4 = q
    # R = [[1-2*(q2*q2 + q3*q3),    2*(q1*q2 - q3*q4),    2*(q1*q3 + q2*q4)],
    #      [  2*(q1*q2 + q3*q4),  1-2*(q1*q1 + q3*q3),    2*(q2*q3 - q1*q4)],
    #      [  2*(q1*q3 - q2*q4),    2*(q1*q4 + q2*q3),  1-2*(q2*q2 + q1*q1)]]

    q0,q1,q2,q3 = q
    R = [[1-2*(q2*q2 + q3*q3),    2*(q1*q2 - q3*q0),    2*(q1*q3 + q2*q0)],
         [  2*(q1*q2 + q3*q0),  1-2*(q1*q1 + q3*q3),    2*(q2*q3 - q1*q0)],
         [  2*(q1*q3 - q2*q0),    2*(q1*q0 + q2*q3),  1-2*(q2*q2 + q1*q1)]]

    return np.array(R)
cmaffeo2's avatar
cmaffeo2 committed
111
112
113
114
115
116
117
118

def rotationAboutAxis(axis,angle, normalizeAxis=True):
    if normalizeAxis: axis = axis / np.linalg.norm(axis)
    angle *= 0.5 * np.pi/180
    cos = np.cos( angle )
    sin = np.sin( angle )
    q = [cos] + [sin*x for x in axis]
    return quaternionToMatrix3(q)