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