Added type hints, adjusted imports, and doc strings

src/mdevaluate/coordinates.py

@@ -1,13 +1,14 @@
-from functools import partial, lru_cache, wraps
+from functools import partial, wraps
 from copy import copy
 from .logging import logger
+from typing import Optional, Callable
 
 import numpy as np
-from scipy.spatial import cKDTree, KDTree
+import numpy.typing as npt
 
 from .atoms import AtomSubset
 from .pbc import whole, nojump, pbc_diff
-from .utils import mask2indices, singledispatchmethod
+from .utils import singledispatchmethod
 from .checksum import checksum
 
 
@@ -15,104 +16,6 @@ class UnknownCoordinatesMode(Exception):
     pass
 
 
-def rotate_axis(coords, axis):
-    """
-    Rotate a set of coordinates to a given axis.
-    """
-    axis = np.array(axis) / np.linalg.norm(axis)
-    zaxis = np.array([0, 0, 1])
-    if (axis == zaxis).sum() == 3:
-        return coords
-    rotation_axis = np.cross(axis, zaxis)
-    rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
-
-    theta = np.arccos(axis @ zaxis / np.linalg.norm(axis))
-
-    # return theta/pi, rotation_axis
-
-    ux, uy, uz = rotation_axis
-    cross_matrix = np.array([[0, -uz, uy], [uz, 0, -ux], [-uy, ux, 0]])
-    rotation_matrix = (
-        np.cos(theta) * np.identity(len(axis))
-        + (1 - np.cos(theta))
-        * rotation_axis.reshape(-1, 1)
-        @ rotation_axis.reshape(1, -1)
-        + np.sin(theta) * cross_matrix
-    )
-
-    if len(coords.shape) == 2:
-        rotated = np.array([rotation_matrix @ xyz for xyz in coords])
-    else:
-        rotated = rotation_matrix @ coords
-    return rotated
-
-
-def spherical_radius(frame, origin=None):
-    """
-    Transform a frame of cartesian coordinates into the sperical radius.
-    If origin=None the center of the box is taken as the coordinates origin.
-    """
-    if origin is None:
-        origin = frame.box.diagonal() / 2
-    return ((frame - origin) ** 2).sum(axis=-1) ** 0.5
-
-
-def polar_coordinates(x, y):
-    """Convert cartesian to polar coordinates."""
-    radius = (x**2 + y**2) ** 0.5
-    phi = np.arctan2(y, x)
-    return radius, phi
-
-
-def spherical_coordinates(x, y, z):
-    """Convert cartesian to spherical coordinates."""
-    xy, phi = polar_coordinates(x, y)
-    radius = (x**2 + y**2 + z**2) ** 0.5
-    theta = np.arccos(z / radius)
-    return radius, phi, theta
-
-
-def radial_selector(frame, coordinates, rmin, rmax):
-    """
-    Return a selection of all atoms with radius in the interval [rmin, rmax].
-    """
-    crd = coordinates[frame.step]
-    rad, _ = polar_coordinates(crd[:, 0], crd[:, 1])
-    selector = (rad >= rmin) & (rad <= rmax)
-    return mask2indices(selector)
-
-
-def selector_radial_cylindrical(atoms, r_min, r_max, origin=None):
-    box = atoms.box
-    atoms = atoms % np.diag(box)
-    if origin is None:
-        origin = [box[0, 0] / 2, box[1, 1] / 2, box[2, 2] / 2]
-    r_vec = (atoms - origin)[:, :2]
-    r = np.linalg.norm(r_vec, axis=1)
-    index = np.argwhere((r >= r_min) * (r < r_max))
-    return index.flatten()
-
-
-def spatial_selector(frame, transform, rmin, rmax):
-    """
-    Select a subset of atoms which have a radius between rmin and rmax.
-    Coordinates are filtered by the condition::
-
-      rmin <= transform(frame) <= rmax
-
-    Args:
-        frame: The coordinates of the actual trajectory
-        transform:
-            A function that transforms the coordinates of the frames into
-            the one-dimensional spatial coordinate (e.g. radius).
-        rmin: Minimum value of the radius
-        rmax: Maximum value of the radius
-    """
-    r = transform(frame)
-    selector = (rmin <= r) & (rmax >= r)
-    return mask2indices(selector)
-
-
 class CoordinateFrame(np.ndarray):
     _known_modes = ("pbc", "whole", "nojump")
 
