Adjusted type hints

src/mdevaluate/coordinates.py

@@ -1,10 +1,10 @@
 from functools import partial, wraps
 from copy import copy
 from .logging import logger
-from typing import Optional, Callable
+from typing import Optional, Callable, List, Tuple
 
 import numpy as np
-import numpy.typing as npt
+from numpy.typing import ArrayLike, NDArray
 from scipy.spatial import KDTree
 
 from .atoms import AtomSubset
@@ -17,7 +17,7 @@ class UnknownCoordinatesMode(Exception):
     pass
 
 
-class CoordinateFrame(np.ndarray):
+class CoordinateFrame(NDArray):
     _known_modes = ("pbc", "whole", "nojump")
 
     @property
@@ -99,7 +99,7 @@ class CoordinateFrame(np.ndarray):
         box=None,
         mode=None,
     ):
-        obj = np.ndarray.__new__(subtype, shape, dtype, buffer, offset, strides)
+        obj = NDArray.__new__(subtype, shape, dtype, buffer, offset, strides)
 
         obj.coordinates = coordinates
         obj.step = step
@@ -319,7 +319,7 @@ class CoordinatesMap:
         return CoordinatesMap(self.coordinates.pbc, self.function)
 
 
-def rotate_axis(coords: npt.ArrayLike, axis: npt.ArrayLike) -> np.ndarray:
+def rotate_axis(coords: ArrayLike, axis: ArrayLike) -> NDArray:
     """
     Rotate a set of coordinates to a given axis.
     """
@@ -352,8 +352,8 @@ def rotate_axis(coords: npt.ArrayLike, axis: npt.ArrayLike) -> np.ndarray:
 
 
 def spherical_radius(
-    frame: CoordinateFrame, origin: Optional[npt.ArrayLike] = None
-) -> np.ndarray:
+    frame: CoordinateFrame, origin: Optional[ArrayLike] = None
+) -> NDArray:
     """
     Transform a frame of cartesian coordinates into the spherical radius.
     If origin=None, the center of the box is taken as the coordinates' origin.
@@ -363,7 +363,7 @@ def spherical_radius(
     return ((frame - origin) ** 2).sum(axis=-1) ** 0.5
 
 
-def polar_coordinates(x: npt.ArrayLike, y: npt.ArrayLike) -> (np.ndarray, np.ndarray):
+def polar_coordinates(x: ArrayLike, y: ArrayLike) -> (NDArray, NDArray):
     """Convert cartesian to polar coordinates."""
     radius = (x**2 + y**2) ** 0.5
     phi = np.arctan2(y, x)
@@ -371,8 +371,8 @@ def polar_coordinates(x: npt.ArrayLike, y: npt.ArrayLike) -> (np.ndarray, np.nda
 
 
 def spherical_coordinates(
-    x: npt.ArrayLike, y: npt.ArrayLike, z: npt.ArrayLike
-) -> (np.ndarray, np.ndarray, np.ndarray):
+    x: ArrayLike, y: ArrayLike, z: ArrayLike
+) -> (NDArray, NDArray, NDArray):
     """Convert cartesian to spherical coordinates."""
     xy, phi = polar_coordinates(x, y)
     radius = (x**2 + y**2 + z**2) ** 0.5
@@ -384,8 +384,8 @@ def selector_radial_cylindrical(
     atoms: CoordinateFrame,
     r_min: float,
     r_max: float,
-    origin: Optional[npt.ArrayLike] = None,
-) -> np.ndarray:
+    origin: Optional[ArrayLike] = None,
+) -> NDArray:
     box = atoms.box
     atoms = atoms % np.diag(box)
     if origin is None:
@@ -397,7 +397,7 @@ def selector_radial_cylindrical(
 
 
 def map_coordinates(
-    func: Callable[[CoordinateFrame, ...], np.ndarray]
+    func: Callable[[CoordinateFrame, ...], NDArray]
 ) -> Callable[..., CoordinatesMap]:
     @wraps(func)
     def wrapped(coordinates: Coordinates, **kwargs) -> CoordinatesMap:
@@ -408,14 +408,14 @@ def map_coordinates(
 
 @map_coordinates
 def center_of_masses(
-    frame: CoordinateFrame, atoms=None, shear: bool = False
-) -> np.ndarray:
-    if atoms is None:
-        atoms = list(range(len(frame)))
-    res_ids = frame.residue_ids[atoms]
-    masses = frame.masses[atoms]
+    frame: CoordinateFrame, atom_indices=None, shear: bool = False
+) -> NDArray:
+    if atom_indices is None:
+        atom_indices = list(range(len(frame)))
+    res_ids = frame.residue_ids[atom_indices]
+    masses = frame.masses[atom_indices]
     if shear:
-        coords = frame[atoms]
+        coords = frame[atom_indices]
         box = frame.box
         sort_ind = res_ids.argsort(kind="stable")
         i = np.concatenate([[0], np.where(np.diff(res_ids[sort_ind]) > 0)[0] + 1])
@@ -423,7 +423,7 @@ def center_of_masses(
         cor = pbc_diff(coords, coms, box)
         coords = coms + cor
     else:
-        coords = frame.whole[atoms]
+        coords = frame.whole[atom_indices]
     mask = np.bincount(res_ids)[1:] != 0
     positions = np.array(
         [
@@ -437,8 +437,8 @@ def center_of_masses(
 
 @map_coordinates
 def pore_coordinates(
-    frame: CoordinateFrame, origin: npt.ArrayLike, sym_axis: str = "z"
-) -> np.ndarray:
+    frame: CoordinateFrame, origin: ArrayLike, sym_axis: str = "z"
+) -> NDArray:
     """
     Map coordinates of a pore simulation so the pore has cylindrical symmetry.
 
