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 """ 20 This class can be used for two related purposes: 21 22 1. To find all atoms/residues/chains/models/structures within radius 23 of a given query position. 24 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 """ 32 o atom_list - list of atoms. This list is used in the queries. 33 It can contain atoms from different structures. 34 o bucket_size - bucket size of KD tree. You can play around 35 with this to optimize speed if you feel like it. 36 """ 37 self.atom_list=atom_list 38 # get the coordinates 39 coord_list = [a.get_coord() for a in atom_list] 40 # to Nx3 array of type float 41 self.coords=numpy.array(coord_list).astype("f") 42 assert(bucket_size>1) 43 assert(self.coords.shape[1]==3) 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 numpy.random import random 127
128 - class Atom(object):
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 #Make a list of 100 atoms 137 al = [Atom() for j in range(100)] 138 ns = NeighborSearch(al) 139 print("Found %i" % len(ns.search_all(5.0))) 140