From e3966303776577e15a043daeceff5a591370398a Mon Sep 17 00:00:00 2001 From: ayushsuhane <34154224+ayushsuhane@users.noreply.github.com> Date: Tue, 14 Aug 2018 07:13:00 -0700 Subject: [PATCH] Using Capped distance in distance selections (#2035) * modified capped function for faster computations, added search_tree in pkdtree, added tests, updated Changlog * updated docs, added cutoff in determine methods * safety condition added for bruteforce method * added _apply_nsgrid to selections * Removed dependency on kdtree flags, substituted capped_distance in class Distance based selection + point selection except CylindricalSelection * changed determine_method * minor correction for failed tests * updated CHANGELOG * minor flake8 corrections --- package/CHANGELOG | 4 + package/MDAnalysis/core/selection.py | 135 +++++------------- .../core/test_atomselections.py | 52 +++---- 3 files changed, 66 insertions(+), 125 deletions(-) diff --git a/package/CHANGELOG b/package/CHANGELOG index 112d6d33764..bc64f31b466 100644 --- a/package/CHANGELOG +++ b/package/CHANGELOG @@ -21,6 +21,10 @@ The rules for this file: Enhancements + * Replaced multiple apply (_apply_distmat, _apply_kdtree) + methods in distance based selections with + lib.distances.capped_distance for automatic selection of + distance evaluation method. (PR #2035) * Added return_distances argument in lib.distances.capped_distances to evaluate and return distances only when required. Modified the optimization rules in lib.distances._determine_method for diff --git a/package/MDAnalysis/core/selection.py b/package/MDAnalysis/core/selection.py index 4532ac854a7..bac71597b08 100644 --- a/package/MDAnalysis/core/selection.py +++ b/package/MDAnalysis/core/selection.py @@ -49,6 +49,7 @@ import numpy as np from numpy.lib.utils import deprecate + from MDAnalysis.lib.pkdtree import PeriodicKDTree from MDAnalysis.lib.util import unique_int_1d from MDAnalysis.core import flags @@ -237,10 +238,6 @@ class DistanceSelection(Selection): - _apply_distmat """ def __init__(self): - if flags['use_KDTree_routines'] in (True, 'fast', 'always'): - self.apply = self._apply_KDTree - else: - self.apply = self._apply_distmat self.periodic = flags['use_periodic_selections'] @@ -271,39 +268,23 @@ def __init__(self, parser, tokens): self.cutoff = float(tokens.popleft()) self.sel = parser.parse_expression(self.precedence) - def _apply_KDTree(self, group): - """KDTree based selection is about 7x faster than distmat - for typical problems. - """ + def apply(self, group): + indices = [] sel = self.sel.apply(group) # All atoms in group that aren't in sel sys = group[~np.in1d(group.indices, sel.indices)] - if not sys: + if not sys or not sel: return sys[[]] box = self.validate_dimensions(group.dimensions) + pairs = distances.capped_distance(sel.positions, sys.positions, + self.cutoff, box=box, + return_distances=False) + if pairs.size > 0: + indices = np.sort(pairs[:, 1]) - cut = self.cutoff if box is not None else None - kdtree = PeriodicKDTree(box=box, leafsize=10) - kdtree.set_coords(sys.positions, cutoff=cut) - kdtree.search(sel.positions, self.cutoff) - unique_idx = np.asarray(kdtree.get_indices()) - - return sys[unique_idx.astype(np.int32)].unique - - def _apply_distmat(self, group): - sel = self.sel.apply(group) - sys = group[~np.in1d(group.indices, sel.indices)] - - box = self.validate_dimensions(group.dimensions) - dist = distances.distance_array( - sys.positions, sel.positions, box) - - mask = (dist <= self.cutoff).any(axis=1) - - return sys[mask].unique - + return sys[np.asarray(indices, dtype=np.int64)].unique class SphericalLayerSelection(DistanceSelection): token = 'sphlayer' @@ -314,39 +295,22 @@ def __init__(self, parser, tokens): self.inRadius = float(tokens.popleft()) self.exRadius = float(tokens.popleft()) self.sel = parser.parse_expression(self.precedence) - - 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) - periodic = box is not None - ref = sel.center_of_geometry(pbc=periodic) - kdtree = PeriodicKDTree(box=box) - - cutoff = self.exRadius if box is not None else None - kdtree.set_coords(group.positions, cutoff=cutoff) - 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 group[found_indices].unique - - def _apply_distmat(self, group): + + def apply(self, group): + indices = [] sel = self.sel.apply(group) box = self.validate_dimensions(group.dimensions) periodic = box is not None ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).astype( np.float32) - d = distances.distance_array(ref, - group.positions, - box=box)[0] - mask = d < self.exRadius - mask &= d > self.inRadius + pairs = distances.capped_distance(ref, group.positions, self.exRadius, + min_cutoff=self.inRadius, + box=box, + return_distances=False) + if pairs.size > 0: + indices = np.sort(pairs[:, 1]) - return group[mask].unique + return group[np.asarray(indices, dtype=np.int64)].unique class SphericalZoneSelection(DistanceSelection): @@ -358,39 +322,27 @@ def __init__(self, parser, tokens): self.cutoff = float(tokens.popleft()) self.sel = parser.parse_expression(self.precedence) - def _apply_KDTree(self, group): - """Selection using KDTree and PeriodicKDTree for aperiodic and - fully-periodic systems, respectively. - """ + def apply(self, group): + indices = [] sel = self.sel.apply(group) box = self.validate_dimensions(group.dimensions) periodic = box is not None - ref = sel.center_of_geometry(pbc=periodic) + ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).astype( + np.float32) + pairs = distances.capped_distance(ref, group.positions, self.cutoff, + box=box, + return_distances=False) + if pairs.size > 0: + indices = np.sort(pairs[:, 1]) - cut = self.cutoff if box is not None else None - kdtree = PeriodicKDTree(box=box) - kdtree.set_coords(group.positions, cutoff=cut) - kdtree.search(ref, self.cutoff) - found_indices = kdtree.get_indices() - return group[found_indices].unique - - def _apply_distmat(self, group): - sel = self.sel.apply(group) - box = self.validate_dimensions(group.dimensions) - periodic = box is not None - ref = sel.center_of_geometry(pbc=periodic).reshape(1, 3).\ - astype(np.float32) - d = distances.distance_array(ref, - group.positions, - box=box)[0] - idx = d < self.cutoff - return group[idx].unique + return group[np.asarray(indices, dtype=np.int64)].unique class CylindricalSelection(Selection): def __init__(self): self.periodic = flags['use_periodic_selections'] + def apply(self, group): sel = self.sel.apply(group) @@ -484,25 +436,16 @@ def __init__(self, parser, tokens): self.ref = np.array([x, y, z], dtype=np.float32) self.cutoff = float(tokens.popleft()) - def _apply_KDTree(self, group): - box = self.validate_dimensions(group.dimensions) - kdtree = PeriodicKDTree(box=box) - cut = self.cutoff if box is not None else None - kdtree.set_coords(group.positions, cutoff=cut) - kdtree.search(self.ref, self.cutoff) - found_indices = kdtree.get_indices() - - return group[found_indices].unique - - def _apply_distmat(self, group): - ref_coor = self.ref[np.newaxis, ...] - - ref_coor = np.asarray(ref_coor, dtype=np.float32) + def apply(self, group): + indices = [] box = self.validate_dimensions(group.dimensions) + pairs = distances.capped_distance(self.ref[None, :], group.positions, self.cutoff, + box=box, + return_distances=False) + if pairs.size > 0: + indices = np.sort(pairs[:, 1]) - dist = distances.distance_array(group.positions, ref_coor, box) - mask = (dist <= self.cutoff).any(axis=1) - return group[mask].unique + return group[np.asarray(indices, dtype=np.int64)].unique class AtomSelection(Selection): diff --git a/testsuite/MDAnalysisTests/core/test_atomselections.py b/testsuite/MDAnalysisTests/core/test_atomselections.py index 33178097ca0..c60e094aee5 100644 --- a/testsuite/MDAnalysisTests/core/test_atomselections.py +++ b/testsuite/MDAnalysisTests/core/test_atomselections.py @@ -496,31 +496,14 @@ class BaseDistanceSelection(object): Cylindrical methods don't use KDTree """ - - methods = [('kdtree', False), - ('kdtree', True), - ('distmat', True), - ('distmat', False)] - - @staticmethod - def choosemeth(sel, meth, periodic): - """hack in the desired apply method""" - if meth == 'kdtree': - sel.apply = sel._apply_KDTree - elif meth == 'distmat': - sel.apply = sel._apply_distmat - + @pytest.mark.parametrize('periodic', (True, False)) + def test_around(self, u, periodic): + sel = Parser.parse('around 5.0 resid 1', u.atoms) if periodic: sel.periodic = True else: sel.periodic = False - return sel - - @pytest.mark.parametrize('meth, periodic', methods) - def test_around(self, u, meth, periodic): - sel = Parser.parse('around 5.0 resid 1', u.atoms) - sel = self.choosemeth(sel, meth, periodic) result = sel.apply(u.atoms) r1 = u.select_atoms('resid 1') @@ -535,10 +518,14 @@ def test_around(self, u, meth, periodic): ref.difference_update(set(r1.indices)) assert ref == set(result.indices) - @pytest.mark.parametrize('meth, periodic', methods) - def test_spherical_layer(self, u, meth, periodic): + @pytest.mark.parametrize('periodic', (True, False)) + def test_spherical_layer(self, u, periodic): sel = Parser.parse('sphlayer 2.4 6.0 resid 1', u.atoms) - sel = self.choosemeth(sel, meth, periodic) + if periodic: + sel.periodic = True + else: + sel.periodic = False + result = sel.apply(u.atoms) r1 = u.select_atoms('resid 1') @@ -549,10 +536,14 @@ def test_spherical_layer(self, u, meth, periodic): assert ref == set(result.indices) - @pytest.mark.parametrize('meth, periodic', methods) - def test_spherical_zone(self, u, meth, periodic): + @pytest.mark.parametrize('periodic', (True, False)) + def test_spherical_zone(self, u, periodic): sel = Parser.parse('sphzone 5.0 resid 1', u.atoms) - sel = self.choosemeth(sel, meth, periodic) + if periodic: + sel.periodic = True + else: + sel.periodic = False + result = sel.apply(u.atoms) r1 = u.select_atoms('resid 1') @@ -563,10 +554,13 @@ def test_spherical_zone(self, u, meth, periodic): assert ref == set(result.indices) - @pytest.mark.parametrize('meth, periodic', methods) - def test_point(self, u, meth, periodic): + @pytest.mark.parametrize('periodic', (True, False)) + def test_point(self, u, periodic): sel = Parser.parse('point 5.0 5.0 5.0 3.0', u.atoms) - sel = self.choosemeth(sel, meth, periodic) + if periodic: + sel.periodic = True + else: + sel.periodic = False result = sel.apply(u.atoms) box = u.dimensions if periodic else None