Package mmLib :: Module GeometryDict
[hide private]
[frames] | no frames]

Source Code for Module mmLib.GeometryDict

  1  ## Copyright 2002-2010 by PyMMLib Development Group (see AUTHORS file) 
  2  ## This code is part of the PyMMLib distribution and governed by 
  3  ## its license.  Please see the LICENSE file that should have been 
  4  ## included as part of this package. 
  5  """Geometry hasing/fast lookup classes. 
  6  """ 
  7   
  8  from __future__ import generators 
  9  import random 
 10  import math 
 11  import itertools 
 12   
 13   
14 -class XYZDict(object):
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 """
19 - def __init__(self, resolution):
20 self.resolution = resolution 21 self.geom_dict = {}
22
23 - def calc_geom_key(self, position):
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
65 - def iter_all(self):
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
72 - def iter_cube_intersection(self, position, radius):
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
95 - def iter_sphere_intersection(self, position, radius):
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
110 - def iter_contact_distance(self, distance):
111 """Iterates all items within a given contact distance. 112 """ 113 visited = {} 114 115 for geom_tuple1 in self.iter_all(): 116 visited[geom_tuple1[1]] = True 117 118 for geom_tuple2 in self.iter_cube_intersection( 119 geom_tuple1[0], distance): 120 121 if visited.has_key(geom_tuple2[1]): 122 continue 123 124 p1 = geom_tuple1[0] 125 p2 = geom_tuple2[0] 126 127 d = math.sqrt((p1[0]-p2[0])**2 + \ 128 (p1[1]-p2[1])**2 + \ 129 (p1[2]-p2[2])**2) 130 131 if d>distance: 132 continue 133 134 yield geom_tuple1, geom_tuple2, d
135 136 137 ### <testing>
138 -def test_module():
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 #d = math.sqrt((p1[0]-p2[0])**2 + (p1[1]-p2[1])**2 + (p1[2]-p2[2])**2) 159 #print atm1, atm2, "%6.2f" % (d) 160 161 print "%d Bonds" % (cnt)
162 163 if __name__=="__main__": 164 test_module() 165 ### <testing> 166