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