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