perf: accelerate RAD neighbour search using NumPy sorting and Numba-compiled blocking loop

Author: harryswift01

CodeEntropy/levels/search.py

@@ -6,6 +6,90 @@
 
 import MDAnalysis as mda
 import numpy as np
+from numba import njit
+
+
+@njit
+def _rad_blocking_loop(i_coords, sorted_indices, sorted_distances, coms, dimensions):
+    """
+    Perform RAD neighbor selection using a blocking criterion.
+
+    This is a Numba-compiled implementation of the RAD algorithm, which
+    determines whether a molecule j is a neighbor of molecule i by checking
+    whether any closer molecule k blocks j based on angular and distance
+    relationships.
+
+    The criterion is based on:
+
+        (1 / r_ij)^2 > (1 / r_ik)^2 * cos(theta_jik)
+
+    where k blocks j if the inequality holds.
+
+    Args:
+        i_coords (np.ndarray):
+            Coordinates of the central molecule.
+        sorted_indices (np.ndarray):
+            Indices of molecules sorted by distance from i.
+        sorted_distances (np.ndarray):
+            Distances corresponding to sorted_indices.
+        coms (np.ndarray):
+            Precomputed center of mass coordinates for all molecules.
+        dimensions (np.ndarray):
+            Simulation box dimensions for periodic boundary conditions.
+
+    Returns:
+        list[int]:
+            Indices of molecules that belong to the RAD neighbor shell.
+    """
+    shell = []
+
+    n = sorted_indices.shape[0]
+    limit = min(n, 30)
+
+    for y in range(limit):
+        j_idx = sorted_indices[y]
+        r_ij = sorted_distances[y]
+        j_coords = coms[j_idx]
+
+        blocked = False
+
+        for z in range(y):
+            k_idx = sorted_indices[z]
+            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):
+                break
+
+            LHS = (1.0 / r_ij) ** 2
+            RHS = ((1.0 / r_ik) ** 2) * costheta
+
+            if LHS < RHS:
+                blocked = True
+                break
+
+        if not blocked:
+            shell.append(j_idx)
+
+    return shell
 
 
 class Search:
@@ -15,36 +99,44 @@ class Search:
 
     def __init__(self):
         """
-        Initializes the Search class with a placeholder for the system
-        trajectory.
+        Initialize the Search class.
+
+        This class currently serves as a container for neighbor search
+        methods operating on an MDAnalysis universe.
         """
         self._universe = None
         self._mol_id = None
 
     def _get_fragment_coms(self, universe):
         """
-        Precompute fragment centres of mass.
+        Precompute center of mass for each molecular fragment.
 
         Args:
-            universe: MDAnalysis universe object.
+            universe (MDAnalysis.Universe):
+                MDAnalysis universe object containing the system.
 
         Returns:
-            np.ndarray: Array of fragment COMs.
+            np.ndarray:
+                Array of shape (n_fragments, 3) containing COM coordinates.
         """
         return np.array([frag.center_of_mass() for frag in universe.atoms.fragments])
 
     def _get_distances(self, coms, i_coords, dimensions):
         """
-        Function to calculate distances between a central point and all COMs.
-        Takes periodic boundary conditions into account.
+        Compute distances between a central coordinate and all fragment COMs
+        using periodic boundary conditions.
 
         Args:
-            coms: array of fragment COMs
-            i_coords: coordinates of central molecule
-            dimensions: simulation box dimensions
+            coms (np.ndarray):
+                Array of fragment center of mass coordinates.
+            i_coords (np.ndarray):
+                Coordinates of the reference (central) molecule.
+            dimensions (np.ndarray):
+                Simulation box dimensions.
 
         Returns:
-            np.ndarray: distances to all molecules
+            np.ndarray:
+                Distances from the central molecule to all fragments.
         """
         delta = coms - i_coords
         delta = np.where(delta > 0.5 * dimensions, delta - dimensions, delta)
@@ -53,171 +145,78 @@ def _get_distances(self, coms, i_coords, dimensions):
 
     def get_RAD_neighbors(self, universe, mol_id):
         """
-        Find the neighbors of molecule with index mol_id.
+        Find RAD neighbors of a given molecule.
 
         Args:
-            universe: The MDAnalysis universe of the system.
-            mol_id (int): the index for the central molecule.
+            universe (MDAnalysis.Universe):
+                The MDAnalysis universe of the system.
+            mol_id (int):
+                Index of the central molecule.
 
         Returns:
-            neighbor_indices (list of ints): the list of neighboring molecule
-                indices.
+            list[int]:
+                Indices of neighboring molecules identified via the RAD method.
         """
         number_molecules = len(universe.atoms.fragments)
 
-        # Precompute COMs once
+        # Precompute COMs
         coms = self._get_fragment_coms(universe)
 
         # Central molecule position
         central_position = coms[mol_id]
 
-        # Compute all distances in one vectorised call
+        # Compute distances
         distances_array = self._get_distances(
             coms, central_position, universe.dimensions[:3]
         )
 
-        # Build distance dict excluding self
-        distances = {}
-        for molecule_index_j in range(number_molecules):
-            if molecule_index_j != mol_id:
-                distances[molecule_index_j] = distances_array[molecule_index_j]
+        # Prepare indices
+        indices = np.arange(number_molecules)
 
-        # Sort distances smallest to largest
-        sorted_dist = sorted(distances.items(), key=lambda item: item[1])
+        # Remove self
+        mask = indices != mol_id
+        filtered_indices = indices[mask]
+        filtered_distances = distances_array[mask]
 
