Updated RDF package to enable backend selection (#5038)

* Updated RDF package to enable backend selection

* removed some whitespace to keep flake8 happy

* black made some changes to distances.py

* Added check for nsgrid functions that don't use backend

* added tests for RDF backend

* fixed mypy errors

* updated AUTHORS and CHANGELOG for PR 5038

* moved distopia check function to common location

* updated CHANGELOG

* fix linting issue

* fixed line width, and removed duplicated entry

---------

Co-authored-by: Oliver Beckstein <orbeckst@gmail.com>

Author: Gareth-elliott

package/AUTHORS

@@ -257,6 +257,7 @@ Chronological list of authors
   - Debasish Mohanty  
   - Abdulrahman Elbanna
   - Tulga-Erdene Sodjargal
+  - Gareth Elliott
 
 
 External code

package/CHANGELOG

@@ -15,7 +15,7 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, BHM-Bob, TRY-ER, Abdulrahman-PROG, pbuslaev,
-         yuxuanzhuang, yuyuan871111, tanishy7777, tulga-rdn
+         yuxuanzhuang, yuyuan871111, tanishy7777, tulga-rdn, Gareth-elliott
 
 
  * 2.10.0
@@ -48,11 +48,18 @@ Enhancements
    (Issue #4948 #2874, PR #4965)
  * Enables parallelization for analysis.lineardensity.LinearDensity
    (Issue #4678, PR #5007)
- * Remove `default` channel from RTD conda env. (Issue # 5036, PR # 5037)
+ * Enabled selection of backend for distance calculations in RDF classes
+   (PR #5038)
+ * Modified functions in lib.distances to add the backend keyword,
+   so that it gets passed through from the calling functions and classes
+   (PR #5038)
+ * Moved distopia checking function to common import location in
+   MDAnalysisTest.util (PR #5038)
 
 Changes
  * Removed undocumented and unused attribute
    `analysis.lineardensity.LinearDensity.totalmass` (PR #5007)
+ * Remove `default` channel from RTD conda env. (Issue # 5036, PR # 5037)   
 
 Deprecations
 

package/MDAnalysis/analysis/rdf.py

@@ -131,6 +131,11 @@ class InterRDF(AnalysisBase):
         spatially correlated due to direct bonded connections.
     verbose : bool
         Show detailed progress of the calculation if set to `True`
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration of the distance calculations.
+        Default is 'serial'.
+
+        .. versionadded:: 2.10.0
 
     Attributes
     ----------
@@ -227,6 +232,7 @@ def __init__(
         norm="rdf",
         exclusion_block=None,
         exclude_same=None,
+        backend="serial",
         **kwargs,
     ):
         super(InterRDF, self).__init__(g1.universe.trajectory, **kwargs)
@@ -262,6 +268,8 @@ def __init__(
                 "Use 'rdf', 'density' or 'none'."
             )
 
+        self.backend = backend
+
     def _prepare(self):
         # Empty histogram to store the RDF
         count, edges = np.histogram([-1], **self.rdf_settings)
@@ -283,6 +291,7 @@ def _single_frame(self):
             self.g2.positions,
             self._maxrange,
             box=self._ts.dimensions,
+            backend=self.backend,
         )
         # Maybe exclude same molecule distances
         if self._exclusion_block is not None:
@@ -424,6 +433,12 @@ class InterRDF_s(AnalysisBase):
         .. deprecated:: 2.3.0
             Instead of `density=True` use `norm='density'`
 
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration of the distance calculations.
+        Default is 'serial'.
+
+        .. versionadded:: 2.10.0
+
     Attributes
     ----------
     results.bins : numpy.ndarray
@@ -571,6 +586,7 @@ def __init__(
         range=(0.0, 15.0),
         norm="rdf",
         density=False,
+        backend="serial",
         **kwargs,
     ):
         super(InterRDF_s, self).__init__(
@@ -603,6 +619,8 @@ def __init__(
             )
             self.norm = "density"
 
+        self.backend = backend
+
     def _prepare(self):
         count, edges = np.histogram([-1], **self.rdf_settings)
         self.results.count = [
@@ -624,6 +642,7 @@ def _single_frame(self):
                 ag2.positions,
                 self._maxrange,
                 box=self._ts.dimensions,
+                backend=self.backend,
             )
 
             for j, (idx1, idx2) in enumerate(pairs):

package/MDAnalysis/lib/distances.py

@@ -519,6 +519,7 @@ def capped_distance(
     box: Optional[npt.NDArray] = None,
     method: Optional[str] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Calculates pairs of indices corresponding to entries in the `reference`
     and `configuration` arrays which are separated by a distance lying within
@@ -557,6 +558,8 @@ def capped_distance(
         method.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -599,6 +602,9 @@ def capped_distance(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     if box is not None:
         box = np.asarray(box, dtype=np.float32)
@@ -622,14 +628,26 @@ def capped_distance(
         box=box,
         method=method,
     )
-    return function(
-        reference,
-        configuration,
-        max_cutoff,
-        min_cutoff=min_cutoff,
-        box=box,
-        return_distances=return_distances,
-    )
+
+    if function.__name__ == "_nsgrid_capped":
+        return function(
+            reference,
+            configuration,
+            max_cutoff,
+            min_cutoff=min_cutoff,
+            box=box,
+            return_distances=return_distances,
+        )
+    else:
+        return function(
+            reference,
+            configuration,
+            max_cutoff,
+            min_cutoff=min_cutoff,
+            box=box,
+            return_distances=return_distances,
+            backend=backend,
+        )
 
 
 def _determine_method(
@@ -727,6 +745,7 @@ def _bruteforce_capped(
     min_cutoff: Optional[float] = None,
     box: Optional[npt.NDArray] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Capped distance evaluations using a brute force method.
 
@@ -764,6 +783,8 @@ def _bruteforce_capped(
         ``[lx, ly, lz, alpha, beta, gamma]``.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -784,13 +805,18 @@ def _bruteforce_capped(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     # Default return values (will be overwritten only if pairs are found):
     pairs = np.empty((0, 2), dtype=np.intp)
     distances = np.empty((0,), dtype=np.float64)
 
     if len(reference) > 0 and len(configuration) > 0:
-        _distances = distance_array(reference, configuration, box=box)
+        _distances = distance_array(
+            reference, configuration, box=box, backend=backend
+        )
         if min_cutoff is not None:
             mask = np.where(
                 (_distances <= max_cutoff) & (_distances > min_cutoff)
@@ -823,6 +849,7 @@ def _pkdtree_capped(
     min_cutoff: Optional[float] = None,
     box: Optional[npt.NDArray] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Capped distance evaluations using a KDtree method.
 
@@ -860,6 +887,8 @@ def _pkdtree_capped(
         ``[lx, ly, lz, alpha, beta, gamma]``.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -880,6 +909,9 @@ def _pkdtree_capped(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     from .pkdtree import (
         PeriodicKDTree,
@@ -899,7 +931,10 @@ def _pkdtree_capped(
             if return_distances or (min_cutoff is not None):
                 refA, refB = pairs[:, 0], pairs[:, 1]
                 distances = calc_bonds(
-                    reference[refA], configuration[refB], box=box
+                    reference[refA],
+                    configuration[refB],
+                    box=box,
+                    backend=backend,
                 )
                 if min_cutoff is not None:
                     mask = np.where(distances > min_cutoff)
@@ -1044,6 +1079,7 @@ def self_capped_distance(
     box: Optional[npt.NDArray] = None,
     method: Optional[str] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Calculates pairs of indices corresponding to entries in the `reference`
     array which are separated by a distance lying within the specified
@@ -1078,6 +1114,8 @@ def self_capped_distance(
         method.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -1126,6 +1164,9 @@ def self_capped_distance(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     if box is not None:
         box = np.asarray(box, dtype=np.float32)
@@ -1146,13 +1187,24 @@ def self_capped_distance(
         box=box,
         method=method,
     )
-    return function(
-        reference,
-        max_cutoff,
-        min_cutoff=min_cutoff,
-        box=box,
-        return_distances=return_distances,
-    )
+
+    if function.__name__ == "_nsgrid_capped_self":
+        return function(
+            reference,
+            max_cutoff,
+            min_cutoff=min_cutoff,
+            box=box,
+            return_distances=return_distances,
+        )
+    else:
+        return function(
+            reference,
+            max_cutoff,
+            min_cutoff=min_cutoff,
+            box=box,
+            return_distances=return_distances,
+            backend=backend,
+        )
 
 
 def _determine_method_self(
@@ -1236,6 +1288,7 @@ def _bruteforce_capped_self(
     min_cutoff: Optional[float] = None,
     box: Optional[npt.NDArray] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Capped distance evaluations using a brute force method.
 
@@ -1266,6 +1319,8 @@ def _bruteforce_capped_self(
         ``[lx, ly, lz, alpha, beta, gamma]``.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -1287,6 +1342,9 @@ def _bruteforce_capped_self(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     # Default return values (will be overwritten only if pairs are found):
     pairs = np.empty((0, 2), dtype=np.intp)
@@ -1296,7 +1354,7 @@ def _bruteforce_capped_self(
     # We're searching within a single coordinate set, so we need at least two
     # coordinates to find distances between them.
     if N > 1:
-        distvec = self_distance_array(reference, box=box)
+        distvec = self_distance_array(reference, box=box, backend=backend)
         dist = np.full((N, N), np.finfo(np.float64).max, dtype=np.float64)
         dist[np.triu_indices(N, 1)] = distvec
 
@@ -1325,6 +1383,7 @@ def _pkdtree_capped_self(
     min_cutoff: Optional[float] = None,
     box: Optional[npt.NDArray] = None,
     return_distances: Optional[bool] = True,
+    backend: Optional[str] = "serial",
 ):
     """Capped distance evaluations using a KDtree method.
 
@@ -1355,6 +1414,8 @@ def _pkdtree_capped_self(
         ``[lx, ly, lz, alpha, beta, gamma]``.
     return_distances : bool, optional
         If set to ``True``, distances will also be returned.
+    backend : {'serial', 'OpenMP', 'distopia'}, optional
+        Keyword selecting the type of acceleration.
 
     Returns
     -------
@@ -1376,6 +1437,9 @@ def _pkdtree_capped_self(
     .. versionchanged:: 2.3.0
        Can now accept an :class:`~MDAnalysis.core.groups.AtomGroup` as an
        argument in any position and checks inputs using type hinting.
+    .. versionchanged:: 2.10.0
+       Added the "backend" argument to select the type of acceleration of
+       the distance calculations.
     """
     from .pkdtree import (
         PeriodicKDTree,
@@ -1397,7 +1461,7 @@ def _pkdtree_capped_self(
             if return_distances or (min_cutoff is not None):
                 refA, refB = pairs[:, 0], pairs[:, 1]
                 distances = calc_bonds(
-                    reference[refA], reference[refB], box=box
+                    reference[refA], reference[refB], box=box, backend=backend
                 )
                 if min_cutoff is not None:
                     idx = distances > min_cutoff

testsuite/MDAnalysisTests/analysis/test_rdf.py

@@ -24,6 +24,7 @@
 
 import MDAnalysis as mda
 from MDAnalysis.analysis.rdf import InterRDF
+from MDAnalysisTests.util import distopia_conditional_backend
 
 from MDAnalysisTests.datafiles import two_water_gro
 
@@ -152,3 +153,10 @@ def test_unknown_norm(sels):
     s1, s2 = sels
     with pytest.raises(ValueError, match="invalid norm"):
         InterRDF(s1, s2, sels, norm="foo")
+
+
+@pytest.mark.parametrize("backend", distopia_conditional_backend())
+def test_norm(sels, backend):
+    s1, s2 = sels
+    rdf = InterRDF(s1, s2, norm="none", backend=backend).run()
+    assert_allclose(max(rdf.results.rdf), 4)