Replace distance array with capped distances

pmda/rdf.py

@@ -30,8 +30,7 @@
 
 import numpy as np
 
-from MDAnalysis.lib.distances import distance_array
-from MDAnalysis.lib.util import blocks_of
+from MDAnalysis.lib import distances
 
 from .parallel import ParallelAnalysisBase
 
@@ -117,23 +116,24 @@ def _prepare(self):
 
         # Need to know average volume
         self.volume = 0.0
+        # Set the max range to filter the search radius
+        self._maxrange = self.rdf_settings['range'][1]
 
     def _single_frame(self, ts, atomgroups):
         g1, g2 = atomgroups
         u = g1.universe
-        d = distance_array(g1.positions, g2.positions,
-                           box=u.dimensions)
+        pairs, dist = distances.capped_distance(g1.positions,
+                                                g2.positions,
+                                                self._maxrange,
+                                                box=u.dimensions)
         # If provided exclusions, create a mask of _result which
-        # lets us take these out
-        exclusion_mask = None
+        # lets us take these out.
         if self._exclusion_block is not None:
-            exclusion_mask = blocks_of(d, *self._exclusion_block)
-            maxrange = self.rdf_settings['range'][1] + 1.0
-        # Maybe exclude same molecule distances
-        if exclusion_mask is not None:
-            exclusion_mask[:] = maxrange
-        count = []
-        count = np.histogram(d, **self.rdf_settings)[0]
+            idxA = pairs[:, 0]//self._exclusion_block[0]
+            idxB = pairs[:, 1]//self._exclusion_block[1]
+            mask = np.where(idxA != idxB)[0]
+            dist = dist[mask]
+        count = np.histogram(dist, **self.rdf_settings)[0]
         volume = u.trajectory.ts.volume
 
         return np.array([count, np.array(volume, dtype=np.float64)])