Back to index

python-biopython  1.60
NeighborSearch.py
Go to the documentation of this file.
00001 # Copyright (C) 2002, Thomas Hamelryck (thamelry@binf.ku.dk)
00002 # This code is part of the Biopython distribution and governed by its
00003 # license.  Please see the LICENSE file that should have been included
00004 # as part of this package.
00005 
00006 """Fast atom neighbor lookup using a KD tree (implemented in C++)."""
00007 
00008 import numpy
00009 
00010 from Bio.KDTree import KDTree
00011 
00012 from Bio.PDB.PDBExceptions import PDBException
00013 from Bio.PDB.Selection import unfold_entities, entity_levels, uniqueify
00014 
00015 
00016 class NeighborSearch(object):
00017     """
00018     This class can be used for two related purposes:
00019 
00020     1. To find all atoms/residues/chains/models/structures within radius 
00021     of a given query position. 
00022 
00023     2. To find all atoms/residues/chains/models/structures that are within 
00024     a fixed radius of each other.
00025 
00026     NeighborSearch makes use of the Bio.KDTree C++ module, so it's fast.
00027     """
00028     def __init__(self, atom_list, bucket_size=10):
00029         """
00030         o atom_list - list of atoms. This list is used in the queries.
00031         It can contain atoms from different structures.
00032         o bucket_size - bucket size of KD tree. You can play around 
00033         with this to optimize speed if you feel like it.
00034         """
00035         self.atom_list=atom_list
00036         # get the coordinates
00037         coord_list = [a.get_coord() for a in atom_list]
00038         # to Nx3 array of type float
00039         self.coords=numpy.array(coord_list).astype("f")
00040         assert(bucket_size>1)
00041         assert(self.coords.shape[1]==3)
00042         self.kdt=KDTree(3, bucket_size)
00043         self.kdt.set_coords(self.coords)
00044 
00045     # Private
00046 
00047     def _get_unique_parent_pairs(self, pair_list):
00048         # translate a list of (entity, entity) tuples to 
00049         # a list of (parent entity, parent entity) tuples,
00050         # thereby removing duplicate (parent entity, parent entity)
00051         # pairs.
00052         # o pair_list - a list of (entity, entity) tuples
00053         parent_pair_list=[]
00054         for (e1, e2) in pair_list:
00055             p1=e1.get_parent()
00056             p2=e2.get_parent()
00057             if p1==p2:
00058                 continue
00059             elif p1<p2:
00060                 parent_pair_list.append((p1, p2))
00061             else:
00062                 parent_pair_list.append((p2, p1))
00063         return uniqueify(parent_pair_list)
00064 
00065     # Public
00066 
00067     def search(self, center, radius, level="A"):
00068         """Neighbor search.
00069 
00070         Return all atoms/residues/chains/models/structures
00071         that have at least one atom within radius of center.
00072         What entitity level is returned (e.g. atoms or residues)
00073         is determined by level (A=atoms, R=residues, C=chains,
00074         M=models, S=structures).
00075 
00076         o center - Numeric array 
00077         o radius - float
00078         o level - char (A, R, C, M, S)
00079         """
00080         if not level in entity_levels:
00081             raise PDBException("%s: Unknown level" % level)
00082         self.kdt.search(center, radius)
00083         indices=self.kdt.get_indices()
00084         n_atom_list=[]
00085         atom_list=self.atom_list
00086         for i in indices:
00087             a=atom_list[i]
00088             n_atom_list.append(a)
00089         if level=="A":
00090             return n_atom_list
00091         else:
00092             return unfold_entities(n_atom_list, level)
00093             
00094     def search_all(self, radius, level="A"):
00095         """All neighbor search.
00096 
00097         Search all entities that have atoms pairs within
00098         radius. 
00099 
00100         o radius - float
00101         o level - char (A, R, C, M, S)
00102         """
00103         if not level in entity_levels:
00104             raise PDBException("%s: Unknown level" % level)
00105         self.kdt.all_search(radius)
00106         indices=self.kdt.all_get_indices()
00107         atom_list=self.atom_list
00108         atom_pair_list=[]
00109         for i1, i2 in indices:
00110             a1=atom_list[i1]
00111             a2=atom_list[i2]
00112             atom_pair_list.append((a1, a2))
00113         if level=="A":
00114             # return atoms
00115             return atom_pair_list
00116         next_level_pair_list=atom_pair_list
00117         for l in ["R", "C", "M", "S"]:
00118             next_level_pair_list=self._get_unique_parent_pairs(next_level_pair_list)
00119             if level==l:
00120                 return next_level_pair_list 
00121 
00122 if __name__=="__main__":
00123 
00124     from numpy.random import random
00125 
00126     class Atom(object):
00127         def __init__(self):
00128             self.coord=(100*random(3))
00129 
00130         def get_coord(self):
00131             return self.coord
00132 
00133     for i in range(0, 20):
00134         #Make a list of 100 atoms
00135         al = [Atom() for j in range(100)]
00136 
00137         ns=NeighborSearch(al)
00138 
00139         print "Found ", len(ns.search_all(5.0))
00140