def test_search(b, qns):
"""
Test finding neighbors for a given query vector and type of box.
Parameters
----------
b : list
MDAnalysis dimensions like list
qns : tuple
a query point and a list of expected neighbors.
"""
b = np.array(b, dtype=np.float32)
q = transform_StoR(np.array(qns[0], dtype=np.float32), b)
# Setting up the periodic tree
tree = PeriodicKDTree(b)
coords = transform_StoR(f_dataset, b)
tree.set_coords(coords) # Input real space coordinates
# Carry out the search and retrieve results
tree.search(q, radius)
indices = tree.get_indices()
if indices:
found_neighbors = np.sort(coords[indices], axis=0)
else:
found_neighbors = list()
if qns[1]:
expected_neighbors = transform_StoR(np.array(qns[1],dtype=np.float32), b)
expected_neighbors = np.sort(expected_neighbors, axis=0)
else:
expected_neighbors = list()
assert_equal(found_neighbors, expected_neighbors)
def test_searchfail():
coords = np.array([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
b = np.array([10, 10, 10, 90, 90, 90], dtype=np.float32)
cutoff = 1.0
search_radius = 2.0
query = np.array([1, 1, 1], dtype=np.float32)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
match = 'Set cutoff greater or equal to the radius.'
with pytest.raises(RuntimeError, match=match):
tree.search(query, search_radius)
def test_searchfail():
coords = np.array([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
b = np.array([10, 10, 10, 90, 90, 90], dtype=np.float32)
cutoff = 1.0
search_radius = 2.0
query = np.array([1, 1, 1], dtype=np.float32)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
match = 'Set cutoff greater or equal to the radius.'
with pytest.raises(RuntimeError, match=match):
tree.search(query, search_radius)
def test_nopbc():
cutoff = 0.3
q = np.array([0.2, 0.3, 0.1])
coords = f_dataset.copy()
tree = PeriodicKDTree(box=None)
tree.set_coords(coords)
indices = tree.search(q, cutoff)
assert_equal(indices, [0, 2])
def test_nopbc():
cutoff = 0.3
q = np.array([0.2, 0.3, 0.1])
coords = f_dataset.copy()
tree = PeriodicKDTree(box=None)
tree.set_coords(coords)
indices = tree.search(q, cutoff)
assert_equal(indices, [0, 2])
class CollisionDetector(object):
"""Given a set of N particles this class allows to query if any new particle
will overlap. If particles overlap this is called a collision
Example
-------
>>> cd = CollisionDetector(coords, 3.81, [100, 100, 100])
>>> cd.collision([50, 50, 50])
"""
def __init__(self, coords, diameter, box):
"""
Parameters
----------
coords : array_like (N, 3)
coordinates of N particles
diameter : float
diameter of particles
box : array_like (3)
box dimensions
"""
fullbox = 90 * np.ones(6, dtype=np.float32)
fullbox[:3] = box
coords = np.asarray(coords, dtype=np.float32)
max_cutoff = np.sum(np.sqrt(fullbox**2)) / 2
self._kdt = PKDTree(fullbox)
self._kdt.set_coords(coords, cutoff=max_cutoff)
self._diameter = diameter
def collision(self, pos):
"""check is there is a collision
Parameters
----------
pos : array_like (3)
coordinates of new particles
Returns
-------
collision : bool
``True`` if there is a collision,
"""
pos = np.asarray(pos, dtype=np.float32)
self._kdt.search(pos, self._diameter)
indices = self._kdt.get_indices()
return len(indices) != 0
def test_search(b, q, result):
b = np.array(b, dtype=np.float32)
q = transform_StoR(np.array(q, dtype=np.float32), b)
cutoff = 3.0
coords = transform_StoR(f_dataset, b)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
indices = tree.search(q, cutoff)
assert_equal(indices, result)
def test_search(b, q, result):
b = np.array(b, dtype=np.float32)
q = transform_StoR(np.array(q, dtype=np.float32), b)
cutoff = 3.0
coords = transform_StoR(f_dataset, b)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
indices = tree.search(q, cutoff)
assert_equal(indices, result)
class AtomNeighborSearch(object):
"""This class can be used to find all atoms/residues/segments within the
radius of a given query position.
For the neighbor search, this class uses the BioPython KDTree and its
wrapper PeriodicKDTree for non-periodic and periodic systems, respectively.
"""
def __init__(self, atom_group, box=None, bucket_size=10):
"""
Parameters
----------
atom_list : AtomGroup
list of atoms
box : array-like or ``None``, optional, default ``None``
Simulation cell dimensions in the form of
:attr:`MDAnalysis.trajectory.base.Timestep.dimensions` when
periodic boundary conditions should be taken into account for
the calculation of contacts.
bucket_size : int
Number of entries in leafs of the KDTree. If you suffer poor
performance you can play around with this number. Increasing the
`bucket_size` will speed up the construction of the KDTree but
slow down the search.
"""
self.atom_group = atom_group
self._u = atom_group.universe
self._box = box
if box is None:
self.kdtree = KDTree(dim=3, bucket_size=bucket_size)
else:
self.kdtree = PeriodicKDTree(box, bucket_size=bucket_size)
self.kdtree.set_coords(atom_group.positions)
def search(self, atoms, radius, level='A'):
"""
Return all atoms/residues/segments that are within *radius* of the
atoms in *atoms*.
Parameters
----------
atoms : AtomGroup, MDAnalysis.core.groups.Atom
list of atoms
radius : float
Radius for search in Angstrom.
level : str
char (A, R, S). Return atoms(A), residues(R) or segments(S) within
*radius* of *atoms*.
"""
if isinstance(atoms, Atom):
positions = atoms.position.reshape(1, 3)
else:
positions = atoms.positions
indices = []
for pos in positions:
self.kdtree.search(pos, radius)
indices.append(self.kdtree.get_indices())
unique_idx = np.unique([i for l in indices for i in l]).astype(np.int64)
return self._index2level(unique_idx, level)
def _index2level(self, indices, level):
"""Convert list of atom_indices in a AtomGroup to either the
Atoms or segments/residues containing these atoms.
Parameters
----------
indices
list of atom indices
level : str
char (A, R, S). Return atoms(A), residues(R) or segments(S) within
*radius* of *atoms*.
"""
n_atom_list = self.atom_group[indices]
if level == 'A':
if not n_atom_list:
return []
else:
return n_atom_list
elif level == 'R':
return list({a.residue for a in n_atom_list})
elif level == 'S':
return list(set([a.segment for a in n_atom_list]))
else:
raise NotImplementedError('{0}: level not implemented'.format(level))