def _apply_KDTree(self, group): kdtree = KDTree(dim=3, bucket_size=10) kdtree.set_coords(group.positions) kdtree.search(self.ref, self.cutoff) found_indices = kdtree.get_indices() return unique(group[found_indices])
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
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_all_search(nr_points, dim, bucket_size, query_radius): """Test fixed neighbor search. Search all point pairs that are within radius. Arguments: - nr_points: number of points used in test - dim: dimension of coords - bucket_size: nr of points per tree node - query_radius: radius of search Returns true if the test passes. """ kdt = KDTree(dim, bucket_size) coords = random.random((nr_points, dim)) kdt.set_coords(coords) kdt.all_search(query_radius) indices = kdt.all_get_indices() if indices is None: l1 = 0 else: l1 = len(indices) radii = kdt.all_get_radii() if radii is None: l2 = 0 else: l2 = len(radii) if l1 == l2: return True else: return False
def test_search(nr_points, dim, bucket_size, radius): """Test search all points within radius of center. Search all point pairs that are within radius. Arguments: - nr_points: number of points used in test - dim: dimension of coords - bucket_size: nr of points per tree node - radius: radius of search Returns true if the test passes. """ kdt = KDTree(dim, bucket_size) coords = random.random((nr_points, dim)) kdt.set_coords(coords) kdt.search(coords[0], radius * 100) radii = kdt.get_radii() l1 = 0 for i in range(0, nr_points): p = coords[i] if _dist(p, coords[0]) <= radius * 100: l1 = l1 + 1 if l1 == len(radii): return True else: return False
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 but periodic = True not supported. (KDTree routine is ca 15% slower than the distance matrix one) """ sel = self.sel.apply(group) ref = sel.center_of_geometry() kdtree = KDTree(dim=3, bucket_size=10) kdtree.set_coords(group.positions) kdtree.search(ref, self.cutoff) found_indices = kdtree.get_indices() return unique(group[found_indices])
def test_KDTree_exceptions(self): kdt = KDTree(dim, bucket_size) with self.assertRaises(Exception) as context: kdt.set_coords(random.random((nr_points, dim)) * 100000000000000) self.assertTrue( "Points should lie between -1e6 and 1e6" in str(context.exception)) with self.assertRaises(Exception) as context: kdt.set_coords(random.random((nr_points, dim - 2))) self.assertTrue("Expected a Nx%i NumPy array" % dim in str(context.exception)) with self.assertRaises(Exception) as context: kdt.search(array([0, 0, 0]), radius) self.assertTrue("No point set specified" in str(context.exception))
def _apply_KDTree(self, group): """Selection using KDTree but periodic = True not supported. """ sel = self.sel.apply(group) ref = sel.center_of_geometry() kdtree = KDTree(dim=3, bucket_size=10) kdtree.set_coords(group.positions) kdtree.search(ref, self.exRadius) found_ExtIndices = kdtree.get_indices() kdtree.search(ref, self.inRadius) found_IntIndices = kdtree.get_indices() found_indices = list(set(found_ExtIndices) - set(found_IntIndices)) return unique(group[found_indices])
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 __init__(self, atom_list, bucket_size=10): """ o atom_list - list of atoms. This list is used in the queries. It can contain atoms from different structures. o bucket_size - bucket size of KD tree. You can play around with this to optimize speed if you feel like it. """ self.atom_list = atom_list # get the coordinates coord_list = [a.get_coord() for a in atom_list] # to Nx3 array of type float self.coords = numpy.array(coord_list).astype("f") assert (bucket_size > 1) assert (self.coords.shape[1] == 3) self.kdt = KDTree(3, bucket_size) self.kdt.set_coords(self.coords)
def __init__(self, atom_group, bucket_size=10): """ Parameters ---------- atom_list : AtomGroup list of atoms 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.kdtree = KDTree(dim=3, bucket_size=bucket_size) self.kdtree.set_coords(atom_group.positions)
def __init__(self, atom_group, bucket_size=10): """ :Arguments: *atom_list* list of atoms (:class: `~MDAnalysis.core.AtomGroup.AtomGroup`) *bucket_size* 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 if not hasattr(atom_group, 'coordinates'): raise TypeError('atom_group must have a coordinates() method' '(eq a AtomGroup from a selection)') self.kdtree = KDTree(dim=3, bucket_size=bucket_size) self.kdtree.set_coords(atom_group.coordinates())
def _apply_KDTree(self, group): """KDTree based selection is about 7x faster than distmat for typical problems. Limitations: always ignores periodicity """ sel = self.sel.apply(group) # All atoms in group that aren't in sel sys = group[~np.in1d(group.indices, sel.indices)] 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()) # These are the indices from SYS that were seen when # probing with SEL unique_idx = np.unique(np.concatenate(found_indices)) return unique(sys[unique_idx.astype(np.int32)])
def anal_jointdist(prot, rfirst, mols, cutoff, midz, box, mat, matburr, buried=None): """ Calculates the joint probability of residue interactions Parameters ---------- prot : AtomSelection the protein atoms rfirst : NumpyArray the first atom of each residue in the protein mols : NumpyArray the centroid of the molecule of interest cutoff : float the contact cut-off midz : float the middle of the bilayer box : NumpyArray the box sides mat : NumpyArray the joint probability for non-buried molecules matburr : NumpyArray the joint probability for buried molecules buried : NumpyArray of boolean, optional flag to indicate buried molecule """ imat = np.zeros(mat.shape) imatburr = np.zeros(matburr.shape) # Calculate all distances at once if box is not None: dist_all = distances.distance_array(np.array(mols), prot.get_positions(), box) else: kdtree = KDTree(dim=3, bucket_size=10) kdtree.set_coords(prot.get_positions()) # Loop over all molecules for mi, mol in enumerate(mols): if box is not None: dist = dist_all[mi, :] else: dist = np.ones(prot.get_positions().shape[0]) + cutoff kdtree.search(np.array([mol]), cutoff) for i in kdtree.get_indices(): dist[i] = 0.0 # Check if this molecule is on if dist.min() >= cutoff: continue # Check contacts for reach residue for i in range(len(rfirst) - 1): if dist[rfirst[i]:rfirst[i + 1]].min() >= cutoff: continue if (buried is not None and buried[mi]): imatburr[i, i] = 1 else: imat[i, i] = 1 for j in range(i + 1, len(rfirst) - 1): if dist[rfirst[j]:rfirst[j + 1]].min() >= cutoff: continue if (buried is not None and buried[mi]): imatburr[i, j] = 1 imatburr[j, i] = 1 else: imat[i, j] = 1 imat[j, i] = 1 return (mat + imat, matburr + imatburr)
def anal_contacts(prot, rfirst, mols, cutoff, midz, box, molfile, resfile, burresfile=None, buried=None, reslist=None, reslistfile=None): """ Calculates a range of contacts and write out contact vectors to files Parameters ---------- prot : AtomSelection the protein atoms rfirst : NumpyArray the first atom of each residue in the protein mols : NumpyArray the centroid of the molecule of interest cutoff : float the contact cut-off midz : float the middle of the bilayer box : NumpyArray the box sides molfile : fileobject the file to write molecular contacts to resfile : fileobject the file to write residue contacts to burresfile : fileobject, optional the file to write buried residue contacts to buried : NumpyArray of boolean, optional flag to indicate buried molecule reslist : list a list of residues to write out individual mol contacts reslistfile : fileobject the file to write out individual mol contacts to """ molon = np.zeros(len(mols), dtype=bool) reson = np.zeros(len(rfirst) - 1, dtype=bool) if buried is not None: burreson = np.zeros(len(rfirst) - 1, dtype=bool) if reslist is not None: resonlist = np.zeros([len(reslist), len(mols)], dtype=bool) # Calculate all distances at once if box is not None: dist_all = distances.distance_array(np.array(mols), prot.get_positions(), box) else: kdtree = KDTree(dim=3, bucket_size=10) kdtree.set_coords(prot.get_positions()) # Loop over all molecules for mi, mol in enumerate(mols): if box is not None: dist = dist_all[mi, :] else: dist = np.ones(prot.get_positions().shape[0]) + cutoff kdtree.search(np.array([mol]), cutoff) for i in kdtree.get_indices(): dist[i] = 0.0 # Check if this molecule is on molon[mi] = dist.min() < cutoff if molon[mi]: # Check contacts for reach residue for i in range(len(rfirst) - 1): if dist[rfirst[i]:rfirst[i + 1]].min() < cutoff: if (burresfile is not None and buried[mi]): burreson[i] = True else: reson[i] = True if reslist is not None and i in reslist: resonlist[reslist.index(i), mi] = True # Write state information to file write_booleans(molfile, molon) write_booleans(resfile, reson) if buried is not None: write_booleans(burresfile, burreson) if reslist is not None: write_booleans(reslistfile, resonlist.reshape(len(reslist) * len(mols)))