1
2
3
4
5 """Geometry hasing/fast lookup classes.
6 """
7
8 from __future__ import generators
9 import random
10 import math
11 import itertools
12
13
15 """Hash all objects according to their position, allowing spacial
16 location of objects quickly. This is a brain-dead simple implementation,
17 but it gets the job done.
18 """
20 self.resolution = resolution
21 self.geom_dict = {}
22
24 """Calculates the cube key for the given position.
25 """
26 res = self.resolution
27
28 return ( int(position[0] / res),
29 int(position[1] / res),
30 int(position[2] / res) )
31
32 - def add(self, position, item):
33 """Add a item.
34 """
35 res = self.resolution
36
37 geom_key = ( int(position[0] / res),
38 int(position[1] / res),
39 int(position[2] / res) )
40
41 geom_tuple = (position, item)
42
43 if self.geom_dict.has_key(geom_key):
44 self.geom_dict[geom_key].append(geom_tuple)
45 else:
46 self.geom_dict[geom_key] = [geom_tuple]
47
48 - def remove(self, position, item):
49 """Remove an item.
50 """
51 geom_key = self.calc_geom_key(position)
52
53 try:
54 geom_list = self.geom_dict[geom_key]
55 except KeyError:
56 pass
57 else:
58 for geom_tuple in geom_list:
59 if geom_tuple[1]==item:
60 geom_list.remove(geom_tuple)
61 return
62
63 raise ValueError, "GeometryDict.remove(x) x not in GeometryDict"
64
66 """Iter all items
67 """
68 for geom_list in self.geom_dict.values():
69 for geom_tuple in geom_list:
70 yield geom_tuple
71
73 """Iterate all objects which intersect the cube in no particular
74 order.
75 """
76 center_geom_key = self.calc_geom_key(position)
77 bounding_cube_blocks = int(radius / self.resolution) + 1
78
79 for i in xrange(-bounding_cube_blocks, bounding_cube_blocks+1):
80 for j in xrange(-bounding_cube_blocks, bounding_cube_blocks+1):
81 for k in xrange(-bounding_cube_blocks, bounding_cube_blocks+1):
82
83 geom_key = ( center_geom_key[0] + i,
84 center_geom_key[1] + j,
85 center_geom_key[2] + k )
86
87 try:
88 geom_list = self.geom_dict[geom_key]
89 except KeyError:
90 continue
91
92 for geom_tuple in geom_list:
93 yield geom_tuple
94
96 """Iterate all objects which intersect the cube in no particular
97 order.
98 """
99 for geom_tuple in self.iter_cube_intersection(position, radius):
100 ipos = geom_tuple[0]
101
102 x = ipos[0] - position[0]
103 y = ipos[1] - position[1]
104 z = ipos[2] - position[2]
105 d = math.sqrt(x*x + y*y + z*z)
106
107 if d <= radius:
108 yield geom_tuple
109
135
136
137
139 import sys
140 import FileIO
141
142 struct = FileIO.LoadStructure(fil=sys.argv[1])
143
144 print "Structure Loaded"
145
146 gdict = XYZDict(2.0)
147
148 cnt = 0
149 for atm in struct.iter_atoms():
150 gdict.add(atm.position, atm)
151 cnt += 1
152
153 print "Hashed %d Atoms" % (cnt)
154
155 cnt = 0
156 for (p1, atm1),(p2,atm2), dist in gdict.iter_contact_distance(1.9):
157 cnt += 1
158
159
160
161 print "%d Bonds" % (cnt)
162
163 if __name__=="__main__":
164 test_module()
165
166