1
2
3
4
5
6 import numpy
7 import sys
8
9
10 __doc__="Vector class, including rotation-related functions."
11
13 """
14 Return angles, axis pair that corresponds to rotation matrix m.
15 """
16
17
18 t=0.5*(numpy.trace(m)-1)
19 t=max(-1, t)
20 t=min(1, t)
21 angle=numpy.arccos(t)
22 if angle<1e-15:
23
24 return 0.0, Vector(1,0,0)
25 elif angle<numpy.pi:
26
27 x=m[2,1]-m[1,2]
28 y=m[0,2]-m[2,0]
29 z=m[1,0]-m[0,1]
30 axis=Vector(x,y,z)
31 axis.normalize()
32 return angle, axis
33 else:
34
35 m00=m[0,0]
36 m11=m[1,1]
37 m22=m[2,2]
38 if m00>m11 and m00>m22:
39 x=numpy.sqrt(m00-m11-m22+0.5)
40 y=m[0,1]/(2*x)
41 z=m[0,2]/(2*x)
42 elif m11>m00 and m11>m22:
43 y=numpy.sqrt(m11-m00-m22+0.5)
44 x=m[0,1]/(2*y)
45 z=m[1,2]/(2*y)
46 else:
47 z=numpy.sqrt(m22-m00-m11+0.5)
48 x=m[0,2]/(2*z)
49 y=m[1,2]/(2*z)
50 axis=Vector(x,y,z)
51 axis.normalize()
52 return numpy.pi, axis
53
54
56 """
57 Returns the vector between a point and
58 the closest point on a line (ie. the perpendicular
59 projection of the point on the line).
60
61 @type line: L{Vector}
62 @param line: vector defining a line
63
64 @type point: L{Vector}
65 @param point: vector defining the point
66 """
67 line=line.normalized()
68 np=point.norm()
69 angle=line.angle(point)
70 return point-line**(np*numpy.cos(angle))
71
72
74 """
75 Calculate a left multiplying rotation matrix that rotates
76 theta rad around vector.
77
78 Example:
79
80 >>> m=rotaxis(pi, Vector(1,0,0))
81 >>> rotated_vector=any_vector.left_multiply(m)
82
83 @type theta: float
84 @param theta: the rotation angle
85
86
87 @type vector: L{Vector}
88 @param vector: the rotation axis
89
90 @return: The rotation matrix, a 3x3 Numeric array.
91 """
92 vector=vector.copy()
93 vector.normalize()
94 c=numpy.cos(theta)
95 s=numpy.sin(theta)
96 t=1-c
97 x,y,z=vector.get_array()
98 rot=numpy.zeros((3,3))
99
100 rot[0,0]=t*x*x+c
101 rot[0,1]=t*x*y-s*z
102 rot[0,2]=t*x*z+s*y
103
104 rot[1,0]=t*x*y+s*z
105 rot[1,1]=t*y*y+c
106 rot[1,2]=t*y*z-s*x
107
108 rot[2,0]=t*x*z-s*y
109 rot[2,1]=t*y*z+s*x
110 rot[2,2]=t*z*z+c
111 return rot
112
113 rotaxis=rotaxis2m
114
116 """
117 Return a (left multiplying) matrix that mirrors p onto q.
118
119 Example:
120 >>> mirror=refmat(p,q)
121 >>> qq=p.left_multiply(mirror)
122 >>> print q, qq # q and qq should be the same
123
124 @type p,q: L{Vector}
125 @return: The mirror operation, a 3x3 Numeric array.
126 """
127 p.normalize()
128 q.normalize()
129 if (p-q).norm()<1e-5:
130 return numpy.identity(3)
131 pq=p-q
132 pq.normalize()
133 b=pq.get_array()
134 b.shape=(3, 1)
135 i=numpy.identity(3)
136 ref=i-2*numpy.dot(b, numpy.transpose(b))
137 return ref
138
140 """
141 Return a (left multiplying) matrix that rotates p onto q.
142
143 Example:
144 >>> r=rotmat(p,q)
145 >>> print q, p.left_multiply(r)
146
147 @param p: moving vector
148 @type p: L{Vector}
149
150 @param q: fixed vector
151 @type q: L{Vector}
152
153 @return: rotation matrix that rotates p onto q
154 @rtype: 3x3 Numeric array
155 """
156 rot=numpy.dot(refmat(q, -p), refmat(p, -p))
157 return rot
158
160 """
161 Calculate the angle between 3 vectors
162 representing 3 connected points.
163
164 @param v1, v2, v3: the tree points that define the angle
165 @type v1, v2, v3: L{Vector}
166
167 @return: angle
168 @rtype: float
169 """
170 v1=v1-v2
171 v3=v3-v2
172 return v1.angle(v3)
173
175 """
176 Calculate the dihedral angle between 4 vectors
177 representing 4 connected points. The angle is in
178 ]-pi, pi].
179
180 @param v1, v2, v3, v4: the four points that define the dihedral angle
181 @type v1, v2, v3, v4: L{Vector}
182 """
183 ab=v1-v2
184 cb=v3-v2
185 db=v4-v3
186 u=ab**cb
187 v=db**cb
188 w=u**v
189 angle=u.angle(v)
190
191 try:
192 if cb.angle(w)>0.001:
193 angle=-angle
194 except ZeroDivisionError:
195
196 pass
197 return angle
198
200 "3D vector"
201
203 if y is None and z is None:
204
205 if len(x)!=3:
206 raise "Vector: x is not a list/tuple/array of 3 numbers"
207 self._ar=numpy.array(x, 'd')
208 else:
209
210 self._ar=numpy.array((x, y, z), 'd')
211
213 x,y,z=self._ar
214 return "<Vector %.2f, %.2f, %.2f>" % (x,y,z)
215
217 "Return Vector(-x, -y, -z)"
218 a=-self._ar
219 return Vector(a)
220
222 "Return Vector+other Vector or scalar"
223 if isinstance(other, Vector):
224 a=self._ar+other._ar
225 else:
226 a=self._ar+numpy.array(other)
227 return Vector(a)
228
230 "Return Vector-other Vector or scalar"
231 if isinstance(other, Vector):
232 a=self._ar-other._ar
233 else:
234 a=self._ar-numpy.array(other)
235 return Vector(a)
236
238 "Return Vector.Vector (dot product)"
239 return sum(self._ar*other._ar)
240
242 "Return Vector(coords/a)"
243 a=self._ar/numpy.array(x)
244 return Vector(a)
245
247 "Return VectorxVector (cross product) or Vectorxscalar"
248 if isinstance(other, Vector):
249 a,b,c=self._ar
250 d,e,f=other._ar
251 c1=numpy.linalg.det(numpy.array(((b,c), (e,f))))
252 c2=-numpy.linalg.det(numpy.array(((a,c), (d,f))))
253 c3=numpy.linalg.det(numpy.array(((a,b), (d,e))))
254 return Vector(c1,c2,c3)
255 else:
256 a=self._ar*numpy.array(other)
257 return Vector(a)
258
261
264
266 "Return vector norm"
267 return numpy.sqrt(sum(self._ar*self._ar))
268
270 "Return square of vector norm"
271 return abs(sum(self._ar*self._ar))
272
274 "Normalize the Vector"
275 self._ar=self._ar/self.norm()
276
278 "Return a normalized copy of the Vector"
279 v=self.copy()
280 v.normalize()
281 return v
282
284 "Return angle between two vectors"
285 n1=self.norm()
286 n2=other.norm()
287 c=(self*other)/(n1*n2)
288
289 c=min(c,1)
290 c=max(-1,c)
291 return numpy.arccos(c)
292
294 "Return (a copy of) the array of coordinates"
295 return numpy.array(self._ar)
296
298 "Return Vector=Matrix x Vector"
299 a=numpy.dot(matrix, self._ar)
300 return Vector(a)
301
303 "Return Vector=Vector x Matrix"
304 a=numpy.dot(self._ar, matrix)
305 return Vector(a)
306
308 "Return a deep copy of the Vector"
309 return Vector(self._ar)
310
311 if __name__=="__main__":
312
313 from numpy.random import random
314
315 v1=Vector(0,0,1)
316 v2=Vector(0,0,0)
317 v3=Vector(0,1,0)
318 v4=Vector(1,1,0)
319
320 v4.normalize()
321
322 print v4
323
324 print calc_angle(v1, v2, v3)
325 dih=calc_dihedral(v1, v2, v3, v4)
326
327 assert(dih>0)
328 print "DIHEDRAL ", dih
329
330 ref=refmat(v1, v3)
331 rot=rotmat(v1, v3)
332
333 print v3
334 print v1.left_multiply(ref)
335 print v1.left_multiply(rot)
336 print v1.right_multiply(numpy.transpose(rot))
337
338
339 print v1-v2
340 print v1-1
341 print v1+(1,2,3)
342
343 print v1+v2
344 print v1+3
345 print v1-(1,2,3)
346
347 print v1*v2
348
349 print v1/2
350 print v1/(1,2,3)
351
352 print v1**v2
353 print v1**2
354 print v1**(1,2,3)
355
356 print v1.norm()
357
358 print v1.normsq()
359
360 v1[2]=10
361 print v1
362
363 print v1[2]
364
365 print numpy.array(v1)
366
367 print "ROT"
368
369 angle=random()*numpy.pi
370 axis=Vector(random(3)-random(3))
371 axis.normalize()
372
373 m=rotaxis(angle, axis)
374
375 cangle, caxis=m2rotaxis(m)
376
377 print angle-cangle
378 print axis-caxis
379 print
380