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_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)
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))