@@ -459,17 +459,17 @@ def pore_coordinates(
 @map_coordinates
 def vectors(
     frame: CoordinateFrame,
-    atoms_indices_a: npt.ArrayLike,
-    atoms_indices_b: npt.ArrayLike,
+    atom_indices_a: ArrayLike,
+    atom_indices_b: ArrayLike,
     normed: bool = False,
-) -> np.ndarray:
+) -> NDArray:
     """
     Compute the vectors between the atoms of two subsets.
 
     Args:
         frame: The Coordinates object the atoms will be taken from
-        atoms_indices_a: Mask or indices of the first atom subset
-        atoms_indices_b: Mask or indices of the second atom subset
+        atom_indices_a: Mask or indices of the first atom subset
+        atom_indices_b: Mask or indices of the second atom subset
         normed (opt.): If the vectors should be normed
 
     The definition of atoms_a/b can be any possible subript of a numpy array.
@@ -492,10 +492,10 @@ def vectors(
         ])
     """
     box = frame.box
-    coords_a = frame[atoms_indices_a]
+    coords_a = frame[atom_indices_a]
     if len(coords_a.shape) > 2:
         coords_a = coords_a.mean(axis=0)
-    coords_b = frame[atoms_indices_b]
+    coords_b = frame[atom_indices_b]
     if len(coords_b.shape) > 2:
         coords_b = coords_b.mean(axis=0)
     vec = pbc_diff(coords_a, coords_b, box=box)
@@ -507,8 +507,8 @@ def vectors(
 
 @map_coordinates
 def dipole_vector(
-    frame: CoordinateFrame, atom_indices: npt.ArrayLike, normed: bool = None
-) -> np.ndarray:
+    frame: CoordinateFrame, atom_indices: ArrayLike, normed: bool = None
+) -> NDArray:
     coords = frame.whole[atom_indices]
     res_ids = frame.residue_ids[atom_indices]
     charges = frame.charges[atom_indices]
@@ -525,9 +525,9 @@ def dipole_vector(
 @map_coordinates
 def sum_dipole_vector(
     coordinates: CoordinateFrame,
-    atom_indices: npt.ArrayLike,
+    atom_indices: ArrayLike,
     normed: bool = True,
-) -> np.ndarray:
+) -> NDArray:
     coords = coordinates.whole[atom_indices]
     charges = coordinates.charges[atom_indices]
     dipole = np.array([c * charges for c in coords.T]).T
@@ -539,11 +539,11 @@ def sum_dipole_vector(
 @map_coordinates
 def normal_vectors(
     frame: CoordinateFrame,
-    atom_indices_a: npt.ArrayLike,
-    atom_indices_b: npt.ArrayLike,
-    atom_indices_c: npt.ArrayLike,
+    atom_indices_a: ArrayLike,
+    atom_indices_b: ArrayLike,
+    atom_indices_c: ArrayLike,
     normed: bool = True,
-) -> np.ndarray:
+) -> NDArray:
     coords_a = frame[atom_indices_a]
     coords_b = frame[atom_indices_b]
     coords_c = frame[atom_indices_c]