-        # Get indices of neighbors
-        neighbor_indices = self._get_RAD_indices(
+        # Sort by distance
+        order = np.argsort(filtered_distances)
+
+        sorted_indices = filtered_indices[order]
+        sorted_distances = filtered_distances[order]
+
+        # RAD blocking (Numba)
+        neighbor_indices = _rad_blocking_loop(
             central_position,
-            sorted_dist,
+            sorted_indices,
+            sorted_distances,
             coms,
             universe.dimensions[:3],
-            number_molecules,
         )
 
         return neighbor_indices
 
-    def _get_RAD_indices(
-        self, i_coords, sorted_distances, coms, dimensions, number_molecules
-    ):
-        # pylint: disable=too-many-locals
-        r"""
-        For a given set of atom coordinates, find its RAD shell from the distance
-        sorted list, truncated to the closest 30 molecules.
-
-        This function calculates coordination shells using RAD the relative
-        angular distance, as defined first in DOI:10.1063/1.4961439
-        where molecules are defined as neighbors if
-        they fulfil the following condition:
-
-        .. math::
-            \Bigg(\frac{1}{r_{ij}}\Bigg)^2 >
-            \Bigg(\frac{1}{r_{ik}}\Bigg)^2 \cos \theta_{jik}
-
-        For a given particle :math:`i`, neighbor :math:`j` is in its coordination
-        shell if :math:`k` is not blocking particle :math:`j`. In this implementation
-        of RAD, we enforce symmetry, whereby neighboring particles must be in each
-        others coordination shells.
-
-        Args:
-            i_coords: xyz centre of mass of molecule :math:`i`
-            sorted_distances: list of index and distance pairs sorted by distance
-            coms: precomputed center of mass array
-            dimensions: system box dimensions
-            number_molecules: total number of molecules
-
-        Returns:
-            shell: list of indices of particles in the RAD shell of neighbors.
-        """
-        shell = []
-        count = -1
-        limit = min(number_molecules - 1, 30)
-
-        for y in range(limit):
-            count += 1
-
-            j_idx = sorted_distances[y][0]
-            r_ij = sorted_distances[y][1]
-            j_coords = coms[j_idx]
-
-            blocked = False
-
-            for z in range(count):
-                k_idx = sorted_distances[z][0]
-                r_ik = sorted_distances[z][1]
-                k_coords = coms[k_idx]
-
-                costheta_jik = self.get_angle(j_coords, i_coords, k_coords, dimensions)
-
-                if np.isnan(costheta_jik):
-                    break
-
-                LHS = (1 / r_ij) ** 2
-                RHS = ((1 / r_ik) ** 2) * costheta_jik
-
-                if LHS < RHS:
-                    blocked = True
-                    break
-
-            if not blocked:
-                shell.append(j_idx)
-
-        return shell
-
-    def get_angle(
-        self, a: np.ndarray, b: np.ndarray, c: np.ndarray, dimensions: np.ndarray
-    ):
-        """
-        Get the angle between three atoms, taking into account periodic
-        boundary conditions.
-
-        b is the vertex of the angle.
-
-        Args:
-            a: (3,) array of atom coordinates
-            b: (3,) array of atom coordinates
-            c: (3,) array of atom coordinates
-            dimensions: (3,) array of system box dimensions.
-
-        Returns:
-            cosine_angle: float, cosine of the angle abc.
-        """
-        ba = np.abs(a - b)
-        bc = np.abs(c - b)
-        ac = np.abs(c - a)
-
-        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)
-
-        dist_ba = np.sqrt((ba**2).sum(axis=-1))
-        dist_bc = np.sqrt((bc**2).sum(axis=-1))
-        dist_ac = np.sqrt((ac**2).sum(axis=-1))
-
-        cosine_angle = (dist_ac**2 - dist_bc**2 - dist_ba**2) / (-2 * dist_bc * dist_ba)
-
-        return cosine_angle
-
     def get_grid_neighbors(self, universe, mol_id, highest_level):
         """
-        Use MDAnalysis neighbor search to find neighbors.
-
-        For molecules with just one united atom, use the "A" search level to
-        find neighboring atoms. For larger molecules use the "R" search level
-        to find neighboring residues.
+        Find neighbors using MDAnalysis grid-based neighbor search.
 
-        The atoms/residues of the molecule of interest are removed from the
-        neighbor list.
+        For small molecules (united_atom), atom-level search is used.
+        For larger molecules, residue-level search is used.
 
         Args:
-            universe: MDAnalysis universe object for system.
-            mol_id: int, the index for the molecule of interest.
-            highest_level: str, molecule level.
+            universe (MDAnalysis.Universe):
+                MDAnalysis universe object for the system.
+            mol_id (int):
+                Index of the molecule of interest.
+            highest_level (str):
+                Molecule level ("united_atom" or other).
 
         Returns:
-            neighbors: MDAnalysis atomgroup of the neighboring particles.
+            list[int]:
+                Fragment indices of neighboring molecules.
         """
         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)
 

conda-recipe/meta.yaml

@@ -39,6 +39,7 @@ requirements:
     - distributed >=2026.1.2,<2026.2.0
     - dask-jobqueue >=0.9,<0.10
     - pytest-xdist >=3.8, <3.9
+    - numba >=0.64.0, <0.7
 
 test:
   imports:

pyproject.toml

@@ -50,6 +50,7 @@ dependencies = [
     "waterEntropy>=2,<2.3",
     "requests>=2.32,<3.0",
     "rdkit>=2025.9.5",
+    "numba>=0.65.0,<0.7",
 ]
 
 [project.urls]