def buildAffinityMatrix(self, atoms, cutoff=4):
"""Build the affinity matrix for given *atoms*.
Note that if you do not want to incorporate hydrogen and non-protein
atoms in calculations, make the selection ``"noh and protein"``.
:arg atoms: atoms for which the affinity matrix will be calculated
:type atoms: :class:`~prody.atomic.Atomic`
:arg cutoff: pairwise atomic contact cutoff distance, default is 4 Å
:type cutoff: float
"""
if not isinstance(atoms, prody.Atomic):
raise TypeError('atoms must be an Atomic instance, '
'{0:s} is not valid'.format(type(atoms)))
cutoff = float(cutoff)
assert cutoff > 0, 'cutoff distance must be greater than 0'
from KDTree import KDTree
start = time.time()
if not isinstance(atoms, prody.AtomGroup):
atoms = atoms.getAtomGroup().copy(atoms)
n_atoms = atoms.numAtoms()
hv = prody.HierView(atoms)
n_res = hv.numResidues()
rids = np.zeros(n_atoms, int) # residue indices of atoms
rlen = np.zeros(n_res) # residue lengths
resmap = {} # used for symmetry purposes
for i, res in enumerate(hv.iterResidues()):
rids[ res.getIndices() ] = i
rlen[ i ] = len(res)
res = (res.getChid(), res.getNumber(), res.getIcode())
resmap[i] = res
resmap[res] = i
self._resmap = resmap
LOGGER.debug('Atoms were evalulated in {0:.2f}s.'
.format(time.time()-start))
start = time.time()
kdtree = KDTree(3)
kdtree.set_coords(atoms.getCoords())
kdtree.all_search(cutoff)
LOGGER.debug('KDTree was built in {0:.2f}s.'
.format(time.time()-start))
start = time.time()
affinity = defaultdict(int)
for i, j in kdtree.all_get_indices():
i = rids[i]
j = rids[j]
if i == j:
affinity[(i,j)] += 0.5
else:
affinity[(i,j)] += 1
length = len(affinity)
i = np.zeros(length, int)
j = np.zeros(length, int)
v = np.zeros(length, float)
k = 0
for key, value in affinity.iteritems():
i[k] = key[0]
j[k] = key[1]
v[k] = value
k += 1
rlen = rlen ** -0.5
# = Nij * (1/Ni^0.5) * (1/Nj^0.5)
v = v * rlen[i] * rlen[j]
affinity = sparse.coo_matrix((v, (i,j)), shape=(n_res, n_res))
self._affinity = affinity + affinity.T
LOGGER.debug('Affinity matrix was built in {0:.2f}s.'
.format(time.time()-start))
self._stationary = None
self._n_nodes = n_res
class CoordinateNeighborSearch(object):
"""
This class can be used for two related purposes:
1. To find all indices of a coordinate list within radius
of a given query position.
2. To find all indices of a coordinate list that are within
a fixed radius of each other.
CoordinateNeighborSearch makes use of the KDTree C++ module, so it's fast.
"""
def __init__(self, coordinates, bucket_size=10): # , copy=True):
"""
:Arguments:
*coordinates*
list of N coordinates (Nx3 numpy array)
*bucket_size*
bucket size of KD tree. You can play around with this to
optimize speed if you feel like it.
"""
# to Nx3 array of type float (required for the C++ code)
## (also force a copy by default and make sure that the array order is compatible
## with the C++ code)
##self.coords=numpy.array(coordinates,dtype=numpy.float32,copy=copy,order='C')
self.coords = numpy.asarray(coordinates, dtype=numpy.float32, order='C')
assert (self.coords.dtype == numpy.float32)
assert (bucket_size > 1)
assert (self.coords.shape[1] == 3)
self.kdt = KDTree(3, bucket_size)
self.kdt.set_coords(self.coords)
def search(self, center, radius, distances=False):
"""Neighbor search.
Return all indices in the coordinates list that have at least
one atom within *radius* of *center*.
:Arguments:
* center
numpy array
* radius
float
* distances
bool ``True``: return (indices,distances); ``False``: return indices only
"""
self.kdt.search(center, radius)
if distances:
return self.kdt.get_indices(), self.kdt.get_radii()
else:
return self.kdt.get_indices()
def search_list(self, centers, radius):
"""Search neighbours near all centers.
Returns all indices that are within *radius* of any center listed in
*centers*, i.e. "find all A within R of B" where A are the
coordinates used for setting up the CoordinateNeighborSearch and B
are the centers.
:Arguments:
*centers*
Mx3 numpy array of M centers
*radius*
float
"""
self.kdt.list_search(centers, radius)
return self.kdt.list_get_indices()
def search_all(self, radius, distances=False):
"""All neighbor search.
Return all index pairs corresponding to coordinates within the *radius*.
:Arguments:
*radius*
float
*distances*
bool ``True``: return (indices,distances); ``False``: return indices only
[``False``]
"""
self.kdt.all_search(radius)
if distances:
return self.kdt.all_get_indices(), self.kdt.all_get_radii()
else:
return self.kdt.all_get_indices()
def _distances(self):
"""Return all distances after search()."""
return self.kdt.get_radii()
def _distances_all(self):
"""Return all distances after search_all()."""
return self.kdt.all_get_radii()
class CoordinateNeighborSearch(object):
"""
This class can be used for two related purposes:
1. To find all indices of a coordinate list within radius
of a given query position.
2. To find all indices of a coordinate list that are within
a fixed radius of each other.
CoordinateNeighborSearch makes use of the KDTree C++ module, so it's fast.
"""
def __init__(self, coordinates, bucket_size=10): # , copy=True):
"""
:Arguments:
*coordinates*
list of N coordinates (Nx3 numpy array)
*bucket_size*
bucket size of KD tree. You can play around with this to
optimize speed if you feel like it.
"""
# to Nx3 array of type float (required for the C++ code)
## (also force a copy by default and make sure that the array order is compatible
## with the C++ code)
##self.coords=numpy.array(coordinates,dtype=numpy.float32,copy=copy,order='C')
self.coords = numpy.asarray(coordinates,
dtype=numpy.float32,
order='C')
assert (self.coords.dtype == numpy.float32)
assert (bucket_size > 1)
assert (self.coords.shape[1] == 3)
self.kdt = KDTree(3, bucket_size)
self.kdt.set_coords(self.coords)
def search(self, center, radius, distances=False):
"""Neighbor search.
Return all indices in the coordinates list that have at least
one atom within *radius* of *center*.
:Arguments:
* center
numpy array
* radius
float
* distances
bool ``True``: return (indices,distances); ``False``: return indices only
"""
self.kdt.search(center, radius)
if distances:
return self.kdt.get_indices(), self.kdt.get_radii()
else:
return self.kdt.get_indices()
def search_list(self, centers, radius):
"""Search neighbours near all centers.
Returns all indices that are within *radius* of any center listed in
*centers*, i.e. "find all A within R of B" where A are the
coordinates used for setting up the CoordinateNeighborSearch and B
are the centers.
:Arguments:
*centers*
Mx3 numpy array of M centers
*radius*
float
"""
self.kdt.list_search(centers, radius)
return self.kdt.list_get_indices()
def search_all(self, radius, distances=False):
"""All neighbor search.
Return all index pairs corresponding to coordinates within the *radius*.
:Arguments:
*radius*
float
*distances*
bool ``True``: return (indices,distances); ``False``: return indices only
[``False``]
"""
self.kdt.all_search(radius)
if distances:
return self.kdt.all_get_indices(), self.kdt.all_get_radii()
else:
return self.kdt.all_get_indices()
def _distances(self):
"""Return all distances after search()."""
return self.kdt.get_radii()
def _distances_all(self):
"""Return all distances after search_all()."""
return self.kdt.all_get_radii()
