Refactored next_neighbors and added type hints

src/mdevaluate/atoms.py

@@ -10,7 +10,7 @@ from .checksum import checksum
 from .coordinates import CoordinateFrame
 
 
-def compare_regex(str_list: list[str], exp: str):
+def compare_regex(str_list: list[str], exp: str) -> np.ndarray:
     """
     Compare a list of strings with a regular expression.
     """
@@ -79,10 +79,10 @@ class AtomSubset:
         the keyworss below. Names are matched as a regular expression with `re.match`.
 
         Args:
-            atom_name:      Specification of the atom name
-            residue_name:   Specification of the resiude name
-            residue_id:     Residue ID or list of IDs
-            indices:        List of atom indices
+            atom_name: Specification of the atom name
+            residue_name: Specification of the resiude name
+            residue_id: Residue ID or list of IDs
+            indices: List of atom indices
         """
         new_subset = self
         if atom_name is not None:
@@ -176,7 +176,7 @@ class AtomSubset:
         return checksum(self.description)
 
 
-def gyration_radius(position: CoordinateFrame):
+def gyration_radius(position: CoordinateFrame) -> np.ndarray:
     r"""
     Calculates a list of all radii of gyration of all molecules given in the coordinate
     frame, weighted with the masses of the individual atoms.
@@ -185,7 +185,8 @@ def gyration_radius(position: CoordinateFrame):
         position: Coordinate frame object
 
     ..math::
-        R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i} - \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i }
+        R_G = \left(\frac{\sum_{i=1}^{n} m_i |\vec{r_i}
+              - \vec{r_{COM}}|^2 }{\sum_{i=1}^{n} m_i }
         \rigth)^{\frac{1}{2}}
     """
     gyration_radii = np.array([])
@@ -207,7 +208,7 @@ def layer_of_atoms(
     thickness: float,
     plane_normal: npt.ArrayLike,
     plane_offset: Optional[npt.ArrayLike] = np.array([0, 0, 0]),
-):
+) -> np.array:
     if plane_offset is None:
         np.array([0, 0, 0])
     atoms = atoms - plane_offset
@@ -215,45 +216,14 @@ def layer_of_atoms(
     return np.abs(distance) <= thickness
 
 
-def distance_to_atoms(ref, atoms, box=None):
-    """Get the minimal distance from atoms to ref.
-    The result is an array of with length == len(atoms)
-    """
-    out = np.empty(atoms.shape[0])
-    for i, atom in enumerate(atoms):
-        diff = (pbc_diff(atom, ref, box) ** 2).sum(axis=1).min()
-        out[i] = np.sqrt(diff)
-    return out
-
-
-def distance_to_atoms_KDtree(ref, atoms, box=None, thickness=None):
-    """
-    Get the minimal distance from atoms to ref.
-    The result is an array of with length == len(atoms)
-    Can be faster than distance_to_atoms.
-    Thickness defaults to box/5. If this is too small results may be wrong.
-    If box is not given then periodic boundary conditions are not applied!
-    """
-    if thickness == None:
-        thickness = box / 5
-    if box is not None:
-        start_coords = np.copy(atoms) % box
-        all_frame_coords = pbc_points(ref, box, thickness=thickness)
-    else:
-        start_coords = atoms
-        all_frame_coords = ref
-
-    tree = KDTree(all_frame_coords)
-    return tree.query(start_coords)[0]
-
-
 def next_neighbors(
-    atoms,
-    query_atoms=None,
-    number_of_neighbors=1,
-    distance_upper_bound=np.inf,
-    distinct=False,
-):
+    atoms: CoordinateFrame,
+    query_atoms: Optional[CoordinateFrame] = None,
+    number_of_neighbors: int = 1,
+    distance_upper_bound: float = np.inf,
+    distinct: bool = False,
+    **kwargs
+) -> (np.ndarray, np.ndarray):
     """
     Find the N next neighbors of a set of atoms.
 
@@ -269,15 +239,29 @@ def next_neighbors(
             If this is true, the atoms and query atoms are taken as distinct sets of
             atoms
     """
-    tree = KDTree(atoms)
     dnn = 0
     if query_atoms is None:
         query_atoms = atoms
+        dnn = 1
     elif not distinct:
         dnn = 1
-    dist, indices = tree.query(
-        query_atoms,
-        number_of_neighbors + dnn,
-        distance_upper_bound=distance_upper_bound,
-    )
-    return indices[:, dnn:]
+
+    box = atoms.box
+    if np.all(np.diag(np.diag(box)) == box):
+        atoms = atoms % np.diag(box)
+        tree = KDTree(atoms, boxsize=np.diag(box))
+        distances, indices = tree.query(
+            query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound
+        )
+    else:
+        atoms_pbc, atoms_pbc_index = pbc_points(
+            query_atoms, box, thickness=distance_upper_bound+0.1, index=True, **kwargs
+        )
+        tree = KDTree(atoms_pbc)
+        distances, indices = tree.query(
+            query_atoms, number_of_neighbors, distance_upper_bound=distance_upper_bound
+        )
+        indices = atoms_pbc_index[indices]
+
+    return distances[:, dnn:], indices[:, dnn:]
+