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)