@@ -571,8 +571,8 @@ def displacements_without_drift(
 
 @map_coordinates
 def cylindrical_coordinates(
-    frame: CoordinateFrame, origin: npt.ArrayLike = None
-) -> np.ndarray:
+    frame: CoordinateFrame, origin: ArrayLike = None
+) -> NDArray:
     if origin is None:
         origin = np.diag(frame.box) / 2
     x = frame[:, 0] - origin[0]
@@ -584,10 +584,10 @@ def cylindrical_coordinates(
 
 
 def layer_of_atoms(
-        atoms: CoordinateFrame,
-        thickness: float,
-        plane_normal: npt.ArrayLike,
-        plane_offset: Optional[npt.ArrayLike] = np.array([0, 0, 0]),
+    atoms: CoordinateFrame,
+    thickness: float,
+    plane_normal: ArrayLike,
+    plane_offset: Optional[ArrayLike] = np.array([0, 0, 0]),
 ) -> np.array:
     if plane_offset is None:
         np.array([0, 0, 0])
@@ -597,13 +597,13 @@ def layer_of_atoms(
 
 
 def next_neighbors(
-        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):
+    atoms: CoordinateFrame,
+    query_atoms: Optional[CoordinateFrame] = None,
+    number_of_neighbors: int = 1,
+    distance_upper_bound: float = np.inf,
+    distinct: bool = False,
+    **kwargs
+) -> Tuple[List, List]:
     """
     Find the N next neighbors of a set of atoms.
 
@@ -634,7 +634,15 @@ def next_neighbors(
             query_atoms,
             number_of_neighbors + dnn,
             distance_upper_bound=distance_upper_bound,
-            )
+        )
+        distances = distances[:, dnn:]
+        indices = indices[:, dnn:]
+        distances_new = []
+        indices_new = []
+        for dist, ind in zip(distances, indices):
+            distances_new.append(dist[dist <= distance_upper_bound])
+            indices_new.append(ind[dist <= distance_upper_bound])
+        return distances_new, indices_new
     else:
         atoms_pbc, atoms_pbc_index = pbc_points(
             query_atoms, box, thickness=distance_upper_bound + 0.1, index=True, **kwargs
@@ -644,7 +652,52 @@ def next_neighbors(
             query_atoms,
             number_of_neighbors + dnn,
             distance_upper_bound=distance_upper_bound,
-            )
-        indices = atoms_pbc_index[indices]
+        )
+        distances = distances[:, dnn:]
+        indices = indices[:, dnn:]
+        distances_new = []
+        indices_new = []
+        for dist, ind in zip(distances, indices):
+            distances_new.append(dist[dist <= distance_upper_bound])
+            indices_new.append(atoms_pbc_index[ind[dist <= distance_upper_bound]])
+        return distances_new, indices_new
 
-    return distances[:, dnn:], indices[:, dnn:]
+
+def number_of_neighbors(
+    atoms: CoordinateFrame,
+    query_atoms: Optional[CoordinateFrame] = None,
+    r_max: float = 1,
+    distinct: bool = False,
+    **kwargs
+) -> Tuple[List, List]:
+    """
+    Find the N next neighbors of a set of atoms.
+
+    Args:
+        atoms:
+            The reference atoms and also the atoms which are queried if `query_atoms`
+            is net provided
+        query_atoms (opt.): If this is not None, these atoms will be queried
+        r_max (float, opt.):
+            Upper bound of the distance between neighbors
+        distinct (bool, opt.):
+            If this is true, the atoms and query atoms are taken as distinct sets of
+            atoms
+    """
+    dnn = 0
+    if query_atoms is None:
+        query_atoms = atoms
+        dnn = 1
+    elif not distinct:
+        dnn = 1
+
+    box = atoms.box
+    if np.all(np.diag(np.diag(box)) == box):
+        atoms = atoms % np.diag(box)
+        tree = KDTree(atoms, boxsize=np.diag(box))
+    else:
+        atoms_pbc = pbc_points(query_atoms, box, thickness=r_max + 0.1, **kwargs)
+        tree = KDTree(atoms_pbc)
+
+    num_of_neighbors = tree.query_ball_point(query_atoms, r_max, return_length=True)
+    return num_of_neighbors - dnn