Improve performance of wrap method (#5226)

* Faster wrap (see PR #5226 for details and benchmark data)
* Changelog and linting
* Add test
* Add benchmarks

Author: hejamu

benchmarks/benchmarks/ag_methods.py

@@ -189,6 +189,11 @@ def time_wrap(self, num_atoms):
         """
         self.ag.wrap()
 
+    def time_wrap_compound(self, num_atoms):
+        """Benchmark wrap() operation on
+        atomgroup with default params.
+        """
+        self.ag.wrap(compound="residues")
 
 class AtomGroupAttrsBench(object):
     """Benchmarks for the various MDAnalysis

package/CHANGELOG

@@ -15,7 +15,7 @@ The rules for this file:
 
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, orbeckst, marinegor, tylerjereddy, ljwoods2, marinegor,
-         spyke7, talagayev, tanii1125, BradyAJohnston
+         spyke7, talagayev, tanii1125, BradyAJohnston, hejamu
 
  * 2.11.0
 
@@ -37,6 +37,7 @@ Fixes
    DSSP by porting upstream PyDSSP 0.9.1 fix (Issue #4913)
 
 Enhancements
+ * Improved performance of `AtomGroup.wrap()` with compounds (PR #5220)
  * Adds support for parsing `.tpr` files produced by GROMACS 2026.0
  * Enables parallelization for analysis.diffusionmap.DistanceMatrix
    (Issue #4679, PR #4745)

package/MDAnalysis/core/groups.py

@@ -1899,13 +1899,17 @@ def wrap(self, compound="atoms", center="com", box=None, inplace=True):
             else:
                 compound_indices = atoms._get_compound_indices(comp)
 
-                # apply the shifts:
+                # Build mapping from compound index to shift index
                 unique_compound_indices = unique_int_1d(compound_indices)
-                shift_idx = 0
-                for i in unique_compound_indices:
-                    mask = np.where(compound_indices == i)
-                    positions[mask] += shifts[shift_idx]
-                    shift_idx += 1
+                index_to_shift = np.empty(
+                    compound_indices.max() + 1, dtype=int
+                )
+                index_to_shift[unique_compound_indices] = np.arange(
+                    len(unique_compound_indices)
+                )
+
+                # Apply shifts
+                positions += shifts[index_to_shift[compound_indices]]
 
         if inplace:
             atoms.positions = positions

testsuite/MDAnalysisTests/core/test_wrap.py

@@ -444,6 +444,30 @@ def test_wrap_no_molnums_exception_safety(self, level, compound, center):
                 group.atoms.positions, ref_wrapped_pos, decimal=self.precision
             )
 
+    @pytest.mark.parametrize(
+        "compound",
+        ("segments", "residues", "molecules", "fragments"),
+    )
+    @pytest.mark.parametrize("center", ("com", "cog"))
+    @pytest.mark.parametrize("is_triclinic", (False, True))
+    def test_wrap_shuffled_index(self, compound, center, is_triclinic):
+        # get a pristine test universe
+        u = UnWrapUniverse(is_triclinic=is_triclinic)
+        # shuffle atom order to make compound indices non-contiguous
+        rng = np.random.RandomState(1234)
+        shuffle = rng.permutation(u.atoms.n_atoms)
+        group = u.atoms[shuffle]
+        # get expected wrapped coordinates in shuffled order
+        ref_wrapped_pos = u.wrapped_coords(compound, center)[shuffle]
+        # wrap the shuffled group:
+        wrapped_pos = group.wrap(
+            compound=compound, center=center, inplace=False
+        )
+        # check for correct result:
+        assert_almost_equal(
+            wrapped_pos, ref_wrapped_pos, decimal=self.precision
+        )
+
 
 class TestWrapTRZ(object):
     """Tests the functionality of AtomGroup.wrap() using a TRZ universe."""