Package Bio :: Package PDB :: Module NeighborSearch
[hide private]
[frames] | no frames]

Source Code for Module Bio.PDB.NeighborSearch

  1  # Copyright (C) 2002, Thomas Hamelryck (thamelry@binf.ku.dk) 
  2  # This code is part of the Biopython distribution and governed by its 
  3  # license.  Please see the LICENSE file that should have been included 
  4  # as part of this package. 
  5   
  6  from Numeric import array 
  7   
  8  from Bio.KDTree import * 
  9  from PDBExceptions import PDBException 
 10  from Selection import unfold_entities, get_unique_parents, entity_levels, uniqueify 
 11   
 12   
 13  __doc__="Fast atom neighbor lookup using a KD tree (implemented in C++)." 
 14   
 15   
 16   
17 -class NeighborSearch:
18 """ 19 This class can be used for two related purposes: 20 21 1. To find all atoms/residues/chains/models/structures within radius 22 of a given query position. 23 24 2. To find all atoms/residues/chains/models/structures that are within 25 a fixed radius of each other. 26 27 NeighborSearch makes use of the Bio.KDTree C++ module, so it's fast. 28 """
29 - def __init__(self, atom_list, bucket_size=10):
30 """ 31 o atom_list - list of atoms. This list is used in the queries. 32 It can contain atoms from different structures. 33 o bucket_size - bucket size of KD tree. You can play around 34 with this to optimize speed if you feel like it. 35 """ 36 self.atom_list=atom_list 37 # get the coordinates 38 coord_list=map(lambda a: a.get_coord(), atom_list) 39 # to Nx3 array of type float 40 self.coords=array(coord_list).astype("f") 41 assert(bucket_size>1) 42 assert(self.coords.shape[1]==3) 43 assert(self.coords.typecode()=="f") 44 self.kdt=KDTree(3, bucket_size) 45 self.kdt.set_coords(self.coords)
46 47 # Private 48
49 - def _get_unique_parent_pairs(self, pair_list):
50 # translate a list of (entity, entity) tuples to 51 # a list of (parent entity, parent entity) tuples, 52 # thereby removing duplicate (parent entity, parent entity) 53 # pairs. 54 # o pair_list - a list of (entity, entity) tuples 55 parent_pair_list=[] 56 for (e1, e2) in pair_list: 57 p1=e1.get_parent() 58 p2=e2.get_parent() 59 if p1==p2: 60 continue 61 elif p1<p2: 62 parent_pair_list.append((p1, p2)) 63 else: 64 parent_pair_list.append((p2, p1)) 65 return uniqueify(parent_pair_list)
66 67 # Public 68
69 - def search(self, center, radius, level="A"):
70 """Neighbor search. 71 72 Return all atoms/residues/chains/models/structures 73 that have at least one atom within radius of center. 74 What entitity level is returned (e.g. atoms or residues) 75 is determined by level (A=atoms, R=residues, C=chains, 76 M=models, S=structures). 77 78 o center - Numeric array 79 o radius - float 80 o level - char (A, R, C, M, S) 81 """ 82 if not level in entity_levels: 83 raise PDBException, "%s: Unknown level" % level 84 self.kdt.search(center, radius) 85 indices=self.kdt.get_indices() 86 n_atom_list=[] 87 atom_list=self.atom_list 88 for i in indices: 89 a=atom_list[i] 90 n_atom_list.append(a) 91 if level=="A": 92 return n_atom_list 93 else: 94 return unfold_entities(n_atom_list, level)
95
96 - def search_all(self, radius, level="A"):
97 """All neighbor search. 98 99 Search all entities that have atoms pairs within 100 radius. 101 102 o radius - float 103 o level - char (A, R, C, M, S) 104 """ 105 if not level in entity_levels: 106 raise PDBException, "%s: Unknown level" % level 107 self.kdt.all_search(radius) 108 indices=self.kdt.all_get_indices() 109 atom_list=self.atom_list 110 atom_pair_list=[] 111 for i1, i2 in indices: 112 a1=atom_list[i1] 113 a2=atom_list[i2] 114 atom_pair_list.append((a1, a2)) 115 if level=="A": 116 # return atoms 117 return atom_pair_list 118 next_level_pair_list=atom_pair_list 119 for l in ["R", "C", "M", "S"]: 120 next_level_pair_list=self._get_unique_parent_pairs(next_level_pair_list) 121 if level==l: 122 return next_level_pair_list
123 124 if __name__=="__main__": 125 126 from RandomArray import * 127
128 - class Atom:
129 - def __init__(self):
130 self.coord=(100*random(3))
131
132 - def get_coord(self):
133 return self.coord
134 135 for i in range(0, 20): 136 al=[] 137 for i in range(0, 100): 138 al.append(Atom()) 139 140 ns=NeighborSearch(al) 141 142 print "Found ", len(ns.search_all(5.0)) 143