@@ -250,7 +153,8 @@ class Coordinates:
     def mode(self, val):
         if val in CoordinateFrame._known_modes:
             logger.warn(
-                "Changing the Coordinates mode directly is deprecated. Use Coordinates.%s instead, which returns a copy.",
+                "Changing the Coordinates mode directly is deprecated. "
+                "Use Coordinates.%s instead, which returns a copy.",
                 val,
             )
             self._mode = val
@@ -414,30 +318,111 @@ class CoordinatesMap:
         return CoordinatesMap(self.coordinates.pbc, self.function)
 
 
-def map_coordinates(func):
+def rotate_axis(coords: npt.ArrayLike, axis: npt.ArrayLike) -> np.ndarray:
+    """
+    Rotate a set of coordinates to a given axis.
+    """
+    axis = np.array(axis) / np.linalg.norm(axis)
+    zaxis = np.array([0, 0, 1])
+    if (axis == zaxis).sum() == 3:
+        return coords
+    rotation_axis = np.cross(axis, zaxis)
+    rotation_axis = rotation_axis / np.linalg.norm(rotation_axis)
+
+    theta = np.arccos(axis @ zaxis / np.linalg.norm(axis))
+
+    # return theta/pi, rotation_axis
+
+    ux, uy, uz = rotation_axis
+    cross_matrix = np.array([[0, -uz, uy], [uz, 0, -ux], [-uy, ux, 0]])
+    rotation_matrix = (
+        np.cos(theta) * np.identity(len(axis))
+        + (1 - np.cos(theta))
+        * rotation_axis.reshape(-1, 1)
+        @ rotation_axis.reshape(1, -1)
+        + np.sin(theta) * cross_matrix
+    )
+
+    if len(coords.shape) == 2:
+        rotated = np.array([rotation_matrix @ xyz for xyz in coords])
+    else:
+        rotated = rotation_matrix @ coords
+    return rotated
+
+
+def spherical_radius(
+    frame: CoordinateFrame, origin: Optional[npt.ArrayLike] = None
+) -> np.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.
+    """
+    if origin is None:
+        origin = frame.box.diagonal() / 2
+    return ((frame - origin) ** 2).sum(axis=-1) ** 0.5
+
+
+def polar_coordinates(x: npt.ArrayLike, y: npt.ArrayLike) -> (np.ndarray, np.ndarray):
+    """Convert cartesian to polar coordinates."""
+    radius = (x**2 + y**2) ** 0.5
+    phi = np.arctan2(y, x)
+    return radius, phi
+
+
+def spherical_coordinates(
+    x: npt.ArrayLike, y: npt.ArrayLike, z: npt.ArrayLike
+) -> (np.ndarray, np.ndarray, np.ndarray):
+    """Convert cartesian to spherical coordinates."""
+    xy, phi = polar_coordinates(x, y)
+    radius = (x**2 + y**2 + z**2) ** 0.5
+    theta = np.arccos(z / radius)
+    return radius, phi, theta
+
+
+def selector_radial_cylindrical(
+    atoms: CoordinateFrame,
+    r_min: float,
+    r_max: float,
+    origin: Optional[npt.ArrayLike] = None,
+) -> np.ndarray:
+    box = atoms.box
+    atoms = atoms % np.diag(box)
+    if origin is None:
+        origin = [box[0, 0] / 2, box[1, 1] / 2, box[2, 2] / 2]
+    r_vec = (atoms - origin)[:, :2]
+    r = np.linalg.norm(r_vec, axis=1)
+    index = np.argwhere((r >= r_min) * (r < r_max))
+    return index.flatten()
+
+
+def map_coordinates(
+    func: Callable[[CoordinateFrame, ...], np.ndarray]
+) -> Callable[..., CoordinatesMap]:
     @wraps(func)
-    def wrapped(coordinates, **kwargs):
+    def wrapped(coordinates: Coordinates, **kwargs) -> CoordinatesMap:
         return CoordinatesMap(coordinates, partial(func, **kwargs))
 
     return wrapped
 
 
 @map_coordinates
-def center_of_masses(coordinates, atoms=None, shear: bool = False):
+def center_of_masses(
+    frame: CoordinateFrame, atoms=None, shear: bool = False
+) -> np.ndarray:
     if atoms is None:
-        atoms = list(range(len(coordinates)))
+        atoms = list(range(len(frame)))
-    res_ids = coordinates.residue_ids[atoms]
+    res_ids = frame.residue_ids[atoms]
-    masses = coordinates.masses[atoms]
+    masses = frame.masses[atoms]
     if shear:
-        coords = coordinates[atoms]
+        coords = frame[atoms]
-        box = coords.box
+        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])
         coms = coords[sort_ind[i]][res_ids - min(res_ids)]
         cor = pbc_diff(coords, coms, box)
         coords = coms + cor
     else:
-        coords = coordinates.whole[atoms]
+        coords = frame.whole[atoms]
     mask = np.bincount(res_ids)[1:] != 0
     positions = np.array(
         [
@@ -450,12 +435,14 @@ def center_of_masses(coordinates, atoms=None, shear: bool = False):
 
 
 @map_coordinates
-def pore_coordinates(coordinates, origin, sym_axis="z"):
+def pore_coordinates(
+    frame: CoordinateFrame, origin: npt.ArrayLike, sym_axis: str = "z"
+) -> np.ndarray:
     """
     Map coordinates of a pore simulation so the pore has cylindrical symmetry.
 
     Args:
