Started rewriting, adding type hints

src/mdevaluate/__init__.py

@@ -1,5 +1,6 @@
 import os
 from glob import glob
+from typing import Optional
 
 import pandas as pd
 
@@ -11,20 +12,19 @@ from . import functions
 from . import pbc
 from . import autosave
 from . import reader
-from . import chill
 from . import system
-from . import free_energy_landscape
+from .extra import free_energy_landscape, chill
 from .logging import logger
 
 
 def open(
-    directory="",
-    topology="*.tpr",
-    trajectory="*.xtc",
-    nojump=False,
-    index_file=None,
-    charges=None,
-    masses=None,
+    directory: str = "",
+    topology: str = "*.tpr",
+    trajectory: str = "*.xtc",
+    nojump: bool = False,
+    index_file: Optional[str] = None,
+    charges: Optional[list[float]] = None,
+    masses: Optional[list[float]] = None,
 ):
     """
     Open a simulation from a directory.
@@ -34,8 +34,17 @@ def open(
         topology (opt.):
             Descriptor of the topology file (tpr or gro). By default, a tpr file is
             used, if there is exactly one in the directoy.
-        trajectory (opt.): Descriptor of the trajectory (xtc file).
-        nojump (opt.): If nojump matrixes should be generated. They will alwyas be loaded if present
+        trajectory (opt.): Descriptor of the trajectory (xtc or trr file).
+        nojump (opt.):
+            If nojump matrixes should be generated. They will alwyas be loaded
+            if present
+        index_file (opt.): Descriptor of the index file (ndx file).
+        charges (opt.):
+            List with charges for each atom. It Has to be the same length as the number
+            of atoms in the system. Only used if topology file is a gro file.
+        masses (opt.):
+            List with masses for each atom. It Has to be the same length as the number
+            of atoms in the system. Only used if topology file is a gro file.
 
     Returns:
         A Coordinate object of the simulation.
@@ -43,13 +52,14 @@ def open(
     Example:
         Open a simulation located in '/path/to/sim', where the trajectory is
         located in a sub-directory '/path/to/sim/out' and named for Example
-        'nojump_traj.xtc'. All read frames will be cached in memory.
+        'nojump_traj.xtc'.
 
-        >>> open('/path/to/sim', trajectory='out/nojump*.xtc', cached=None)
+        >>> open('/path/to/sim', trajectory='out/nojump*.xtc')
 
     The file descriptors can use unix style pathname expansion to define the filenames.
     The default patterns use the recursive placeholder `**` which matches the base or
-    any subdirctory, thus files in subdirectories with matching file type will be found too.
+    any subdirctory, thus files in subdirectories with matching file type will be found
+    too.
     For example: 'out/nojump*.xtc' would match xtc files in a subdirectory `out` that
     start with `nojump` and end with `.xtc`.
 
@@ -91,9 +101,11 @@ def open(
     return coords
 
 
-def open_energy(file):
+def open_energy(file: str):
     """Reads a gromacs energy file and output the data in a pandas DataFrame.
     Args:
         file: Filename of the energy file
+    Returns:
+        A DataFrame with the different energies doe each time.
     """
     return pd.DataFrame(reader.energy_reader(file).data_dict)

src/mdevaluate/atoms.py

@@ -1,25 +1,21 @@
 import re
+from typing import Optional
 
-from .pbc import pbc_diff
-from .checksum import checksum
 import numpy as np
+import numpy.typing as npt
+from scipy.spatial import KDTree
 
-import scipy
-
-if scipy.version.version >= "0.17.0":
-    from scipy.spatial import cKDTree as KDTree
-else:
-    from scipy.spatial import KDTree
+from .pbc import pbc_diff, pbc_points
+from .checksum import checksum
+from .coordinates import CoordinateFrame
 
 
-def compare_regex(list, exp):
+def compare_regex(str_list: list[str], exp: str):
     """
     Compare a list of strings with a regular expression.
     """
-    if not exp.endswith("$"):
-        exp += "$"
     regex = re.compile(exp)
-    return np.array([regex.match(s) is not None for s in list])
+    return np.array([regex.match(s) is not None for s in str_list])
 
 
 class Atoms:
@@ -180,34 +176,26 @@ class AtomSubset:
         return checksum(self.description)
 
 
-def center_of_mass(position, mass=None):
-    if mass is not None:
-        return 1 / mass.sum() * (mass * position).sum(axis=0)
-    else:
-        return 1 / len(position) * position.sum(axis=0)
-
-
-def gyration_radius(position):
+def gyration_radius(position: CoordinateFrame):
     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.
+    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
+        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 = center_of_mass(pos, mass)
+        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)
+        g_radius = ((r_sq * mass).sum() / mass.sum()) ** 0.5
 
         gyration_radii = np.append(gyration_radii, g_radius)
 
@@ -215,12 +203,16 @@ def gyration_radius(position):
 
 
 def layer_of_atoms(
-    atoms, thickness, plane_offset=np.array([0, 0, 0]), plane_normal=np.array([1, 0, 0])
+    atoms: CoordinateFrame,
+    thickness: float,
+    plane_normal: npt.ArrayLike,
+    plane_offset: Optional[npt.ArrayLike] = np.array([0, 0, 0]),
 ):
-    p_ = atoms - plane_offset
-    distance = np.dot(p_, plane_normal)
-
-    return abs(distance) <= thickness
+    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 distance_to_atoms(ref, atoms, box=None):
@@ -234,7 +226,7 @@ def distance_to_atoms(ref, atoms, box=None):
     return out
 
 
-def distance_to_atoms_cKDtree(ref, atoms, box=None, thickness=None):
+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)
@@ -251,7 +243,7 @@ def distance_to_atoms_cKDtree(ref, atoms, box=None, thickness=None):
         start_coords = atoms
         all_frame_coords = ref
 
-    tree = spatial.cKDTree(all_frame_coords)
+    tree = KDTree(all_frame_coords)
     return tree.query(start_coords)[0]
 
 
@@ -266,11 +258,16 @@ def next_neighbors(
     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
+        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
+        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
     """
     tree = KDTree(atoms)
     dnn = 0

src/mdevaluate/coordinates.py

@@ -217,7 +217,7 @@ class Coordinates:
     """
     Coordinates represent trajectory data, which is used for evaluation functions.
 
-    Atoms may be selected by specifing a atom_subset or a atom_filter.
+    Atoms may be selected by specifying an atom_subset or an atom_filter.
     """
 
     def get_mode(self, mode):
@@ -424,13 +424,6 @@ def map_coordinates(func):
 
 @map_coordinates
 def center_of_masses(coordinates, atoms=None, shear: bool = False):
-    """
-    Example:
-    rd = XTCReader('t.xtc')
-    coordinates = Coordinates(rd)
-    com = centers_of_mass(coordinates, (1.0, 2.0, 1.0, 3.0))
-
-    """
     if atoms is None:
         atoms = list(range(len(coordinates)))
     res_ids = coordinates.residue_ids[atoms]