Moved functions around to circumvent a circular import

.gitignore

@@ -17,3 +17,4 @@ tmp/
 .traj.xtc_offsets.lock
 .traj.xtc_offsets.npz
 *.npy
+.venv/

src/mdevaluate/atoms.py

@@ -1,13 +1,8 @@
 import re
-from typing import Optional
 
 import numpy as np
-import numpy.typing as npt
-from scipy.spatial import KDTree
 
-from .pbc import pbc_diff, pbc_points
 from .checksum import checksum
-from .coordinates import CoordinateFrame
 
 
 def compare_regex(str_list: list[str], exp: str) -> np.ndarray:
@@ -176,95 +171,6 @@ class AtomSubset:
         return checksum(self.description)
 
 
-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.
-
-    Args:
-        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 }
-        \rigth)^{\frac{1}{2}}
-    """
-    gyration_radii = np.array([])
-
-    for resid in np.unique(position.residue_ids):
-        pos = position.whole[position.residue_ids == resid]
-        mass = position.masses[position.residue_ids == resid][:, np.newaxis]
-        COM = 1 / mass.sum() * (mass * position).sum(axis=0)
-        r_sq = ((pbc_diff(pos, COM, pos.box.diagonal())) ** 2).sum(1)[:, np.newaxis]
-        g_radius = ((r_sq * mass).sum() / mass.sum()) ** 0.5
-
-        gyration_radii = np.append(gyration_radii, g_radius)
-
-    return gyration_radii
 
 
-def layer_of_atoms(
-    atoms: CoordinateFrame,
-    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
-    distance = np.dot(atoms, plane_normal)
-    return np.abs(distance) <= thickness
 
-
-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):
-    """
-    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
-        number_of_neighbors (int, opt.): Number of neighboring atoms to find
-        distance_upper_bound (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))
-        distances, indices = tree.query(
-            query_atoms,
-            number_of_neighbors + dnn,
-            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 + dnn,
-            distance_upper_bound=distance_upper_bound,
-        )
-        indices = atoms_pbc_index[indices]
-
-    return distances[:, dnn:], indices[:, dnn:]

src/mdevaluate/coordinates.py

@@ -5,9 +5,10 @@ from typing import Optional, Callable
 
 import numpy as np
 import numpy.typing as npt
+from scipy.spatial import KDTree
 
 from .atoms import AtomSubset
-from .pbc import whole, nojump, pbc_diff
+from .pbc import whole, nojump, pbc_diff, pbc_points
 from .utils import singledispatchmethod
 from .checksum import checksum
 
@@ -458,8 +459,8 @@ def pore_coordinates(
 @map_coordinates
 def vectors(
     frame: CoordinateFrame,
-    atoms_indices_a: npt.ArrayLike[int],
-    atoms_indices_b: npt.ArrayLike[int],
+    atoms_indices_a: npt.ArrayLike,
+    atoms_indices_b: npt.ArrayLike,
     normed: bool = False,
 ) -> np.ndarray:
     """
@@ -506,7 +507,7 @@ def vectors(
 
 @map_coordinates
 def dipole_vector(
-    frame: CoordinateFrame, atom_indices: npt.ArrayLike[int], normed: bool = None
+    frame: CoordinateFrame, atom_indices: npt.ArrayLike, normed: bool = None
 ) -> np.ndarray:
     coords = frame.whole[atom_indices]
     res_ids = frame.residue_ids[atom_indices]
@@ -524,7 +525,7 @@ def dipole_vector(
 @map_coordinates
 def sum_dipole_vector(
     coordinates: CoordinateFrame,
-    atom_indices: npt.ArrayLike[int],
+    atom_indices: npt.ArrayLike,
     normed: bool = True,
 ) -> np.ndarray:
     coords = coordinates.whole[atom_indices]
@@ -538,9 +539,9 @@ def sum_dipole_vector(
 @map_coordinates
 def normal_vectors(
     frame: CoordinateFrame,
-    atom_indices_a: npt.ArrayLike[int],
-    atom_indices_b: npt.ArrayLike[int],
-    atom_indices_c: npt.ArrayLike[int],
+    atom_indices_a: npt.ArrayLike,
+    atom_indices_b: npt.ArrayLike,
+    atom_indices_c: npt.ArrayLike,
     normed: bool = True,
 ) -> np.ndarray:
     coords_a = frame[atom_indices_a]
@@ -580,3 +581,70 @@ def cylindrical_coordinates(
     radius = (x**2 + y**2) ** 0.5
     phi = np.arctan2(y, x)
     return np.array([radius, phi, z]).T
+
+
+def layer_of_atoms(
+        atoms: CoordinateFrame,
+        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
+    distance = np.dot(atoms, plane_normal)
+    return np.abs(distance) <= thickness
+
+
+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):
+    """
+    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
+        number_of_neighbors (int, opt.): Number of neighboring atoms to find
+        distance_upper_bound (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))
+        distances, indices = tree.query(
+            query_atoms,
+            number_of_neighbors + dnn,
+            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 + dnn,
+            distance_upper_bound=distance_upper_bound,
+            )
+        indices = atoms_pbc_index[indices]
+
+    return distances[:, dnn:], indices[:, dnn:]

src/mdevaluate/distribution.py

@@ -4,8 +4,13 @@ from scipy import spatial
 from scipy.spatial import KDTree
 from scipy.sparse.csgraph import connected_components
 
-from .coordinates import rotate_axis, polar_coordinates, Coordinates, CoordinateFrame
-from .atoms import next_neighbors
+from .coordinates import (
+    rotate_axis,
+    polar_coordinates,
+    Coordinates,
+    CoordinateFrame,
+    next_neighbors,
+)
 from .autosave import autosave_data
 from .utils import runningmean
 from .pbc import pbc_diff, pbc_points
@@ -455,3 +460,30 @@ def calc_cluster_sizes(frame, r_max=0.35):
     for i in range(0, np.max(labels) + 1):
         cluster_sizes.append(np.sum(labels == i))
     return np.array(cluster_sizes).flatten()
+
+
+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.
+
+    Args:
+        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 }
+        \rigth)^{\frac{1}{2}}
+    """
+    gyration_radii = np.array([])
+
+    for resid in np.unique(position.residue_ids):
+        pos = position.whole[position.residue_ids == resid]
+        mass = position.masses[position.residue_ids == resid][:, np.newaxis]
+        COM = 1 / mass.sum() * (mass * position).sum(axis=0)
+        r_sq = ((pbc_diff(pos, COM, pos.box.diagonal())) ** 2).sum(1)[:, np.newaxis]
+        g_radius = ((r_sq * mass).sum() / mass.sum()) ** 0.5
+
+        gyration_radii = np.append(gyration_radii, g_radius)
+
+    return gyration_radii