-        coordinates: Coordinates of the simulation
+        frame: Coordinates of the simulation
         origin: Origin of the pore which will be the coordinates origin after mapping
         sym_axis (opt.): Symmtery axis of the pore, may be a literal direction
             'x', 'y' or 'z' or an array of shape (3,)
@@ -465,30 +452,34 @@ def pore_coordinates(coordinates, origin, sym_axis="z"):
         rot_axis[["x", "y", "z"].index(sym_axis)] = 1
     else:
         rot_axis = sym_axis
-
+    return rotate_axis(frame - origin, rot_axis)
-    return rotate_axis(coordinates - origin, rot_axis)
 
 
 @map_coordinates
-def vectors(coordinates, atoms_a, atoms_b, normed=False):
+def vectors(
+    frame: CoordinateFrame,
+    atoms_indices_a: npt.ArrayLike[int],
+    atoms_indices_b: npt.ArrayLike[int],
+    normed: bool = False,
+) -> np.ndarray:
     """
     Compute the vectors between the atoms of two subsets.
 
     Args:
-        coordinates: The Coordinates object the atoms will be taken from
+        frame: The Coordinates object the atoms will be taken from
-        atoms_a: Mask or indices of the first atom subset
+        atoms_indices_a: Mask or indices of the first atom subset
-        atoms_b: Mask or indices of the second atom subset
+        atoms_indices_b: Mask or indices of the second atom subset
         normed (opt.): If the vectors should be normed
-        box (opt.): If not None, the vectors are calcualte with PBC
 
-    The defintion of atoms_a/b can be any possible subript of a numpy array.
+    The definition of atoms_a/b can be any possible subript of a numpy array.
     They can, for example, be given as a masking array of bool values with the
-    same length as the frames of the coordinates. Or they can be a list of
+    same length as the frames of the coordinates.
-    indices selecting the atoms of these indices from each frame.
+    Or there can be a list of indices selecting the atoms of these indices from each
+    frame.
 
-    It is possible to compute the mean of several atoms before calculating the vectors,
+    It is possible to compute the means of several atoms before calculating the vectors,
     by using a two-dimensional list of indices. The following code computes the vectors
-    between atoms 0, 3, 6 and the mean coordinate of atoms 1, 4, 7 and 2, 5, 8::
+    between atoms 0, 3, 6 and the mean coordinate of atoms 1, 4, 7 and 2, 5, 8:
 
         >>> inds_a = [0, 3, 6]
         >>> inds_b = [[1, 4, 7], [2, 5, 8]]
@@ -499,11 +490,11 @@ def vectors(coordinates, atoms_a, atoms_b, normed=False):
             coords[6] - (coords[7] + coords[8])/2,
         ])
     """
