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 test_setcoords(b, cut, result):
coords = np.array([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
if b is not None:
b = np.array(b, dtype=np.float32)
tree = PeriodicKDTree(box=b)
print(b, tree.box, cut, result)
with pytest.raises(RuntimeError, match=result):
tree.set_coords(coords, cutoff=cut)
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_setcoords(b, cut, result):
coords = np.array([[1, 1, 1], [2, 2, 2]], dtype=np.float32)
if b is not None:
b = np.array(b, dtype=np.float32)
tree = PeriodicKDTree(box=b)
print(b, tree.box, cut, result)
with pytest.raises(RuntimeError, match=result):
tree.set_coords(coords, cutoff=cut)
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_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)
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_searchpairs(b, radius, result):
b = np.array(b, dtype=np.float32)
cutoff = 2.0
coords = transform_StoR(f_dataset, b)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
if cutoff < radius:
with pytest.raises(RuntimeError, match=result):
indices = tree.search_pairs(radius)
else:
indices = tree.search_pairs(radius)
assert_equal(len(indices), len(result))
def _apply_KDTree(self, group):
"""KDTree based selection is about 7x faster than distmat
for typical problems.
"""
sel = self.sel.apply(group)
# All atoms in group that aren't in sel
sys = group[~np.in1d(group.indices, sel.indices)]
box = self.validate_dimensions(group.dimensions)
if box is None:
kdtree = KDTree(dim=3, bucket_size=10)
kdtree.set_coords(sys.positions)
found_indices = []
for atom in sel.positions:
kdtree.search(atom, self.cutoff)
found_indices.append(kdtree.get_indices())
unique_idx = np.unique(np.concatenate(found_indices))
else:
kdtree = PeriodicKDTree(box, bucket_size=10)
kdtree.set_coords(sys.positions)
kdtree.search(sel.positions, self.cutoff)
unique_idx = np.asarray(kdtree.get_indices())
# These are the indices from SYS that were seen when
# probing with SEL
return sys[unique_idx.astype(np.int32)].unique
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, 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_find_images(b, qcs):
"""
Test the generation of images 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(qcs[0], dtype=np.float32), b)
tree = PeriodicKDTree(b)
coords = transform_StoR(f_dataset, b)
tree.set_coords(coords) # Input real space coordinates
cs = np.sort(transform_StoR(np.array(qcs[1], dtype=np.float32), b), axis=0)
q_wrapped = apply_PBC(q.reshape((1, 3)), b)
found_images = np.sort(tree.find_images(q_wrapped, radius), axis=0)
assert_almost_equal(found_images, cs, decimal=6)
def _apply_KDTree(self, group):
box = group.dimensions if self.periodic else None
if box is None:
kdtree = KDTree(dim=3, bucket_size=10)
else:
kdtree = PeriodicKDTree(box, bucket_size=10)
kdtree.set_coords(group.positions)
kdtree.search(self.ref, self.cutoff)
found_indices = kdtree.get_indices()
return group[found_indices].unique
def _apply_KDTree(self, group):
"""Selection using KDTree and PeriodicKDTree for aperiodic and
fully-periodic systems, respectively.
"""
sel = self.sel.apply(group)
box = self.validate_dimensions(group.dimensions)
ref = sel.center_of_geometry(pbc=self.periodic)
if box is None:
kdtree = KDTree(dim=3, bucket_size=10)
else:
kdtree = PeriodicKDTree(box, bucket_size=10)
kdtree.set_coords(group.positions)
kdtree.search(ref, self.cutoff)
found_indices = kdtree.get_indices()
return group[found_indices].unique
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_searchpairs(b, radius, result):
b = np.array(b, dtype=np.float32)
cutoff = 2.0
coords = transform_StoR(f_dataset, b)
tree = PeriodicKDTree(box=b)
tree.set_coords(coords, cutoff=cutoff)
if cutoff < radius:
with pytest.raises(RuntimeError, match=result):
indices = tree.search_pairs(radius)
else:
indices = tree.search_pairs(radius)
assert_equal(len(indices), len(result))
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 test_ckd_searchpairs_nopbc(radius, result):
coords = f_dataset.copy()
tree = PeriodicKDTree()
tree.set_coords(coords)
indices = tree.search_pairs(radius)
assert_equal(indices, result)
def test_initialize_bm(box, rm):
"""
Assert the construction of the recripocal box matrix.
"""
assert_almost_equal(PeriodicKDTree(box)._rm, rm, decimal=7)
def test_set_coords():
match = 'coords must be a sequence of 3 dimensional coordinates'
with pytest.raises(ValueError, match=match):
xy = np.array([[2, 2], [5, 5], [1.1, 1.1]], dtype=np.float32)
tree = PeriodicKDTree(boxes_1[0])
tree.set_coords(xy)
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))
def test_ckd_searchpairs_nopbc(radius, result):
coords = f_dataset.copy()
tree = PeriodicKDTree()
tree.set_coords(coords)
indices = tree.search_pairs(radius)
assert_equal(indices, result)
