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

@@ -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

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

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