coords.py 6.31 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
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

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

        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)
33
    assert( np.all(np.isreal( q )) )
34
    return quaternion_to_matrix(q),comA,comB
cmaffeo2's avatar
cmaffeo2 committed
35
36


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

    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)):
50
        assert( s[1] == 3 and s[0] >= s[1] )
cmaffeo2's avatar
cmaffeo2 committed
51
    
52
53
54
55
56
57
58
    # 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
59

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

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

cmaffeo2's avatar
cmaffeo2 committed
78
79
80
81
82
83
84
    # 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
85

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

cmaffeo2's avatar
cmaffeo2 committed
90
91
92
    # 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
93

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

cmaffeo2's avatar
cmaffeo2 committed
97
98
99
100
101
102
103
    ## 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)]]

104
    q = q / np.linalg.norm(q)
cmaffeo2's avatar
cmaffeo2 committed
105
106
107
108
109
110
    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
119
120
121
122
123
def quaternion_from_matrix( R ):
    e1 = R[0]
    e2 = R[1]
    e3 = R[2]
    
    # d1 = 0.5 * np.sqrt( 1+R[0,0]+R[1,1]+R[2,2] )
    # d2 = 0.5 * np.sqrt( 1+R[0,0]-R[1,1]-R[2,2] )
    # d2 = 0.5 * np.sqrt( 1+R[0,0]-R[1,1]-R[2,2] )
    # d2 = 0.5 * np.sqrt( 1+R[0,0]-R[1,1]-R[2,2] )

    d1 = 1+R[0,0]+R[1,1]+R[2,2]
    d2 = 1+R[0,0]-R[1,1]-R[2,2]
cmaffeo2's avatar
cmaffeo2 committed
124
125
    d3 = 1-R[0,0]+R[1,1]-R[2,2]
    d4 = 1-R[0,0]-R[1,1]+R[2,2]
126
    
cmaffeo2's avatar
cmaffeo2 committed
127
    maxD = max((d1,d2,d3,d4))
128
    d = 0.5 / np.sqrt(maxD)
cmaffeo2's avatar
cmaffeo2 committed
129

130
131
132
133
134
135
136
137
138
139
    if d1 == maxD:
        return np.array(( 1.0/(4*d),
                          d * (R[2,1]-R[1,2]),
                          d * (R[0,2]-R[2,0]),
                          d * (R[1,0]-R[0,1]) ))
    elif d2 == maxD:
        return np.array(( d * (R[2,1]-R[1,2]),
                          1.0/(4*d),
                          d * (R[0,1]+R[1,0]),
                          d * (R[0,2]+R[2,0]) ))
cmaffeo2's avatar
cmaffeo2 committed
140
141
142
143
144
145
146
147
148
149
    elif d3 == maxD:
        return np.array(( d * (R[0,2]-R[2,0]),
                          d * (R[0,1]+R[1,0]),
                          1.0/(4*d),
                          d * (R[1,2]+R[2,1]) ))
    elif d4 == maxD:
        return np.array(( d * (R[1,0]-R[0,1]),
                          d * (R[0,2]+R[2,0]),
                          d * (R[1,2]+R[2,1]),
                          1.0/(4*d) ))
150

cmaffeo2's avatar
cmaffeo2 committed
151
152
153
154
155
156
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]
157
158
159
160
161
162
163
164
    return quaternion_to_matrix(q)

def readArbdCoords(fname):
    coords = []
    with open(fname) as fh:
        for line in fh:
            coords.append([float(x) for x in line.split()[1:]])
    return np.array(coords)
cmaffeo2's avatar
cmaffeo2 committed
165
166

def readAvgArbdCoords(psf,pdb,dcd,rmsdThreshold=3.5):
cmaffeo2's avatar
cmaffeo2 committed
167
    import MDAnalysis as mda
cmaffeo2's avatar
cmaffeo2 committed
168

cmaffeo2's avatar
cmaffeo2 committed
169
170
    usel = mda.Universe(psf, dcd)
    sel = usel.select_atoms("name D*")
cmaffeo2's avatar
cmaffeo2 committed
171
172

    # r0 = ref.xyz[0,ids,:]
cmaffeo2's avatar
cmaffeo2 committed
173
174
175
176
177
178
179
180
181
182
183
    ts = usel.trajectory[-1]
    r0 = sel.positions
    pos = []
    for t in range(ts.frame-1,-1,-1):
        usel.trajectory[t]
        R,comA,comB = minimizeRmsd(sel.positions,r0)
        r = np.array( [(r-comA).dot(R)+comB for r in sel.positions] )
        rmsd = np.mean( (r-r0)**2 )
        r = np.array( [(r-comA).dot(R)+comB for r in usel.atoms.positions] )
        pos.append( r )
        if rmsd > rmsdThreshold**2:
cmaffeo2's avatar
cmaffeo2 committed
184
185
186
187
188
            break
    t0=t+1
    print( "Averaging coordinates in %s after frame %d" % (dcd, t0) )

    pos = np.mean(pos, axis=0)
cmaffeo2's avatar
cmaffeo2 committed
189
    return pos
cmaffeo2's avatar
cmaffeo2 committed
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204

def unit_quat_conversions():
    for axis in [[0,0,1],[1,1,1],[1,0,0],[-1,-2,0]]:
        for angle in np.linspace(-180,180,10):
            R = rotationAboutAxis(axis, angle)
            R2 = quaternion_to_matrix( quaternion_from_matrix( R ) )
            if not np.all( np.abs(R-R2) < 0.01 ):
                import pdb
                pdb.set_trace()
                quaternion_to_matrix( quaternion_from_matrix( R ) )


if __name__ == "__main__":
    unit_quat_conversions()