2323
2424import math
2525import random
26- import copy
2726
2827RSMALL4 = 0.0001
2928
3029#------------------------------------------------------------------------------
31- def get_null ():
30+ def get_null () -> list [ float ] :
3231 return [0.0 ,0.0 ,0.0 ]
3332
3433#------------------------------------------------------------------------------
35- def get_identity ():
34+ def get_identity () -> list [ list [ float ]] :
3635 return [[1.0 ,0.0 ,0.0 ],[0.0 ,1.0 ,0.0 ],[0.0 ,0.0 ,1.0 ]]
3736
3837#------------------------------------------------------------------------------
39- def distance_sq (v1 , v2 ) :
38+ def distance_sq (v1 : list [ float ] , v2 : list [ float ]) -> float :
4039 d0 = v2 [0 ] - v1 [0 ]
4140 d1 = v2 [1 ] - v1 [1 ]
4241 d2 = v2 [2 ] - v1 [2 ]
4342 return (d0 * d0 ) + (d1 * d1 ) + (d2 * d2 )
4443
4544#------------------------------------------------------------------------------
46- def distance (v1 , v2 ) :
45+ def distance (v1 : list [ float ] , v2 : list [ float ]) -> float :
4746 d0 = v2 [0 ] - v1 [0 ]
4847 d1 = v2 [1 ] - v1 [1 ]
4948 d2 = v2 [2 ] - v1 [2 ]
5049 return math .sqrt ((d0 * d0 ) + (d1 * d1 ) + (d2 * d2 ))
5150
5251#------------------------------------------------------------------------------
53- def length (v ) :
52+ def length (v : list [ float ]) -> float :
5453 return math .sqrt (v [0 ]* v [0 ] + v [1 ]* v [1 ] + v [2 ]* v [2 ])
5554
5655#------------------------------------------------------------------------------
57- def random_displacement (v ,radius ):
58- r_vect = lambda r = random .random :[r ()- 0.5 ,r ()- 0.5 ,r ()- 0.5 ]
56+ def _random_vector () -> list [float ]:
57+ r = random .random
58+ return [r ()- 0.5 ,r ()- 0.5 ,r ()- 0.5 ]
59+
60+ def random_displacement (v : list [float ], radius : float ) -> list [float ]:
5961 while 1 :
60- vect = r_vect ()
62+ vect = _random_vector ()
6163 v_len = length (vect )
6264 if (v_len <= 0.5 ):
6365 break ;
@@ -68,92 +70,85 @@ def random_displacement(v,radius):
6870 return v
6971
7072#------------------------------------------------------------------------------
71- def random_sphere (v ,radius ):
72- r_vect = lambda r = random .random :[r ()- 0.5 ,r ()- 0.5 ,r ()- 0.5 ]
73+ def random_sphere (v : list [float ], radius : float ) -> list [float ]:
7374 while 1 :
74- vect = r_vect ()
75+ vect = _random_vector ()
7576 v_len = length (vect )
7677 if (v_len <= 0.5 ) and (v_len != 0.0 ):
7778 break ;
7879 return add (v ,scale ([vect [0 ], vect [1 ], vect [2 ]],2 * radius / v_len ))
7980
8081#------------------------------------------------------------------------------
81- def random_vector ():
82- r_vect = lambda r = random .random :[r ()- 0.5 ,r ()- 0.5 ,r ()- 0.5 ]
82+ def random_vector () -> list [float ]:
8383 while 1 :
84- vect = r_vect ()
84+ vect = _random_vector ()
8585 if length (vect )<= 0.5 :
8686 break ;
8787 return scale ([vect [0 ], vect [1 ], vect [2 ]],2.0 )
8888
8989#------------------------------------------------------------------------------
90- def add (v1 , v2 ) :
90+ def add (v1 : list [ float ], v2 : list [ float ]) -> list [ float ] :
9191 return [v1 [0 ]+ v2 [0 ],v1 [1 ]+ v2 [1 ],v1 [2 ]+ v2 [2 ]]
9292
9393#------------------------------------------------------------------------------
94- def average (v1 , v2 ) :
94+ def average (v1 : list [ float ], v2 : list [ float ]) -> list [ float ] :
9595 return [(v1 [0 ]+ v2 [0 ])/ 2.0 ,(v1 [1 ]+ v2 [1 ])/ 2.0 ,(v1 [2 ]+ v2 [2 ])/ 2.0 ]
9696
9797#------------------------------------------------------------------------------
98- def scale (v , factor ) :
98+ def scale (v : list [ float ], factor : float ) -> list [ float ] :
9999 return [v [0 ]* factor ,v [1 ]* factor ,v [2 ]* factor ]
100100
101101#------------------------------------------------------------------------------
102- def negate (v ) :
102+ def negate (v : list [ float ]) -> list [ float ] :
103103 return [- v [0 ],- v [1 ],- v [2 ]]
104104
105105#------------------------------------------------------------------------------
106- def sub (v1 ,v2 ):
106+ def reverse (v : list [float ]) -> list [float ]:
107+ return [ - v [0 ], - v [1 ], - v [2 ] ]
108+
109+ #------------------------------------------------------------------------------
110+ def sub (v1 : list [float ], v2 : list [float ]) -> list [float ]:
107111 return [v1 [0 ]- v2 [0 ],v1 [1 ]- v2 [1 ],v1 [2 ]- v2 [2 ]]
108112
109113#------------------------------------------------------------------------------
110- def dot_product (v1 , v2 ) :
114+ def dot_product (v1 : list [ float ], v2 : list [ float ]) -> float :
111115 return v1 [0 ]* v2 [0 ]+ v1 [1 ]* v2 [1 ]+ v1 [2 ]* v2 [2 ]
112116
113117#------------------------------------------------------------------------------
114- def cross_product (v1 , v2 ) :
118+ def cross_product (v1 : list [ float ], v2 : list [ float ]) -> list [ float ] :
115119 return [(v1 [1 ]* v2 [2 ]) - (v1 [2 ]* v2 [1 ]),
116120 (v1 [2 ]* v2 [0 ]) - (v1 [0 ]* v2 [2 ]),
117121 (v1 [0 ]* v2 [1 ]) - (v1 [1 ]* v2 [0 ])]
118122
119123#------------------------------------------------------------------------------
120- def transform (m , v ) :
124+ def transform (m : list [ list [ float ]], v : list [ float ]) -> list [ float ] :
121125 return [m [0 ][0 ]* v [0 ] + m [0 ][1 ]* v [1 ] + m [0 ][2 ]* v [2 ],
122126 m [1 ][0 ]* v [0 ] + m [1 ][1 ]* v [1 ] + m [1 ][2 ]* v [2 ],
123127 m [2 ][0 ]* v [0 ] + m [2 ][1 ]* v [1 ] + m [2 ][2 ]* v [2 ]]
124128
125129#------------------------------------------------------------------------------
126- def inverse_transform (m , v ) :
130+ def inverse_transform (m : list [ list [ float ]], v : list [ float ]) -> list [ float ] :
127131 return [m [0 ][0 ]* v [0 ] + m [1 ][0 ]* v [1 ] + m [2 ][0 ]* v [2 ],
128132 m [0 ][1 ]* v [0 ] + m [1 ][1 ]* v [1 ] + m [2 ][1 ]* v [2 ],
129133 m [0 ][2 ]* v [0 ] + m [1 ][2 ]* v [1 ] + m [2 ][2 ]* v [2 ]]
130134
131135#------------------------------------------------------------------------------
132- def multiply (m1 ,m2 ): # HAVEN'T YET VERIFIED THAT THIS CONFORMS TO STANDARD DEFT
133- return [[m1 [0 ][0 ]* m2 [0 ][0 ] + m1 [0 ][1 ]* m2 [1 ][0 ] + m1 [0 ][2 ]* m2 [2 ][0 ],
134- m1 [1 ][0 ]* m2 [0 ][0 ] + m1 [1 ][1 ]* m2 [1 ][0 ] + m1 [1 ][2 ]* m2 [2 ][0 ],
135- m1 [2 ][0 ]* m2 [0 ][0 ] + m1 [2 ][1 ]* m2 [1 ][0 ] + m1 [2 ][2 ]* m2 [2 ][0 ]],
136- [m1 [0 ][0 ]* m2 [0 ][1 ] + m1 [0 ][1 ]* m2 [1 ][1 ] + m1 [0 ][2 ]* m2 [2 ][1 ],
137- m1 [1 ][0 ]* m2 [0 ][1 ] + m1 [1 ][1 ]* m2 [1 ][1 ] + m1 [1 ][2 ]* m2 [2 ][1 ],
138- m1 [2 ][0 ]* m2 [0 ][1 ] + m1 [2 ][1 ]* m2 [1 ][1 ] + m1 [2 ][2 ]* m2 [2 ][1 ]],
139- [m1 [0 ][0 ]* m2 [0 ][2 ] + m1 [0 ][1 ]* m2 [1 ][2 ] + m1 [0 ][2 ]* m2 [2 ][2 ],
140- m1 [1 ][0 ]* m2 [0 ][2 ] + m1 [1 ][1 ]* m2 [1 ][2 ] + m1 [1 ][2 ]* m2 [2 ][2 ],
141- m1 [2 ][0 ]* m2 [0 ][2 ] + m1 [2 ][1 ]* m2 [1 ][2 ] + m1 [2 ][2 ]* m2 [2 ][2 ]]]
136+ def multiply (m1 : list [list [float ]], m2 : list [list [float ]]) -> list [list [float ]]:
137+ # HAVEN'T YET VERIFIED THAT THIS CONFORMS TO STANDARD DEFT
138+ # upd: no, it's not(fixed)
142139
143- #------------------------------------------------------------------------------
144- def transpose (m1 ):
145- return [[m1 [0 ][0 ],
146- m1 [1 ][0 ],
147- m1 [2 ][0 ]],
148- [m1 [0 ][1 ],
149- m1 [1 ][1 ],
150- m1 [2 ][1 ]],
151- [m1 [0 ][2 ],
152- m1 [1 ][2 ],
153- m1 [2 ][2 ]]]
140+ return [[m1 [0 ][0 ]* m2 [0 ][0 ] + m1 [0 ][1 ]* m2 [1 ][0 ] + m1 [0 ][2 ]* m2 [2 ][0 ],
141+ m1 [0 ][0 ]* m2 [0 ][1 ] + m1 [0 ][1 ]* m2 [1 ][1 ] + m1 [0 ][2 ]* m2 [2 ][1 ],
142+ m1 [0 ][0 ]* m2 [0 ][2 ] + m1 [0 ][1 ]* m2 [1 ][2 ] + m1 [0 ][2 ]* m2 [2 ][2 ]],
143+ [m1 [1 ][0 ]* m2 [0 ][0 ] + m1 [1 ][1 ]* m2 [1 ][0 ] + m1 [1 ][2 ]* m2 [2 ][0 ],
144+ m1 [1 ][0 ]* m2 [0 ][1 ] + m1 [1 ][1 ]* m2 [1 ][1 ] + m1 [1 ][2 ]* m2 [2 ][1 ],
145+ m1 [1 ][0 ]* m2 [0 ][2 ] + m1 [1 ][1 ]* m2 [1 ][2 ] + m1 [1 ][2 ]* m2 [2 ][2 ]],
146+ [m1 [2 ][0 ]* m2 [0 ][0 ] + m1 [2 ][1 ]* m2 [1 ][0 ] + m1 [2 ][2 ]* m2 [2 ][0 ],
147+ m1 [2 ][0 ]* m2 [0 ][1 ] + m1 [2 ][1 ]* m2 [1 ][1 ] + m1 [2 ][2 ]* m2 [2 ][1 ],
148+ m1 [2 ][0 ]* m2 [0 ][2 ] + m1 [2 ][1 ]* m2 [1 ][2 ] + m1 [2 ][2 ]* m2 [2 ][2 ]]]
154149
155150#------------------------------------------------------------------------------
156- def get_system2 (x , y ) :
151+ def get_system2 (x : list [ float ], y : list [ float ]) -> list [ list [ float ]] :
157152 z = cross_product (x ,y )
158153 z = normalize (z )
159154 y = cross_product (z ,x );
@@ -162,24 +157,24 @@ def get_system2(x,y):
162157 return [x ,y ,z ]
163158
164159#------------------------------------------------------------------------------
165- def scale_system (s , factor ) :
160+ def scale_system (s : list [ list [ float ]], factor : float ) -> list [ list [ float ]] :
166161 r = []
167162 for a in s :
168163 r .append ([a [0 ]* factor ,a [1 ]* factor ,a [2 ]* factor ])
169164 return r
170165
171166#------------------------------------------------------------------------------
172- def transpose (m ) :
167+ def transpose (m : list [ list [ float ]]) -> list [ list [ float ]] :
173168 return [[m [0 ][0 ], m [1 ][0 ], m [2 ][0 ]],
174169 [m [0 ][1 ], m [1 ][1 ], m [2 ][1 ]],
175170 [m [0 ][2 ], m [1 ][2 ], m [2 ][2 ]]]
176171
177172#------------------------------------------------------------------------------
178- def transform_about_point (m , v , p ) :
173+ def transform_about_point (m : list [ list [ float ]], v : list [ float ], p : list [ float ]) -> list [ float ] :
179174 return add (transform (m ,sub (v ,p )),p )
180175
181176#------------------------------------------------------------------------------
182- def get_angle (v1 , v2 ) : # v1,v2 must be unit vectors
177+ def get_angle (v1 : list [ float ], v2 : list [ float ]) -> float : # v1,v2 must be unit vectors
183178 denom = (math .sqrt (((v1 [0 ]* v1 [0 ]) + (v1 [1 ]* v1 [1 ]) + (v1 [2 ]* v1 [2 ]))) *
184179 math .sqrt (((v2 [0 ]* v2 [0 ]) + (v2 [1 ]* v2 [1 ]) + (v2 [2 ]* v2 [2 ]))))
185180 if denom > 1e-10 :
@@ -190,7 +185,7 @@ def get_angle(v1,v2): # v1,v2 must be unit vectors
190185 return result
191186
192187#------------------------------------------------------------------------------
193- def get_angle_formed_by (p1 , p2 , p3 ) : # angle formed by three positions in space
188+ def get_angle_formed_by (p1 : list [ float ], p2 : list [ float ], p3 : list [ float ]) -> float : # angle formed by three positions in space
194189
195190 # based on code submitted by Paul Sherwood
196191 r1 = distance (p1 ,p2 )
@@ -207,37 +202,33 @@ def get_angle_formed_by(p1,p2,p3): # angle formed by three positions in space
207202 return theta ;
208203
209204#------------------------------------------------------------------------------
210- def project (v , n ) :
205+ def project (v : list [ float ], n : list [ float ]) -> list [ float ] :
211206 dot = v [0 ]* n [0 ] + v [1 ]* n [1 ] + v [2 ]* n [2 ]
212207 return [ dot * n [0 ], dot * n [1 ], dot * n [2 ] ]
213208
214209#------------------------------------------------------------------------------
215- def remove_component (v , n ) :
210+ def remove_component (v : list [ float ] , n : list [ float ]) -> list [ float ] :
216211 dot = v [0 ]* n [0 ] + v [1 ]* n [1 ] + v [2 ]* n [2 ]
217212 return [v [0 ] - dot * n [0 ], v [1 ] - dot * n [1 ], v [2 ] - dot * n [2 ]]
218213
219214#------------------------------------------------------------------------------
220- def normalize (v ) :
215+ def normalize (v : list [ float ]) -> list [ float ] :
221216 vlen = math .sqrt ((v [0 ]* v [0 ]) + (v [1 ]* v [1 ]) + (v [2 ]* v [2 ]))
222217 if vlen > RSMALL4 :
223218 return [v [0 ]/ vlen ,v [1 ]/ vlen ,v [2 ]/ vlen ]
224219 else :
225220 return get_null ()
226221
227222#------------------------------------------------------------------------------
228- def reverse (v ):
229- return [ - v [0 ], - v [1 ], - v [2 ] ]
230-
231- #------------------------------------------------------------------------------
232- def normalize_failsafe (v ):
223+ def normalize_failsafe (v : list [float ]) -> list [float ]:
233224 vlen = math .sqrt ((v [0 ]* v [0 ]) + (v [1 ]* v [1 ]) + (v [2 ]* v [2 ]))
234225 if vlen > RSMALL4 :
235226 return [v [0 ]/ vlen ,v [1 ]/ vlen ,v [2 ]/ vlen ]
236227 else :
237228 return [1.0 ,0.0 ,0.0 ]
238229
239230#------------------------------------------------------------------------------
240- def rotation_matrix (angle , axis ) :
231+ def rotation_matrix (angle : float , axis : list [ float ]) -> list [ list [ float ]] :
241232
242233 x = axis [0 ]
243234 y = axis [1 ]
@@ -271,7 +262,7 @@ def rotation_matrix(angle,axis):
271262 [ (one_c * zx ) - ys , (one_c * yz ) + xs , (one_c * zz ) + c ]]
272263
273264#------------------------------------------------------------------------------
274- def transform_array (rot_mtx , vec_array ) :
265+ def transform_array (rot_mtx : list [ list [ float ]], vec_array : list [ list [ float ]]) -> list [ list [ float ]] :
275266
276267 '''transform_array( matrix, vector_array ) -> vector_array
277268
@@ -280,7 +271,7 @@ def transform_array(rot_mtx,vec_array):
280271 return [transform (rot_mtx , x ) for x in vec_array ]
281272
282273#------------------------------------------------------------------------------
283- def translate_array (trans_vec , vec_array ):
274+ def translate_array (trans_vec : list [ float ], vec_array : list [ list [ float ]] ):
284275
285276 '''translate_array(trans_vec,vec_array) -> vec_array
286277
@@ -291,17 +282,17 @@ def translate_array(trans_vec,vec_array):
291282 return [add (trans_vec , x ) for x in vec_array ]
292283
293284#------------------------------------------------------------------------------
294- def fit_apply (fit_result , vec_array ) :
285+ def fit_apply (fit_result : tuple [ list [ float ], list [ float ], list [ list [ float ]], float ], vec_array : list [ list [ float ]]) -> list [ list [ float ]] :
295286 '''fit_apply(fir_result,vec_array) -> vec_array
296-
287+
297288 Applies a fit result to an array of vectors
298289 '''
299290
300291 t1 , mt2 , m = fit_result [:3 ]
301292 return [add (t1 , transform (m , add (mt2 , x ))) for x in vec_array ]
302293
303294#------------------------------------------------------------------------------
304- def fit (target_array , source_array ) :
295+ def fit (target_array : list [ list [ float ]] , source_array : list [ list [ float ]]) -> tuple [ list [ float ], list [ float ], list [ list [ float ]], float ] | None :
305296
306297 '''fit(target_array, source_array) -> (t1, t2, rot_mtx, rmsd) [fit_result]
307298
@@ -370,7 +361,6 @@ def fit(target_array, source_array):
370361 iters = 0
371362 while (iters < maxiter ):
372363 iters = iters + 1
373- ix = (iters - 1 )% ndim
374364 iy = iters % ndim
375365 iz = (iters + 1 )% ndim
376366 sig = corr_mtx [iz ][iy ] - corr_mtx [iy ][iz ]
0 commit comments