perf: make RAD neighbour search frame-safe by caching COMs per frame and enforcing deterministic ordering to restore identical results

Author: harryswift01

CodeEntropy/levels/search.py

@@ -58,22 +58,18 @@ def _rad_blocking_loop(i_coords, sorted_indices, sorted_distances, coms, dimensi
             r_ik = sorted_distances[z]
             k_coords = coms[k_idx]
 
-            # Compute coordinate differences
             ba = np.abs(j_coords - i_coords)
             bc = np.abs(k_coords - i_coords)
             ac = np.abs(k_coords - j_coords)
 
-            # Apply periodic boundary conditions
             ba = np.where(ba > 0.5 * dimensions, ba - dimensions, ba)
             bc = np.where(bc > 0.5 * dimensions, bc - dimensions, bc)
             ac = np.where(ac > 0.5 * dimensions, ac - dimensions, ac)
 
-            # Compute distances
             dist_ba = np.sqrt((ba**2).sum())
             dist_bc = np.sqrt((bc**2).sum())
             dist_ac = np.sqrt((ac**2).sum())
 
-            # Cosine of angle jik
             costheta = (dist_ac**2 - dist_bc**2 - dist_ba**2) / (-2 * dist_bc * dist_ba)
 
             if np.isnan(costheta):
@@ -101,25 +97,39 @@ def __init__(self):
         """
         Initialize the Search class.
 
-        This class currently serves as a container for neighbor search
-        methods operating on an MDAnalysis universe.
+        This class includes frame-safe caching of fragment COMs and
+        system dimensions to avoid recomputation while preserving
+        identical results to the original implementation.
         """
-        self._universe = None
-        self._mol_id = None
+        self._cached_frame = None
+        self._cached_fragments = None
+        self._cached_coms = None
+        self._cached_dimensions = None
 
-    def _get_fragment_coms(self, universe):
+    def _update_cache(self, universe):
         """
-        Precompute center of mass for each molecular fragment.
+        Update cached MDAnalysis data if the simulation frame has changed.
 
         Args:
             universe (MDAnalysis.Universe):
                 MDAnalysis universe object containing the system.
-
-        Returns:
-            np.ndarray:
-                Array of shape (n_fragments, 3) containing COM coordinates.
         """
-        return np.array([frag.center_of_mass() for frag in universe.atoms.fragments])
+        # Get current frame index (MDAnalysis trajectory)
+        current_frame = universe.trajectory.ts.frame
+
+        # Only recompute if frame has changed
+        if self._cached_frame == current_frame:
+            return
+
+        fragments = universe.atoms.fragments
+
+        # Compute COMs once per frame (deterministic snapshot)
+        coms = np.array([frag.center_of_mass() for frag in fragments])
+
+        self._cached_fragments = fragments
+        self._cached_coms = coms
+        self._cached_dimensions = universe.dimensions[:3]
+        self._cached_frame = current_frame
 
     def _get_distances(self, coms, i_coords, dimensions):
         """
@@ -157,40 +167,38 @@ def get_RAD_neighbors(self, universe, mol_id):
             list[int]:
                 Indices of neighboring molecules identified via the RAD method.
         """
-        number_molecules = len(universe.atoms.fragments)
+        # Ensure cache corresponds to current frame
+        self._update_cache(universe)
+
+        fragments = self._cached_fragments
+        coms = self._cached_coms
+        dimensions = self._cached_dimensions
 
-        # Precompute COMs
-        coms = self._get_fragment_coms(universe)
+        number_molecules = len(fragments)
 
-        # Central molecule position
         central_position = coms[mol_id]
 
-        # Compute distances
-        distances_array = self._get_distances(
-            coms, central_position, universe.dimensions[:3]
-        )
+        # Distances computed from same COM snapshot
+        distances_array = self._get_distances(coms, central_position, dimensions)
 
-        # Prepare indices
         indices = np.arange(number_molecules)
 
-        # Remove self
         mask = indices != mol_id
         filtered_indices = indices[mask]
         filtered_distances = distances_array[mask]
 
-        # Sort by distance
-        order = np.argsort(filtered_distances)
+        # Stable sort to avoid ordering ambiguity
+        order = np.argsort(filtered_distances, kind="mergesort")
 
         sorted_indices = filtered_indices[order]
         sorted_distances = filtered_distances[order]
 
-        # RAD blocking (Numba)
         neighbor_indices = _rad_blocking_loop(
             central_position,
             sorted_indices,
             sorted_distances,
             coms,
-            universe.dimensions[:3],
+            dimensions,
         )
 
         return neighbor_indices
@@ -214,8 +222,10 @@ def get_grid_neighbors(self, universe, mol_id, highest_level):
             list[int]:
                 Fragment indices of neighboring molecules.
         """
+        fragments = universe.atoms.fragments
+        fragment = fragments[mol_id]
+
         search_object = mda.lib.NeighborSearch.AtomNeighborSearch(universe.atoms)
-        fragment = universe.atoms.fragments[mol_id]
 
         selection_string = f"index {fragment.indices[0]}:{fragment.indices[-1]}"
         molecule_atom_group = universe.select_atoms(selection_string)