Optimize inverse index mapping in groups.py (Fixes Issue #3387)

package/AUTHORS

@@ -324,4 +324,4 @@ Logo
 
 The MDAnalysis 'Atom' logo was designed by Christian Beckstein; it is
 Copyright (c) 2011 Christian Beckstein and made available under a
-Creative Commons Attribution-NoDerivs 3.0 Unported License.
+Creative Commons Attribution-NoDerivs 3.0 Unported License.

package/CHANGELOG

@@ -45,6 +45,9 @@ Fixes
    DSSP by porting upstream PyDSSP 0.9.1 fix (Issue #4913)
 
 Enhancements
+ * Improved performance of inverse index mapping in AtomGroup using an optimized 
+   Cython implementation in lib._cutils.inverse_int_index()
+   (Issue #3387, PR #5252)
  * Added `select=None` in `analysis.rms.RMSD` to perform no selection on
    the input `atomgroup` and `reference` (Issue #5300, PR #5296)
  * MOL2Parser now reads unit cell dimensions from @<TRIPOS>CRYSIN records (Issue #3341)
@@ -3627,4 +3630,4 @@ Testsuite
     licenses
 
 11/12/07 naveen
-  * prepared for release outside lab
+  * prepared for release outside lab

package/MDAnalysis/core/groups.py

@@ -122,6 +122,7 @@
 from ..exceptions import NoDataError
 from . import topologyobjects
 from ._get_readers import get_writer_for, get_converter_for
+from ..lib._cutil import inverse_int_index
 
 
 def _unpickle(u, ix):
@@ -912,10 +913,7 @@ def _asunique(self, group, sorted=False, set_mask=False):
 
         indices = unique_int_1d_unsorted(self.ix)
         if set_mask:
-            mask = np.zeros_like(self.ix)
-            for i, x in enumerate(indices):
-                values = np.where(self.ix == x)[0]
-                mask[values] = i
+            mask = inverse_int_index(self.ix, indices)
             self._unique_restore_mask = mask
 
         issorted = int_array_is_sorted(indices)

package/MDAnalysis/lib/_cutil.pyx

@@ -37,7 +37,7 @@ from cython.operator cimport dereference as deref
 
 cnp.import_array()
 
-__all__ = ['unique_int_1d', 'make_whole', 'find_fragments',
+__all__ = ['unique_int_1d', 'inverse_int_index', 'make_whole', 'find_fragments',
            '_sarrus_det_single', '_sarrus_det_multiple']
 
 cdef extern from "calc_distances.h":
@@ -91,6 +91,42 @@ def unique_int_1d(cnp.intp_t[:] values):
 
     return np.array(result)
 
+@cython.boundscheck(False)
+@cython.wraparound(False)
+def inverse_int_index(cnp.intp_t[:] values,
+                      cnp.intp_t[:] unique_vals):
+    """
+    Construct inverse index map such that:
+
+        unique_vals[mask] == values
+
+    Parameters
+    ----------
+    values : numpy.ndarray
+        1D array of integers.
+    unique_vals : numpy.ndarray
+        1D array of unique integers (unsorted).
+
+    Returns
+    -------
+    numpy.ndarray
+        Integer mask mapping values -> index in unique_vals.
+    """
+
+    cdef Py_ssize_t n = values.shape[0]
+    cdef Py_ssize_t m = unique_vals.shape[0]
+    cdef Py_ssize_t i
+
+    cdef dict lookup = {}
+    cdef cnp.intp_t[:] mask = np.empty(n, dtype=np.intp)
+
+    for i in range(m):
+        lookup[unique_vals[i]] = i
+
+    for i in range(n):
+        mask[i] = lookup[values[i]]
+
+    return np.array(mask)
 
 @cython.boundscheck(False)
 def _in2d(cnp.intp_t[:, :] arr1, cnp.intp_t[:, :] arr2):
@@ -515,4 +551,4 @@ def find_fragments(atoms, bondlist):
         # Add fragment to output
         frags.append(np.asarray(this_frag))
 
-    return frags
+    return frags

testsuite/MDAnalysisTests/lib/test_cutil.py

@@ -28,6 +28,7 @@
     unique_int_1d,
     find_fragments,
     _in2d,
+    inverse_int_index,
 )
 
 
@@ -103,3 +104,52 @@ def test_in2d_VE(arr1, arr2):
         ValueError, match=r"Both arrays must be \(n, 2\) arrays"
     ):
         _in2d(arr1, arr2)
+
+
+def _python_reference_mask(ix, indices):
+    mask = np.zeros_like(ix)
+    for i, x in enumerate(indices):
+        values = np.where(ix == x)[0]
+        mask[values] = i
+    return mask
+
+
+@pytest.mark.parametrize(
+    "ix,indices",
+    [
+        # unsorted and not unique
+        (
+            np.array([1, 5, 3, 3, 6], dtype=np.intp),
+            np.array([1, 5, 3, 6], dtype=np.intp),
+        ),
+        # sorted and not unique
+        (
+            np.array([1, 3, 3, 5, 6], dtype=np.intp),
+            np.array([1, 3, 5, 6], dtype=np.intp),
+        ),
+        # unsorted and unique
+        (
+            np.array([1, 5, 3, 6], dtype=np.intp),
+            np.array([1, 5, 3, 6], dtype=np.intp),
+        ),
+        # sorted and unique
+        (
+            np.array([1, 3, 5, 6], dtype=np.intp),
+            np.array([1, 3, 5, 6], dtype=np.intp),
+        ),
+        # all elements identical
+        (
+            np.array([5, 5, 5], dtype=np.intp),
+            np.array([5], dtype=np.intp),
+        ),
+        # single element
+        (
+            np.array([7], dtype=np.intp),
+            np.array([7], dtype=np.intp),
+        ),
+    ],
+)
+def test_inverse_int_index(ix, indices):
+    pyref = _python_reference_mask(ix, indices)
+    cy = inverse_int_index(ix, indices)
+    assert_equal(pyref, cy)