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

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