-    box = coordinates.box
+    box = frame.box
-    coords_a = coordinates[atoms_a]
+    coords_a = frame[atoms_indices_a]
     if len(coords_a.shape) > 2:
         coords_a = coords_a.mean(axis=0)
-    coords_b = coordinates[atoms_b]
+    coords_b = frame[atoms_indices_b]
     if len(coords_b.shape) > 2:
         coords_b = coords_b.mean(axis=0)
     vec = pbc_diff(coords_a, coords_b, box=box)
@@ -514,10 +505,12 @@ def vectors(coordinates, atoms_a, atoms_b, normed=False):
 
 
 @map_coordinates
-def dipole_vector(coordinates, atoms, normed=None):
+def dipole_vector(
-    coords = coordinates.whole[atoms]
+    frame: CoordinateFrame, atom_indices: npt.ArrayLike[int], normed: bool = None
-    res_ids = coordinates.residue_ids[atoms]
+) -> np.ndarray:
-    charges = coordinates.charges[atoms]
+    coords = frame.whole[atom_indices]
+    res_ids = frame.residue_ids[atom_indices]
+    charges = frame.charges[atom_indices]
     mask = np.bincount(res_ids)[1:] != 0
     dipoles = np.array(
         [np.bincount(res_ids, weights=c * charges)[1:] for c in coords.T]
@@ -531,11 +524,11 @@ def dipole_vector(coordinates, atoms, normed=None):
 @map_coordinates
 def sum_dipole_vector(
     coordinates: CoordinateFrame,
-    atoms: list[int],
+    atom_indices: npt.ArrayLike[int],
     normed: bool = True,
-):
+) -> np.ndarray:
-    coords = coordinates.whole[atoms]
+    coords = coordinates.whole[atom_indices]
-    charges = coordinates.charges[atoms]
+    charges = coordinates.charges[atom_indices]
     dipole = np.array([c * charges for c in coords.T]).T
     if normed:
         dipole /= np.linalg.norm(dipole)
@@ -544,16 +537,16 @@ def sum_dipole_vector(
 
 @map_coordinates
 def normal_vectors(
-    coordinates: CoordinateFrame,
+    frame: CoordinateFrame,
-    atoms_a: list[int],
+    atom_indices_a: npt.ArrayLike[int],
-    atoms_b: list[int],
+    atom_indices_b: npt.ArrayLike[int],
-    atoms_c: list[int],
+    atom_indices_c: npt.ArrayLike[int],
     normed: bool = True,
-):
+) -> np.ndarray:
-    coords_a = coordinates[atoms_a]
+    coords_a = frame[atom_indices_a]
-    coords_b = coordinates[atoms_b]
+    coords_b = frame[atom_indices_b]
-    coords_c = coordinates[atoms_c]
+    coords_c = frame[atom_indices_c]
-    box = coordinates.box
+    box = frame.box
     vectors_a = pbc_diff(coords_a, coords_b, box=box)
     vectors_b = pbc_diff(coords_a, coords_c, box=box)
     vec = np.cross(vectors_a, vectors_b)
@@ -576,7 +569,9 @@ def displacements_without_drift(
 
 
 @map_coordinates
-def cylindrical_coordinates(frame, origin=None):
+def cylindrical_coordinates(
+    frame: CoordinateFrame, origin: npt.ArrayLike = None
+) -> np.ndarray:
     if origin is None:
         origin = np.diag(frame.box) / 2
     x = frame[:, 0] - origin[0]