#476 Add docstrings and type hints

lindemann/index/online_atoms.py

@@ -7,16 +7,27 @@
 from ovito.pipeline import Pipeline
 
 
-@nb.njit(fastmath=True, error_model="numpy")  # type: ignore # , cache=True) #(parallel=True)
-def calculate_frame(positions, array_mean, array_var, frame, natoms) -> npt.NDArray[np.float32]:
-    """Calculate the contribution of each atom to the lindemann index over the frames
+@nb.njit(fastmath=True, parallel=False)
+def calculate_frame(
+    positions: npt.NDArray[np.float32],
+    array_mean: npt.NDArray[np.float32],
+    array_var: npt.NDArray[np.float32],
+    frame: int,
+    natoms: int,
+) -> npt.NDArray[np.float32]:
+    """
+    Calculates the contribution of the individual atomic positions to the Lindemann Index for a specific frame.
 
     Args:
-        frames: numpy array of shape(frames,atoms)
+        positions (npt.NDArray[np.float32]): Array of atomic positions for the current frame.
+        array_mean (npt.NDArray[np.float32]): Array to store the mean distances.
+        array_var (npt.NDArray[np.float32]): Array to store the variance of distances.
+        frame (int): The current frame index.
+        natoms (int): The number of atoms.
+
     Returns:
-        npt.NDArray[np.float32]: Returns 1D array with the progression of the lindeman index per frame of shape(frames, atoms)
+        npt.NDArray[np.float32]: Array of the individual atomic contributions to the Lindemann indices for the current frame.
     """
-
     frame_count = frame + 1
     for i in range(natoms):
         for j in range(i + 1, natoms):
@@ -46,7 +57,23 @@ def calculate_frame(positions, array_mean, array_var, frame, natoms) -> npt.NDAr
     return lindemann_indices
 
 
-def calculate(pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None):
+def calculate(
+    pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None
+) -> npt.NDArray[np.float32]:
+    """
+    Calculates the contribution of the individual atomic positions to the Lindemann Index for a series of frames from an OVITO pipeline.
+
+    Args:
+        pipeline (Pipeline): The OVITO pipeline object.
+        data (DataCollection): The data collection object from OVITO.
+        nframes (Optional[int]): The number of frames to process. If None, all frames are processed.
+
+    Returns:
+        npt.NDArray[np.float32]: Array of the individual atomic contributions to the Lindemann indices for each frame.
+
+    Raises:
+        ValueError: If the requested number of frames exceeds the available frames in the pipeline.
+    """
     num_particle = data.particles.count
     num_frame = pipeline.source.num_frames
     if nframes is None:

lindemann/index/online_frames.py

@@ -1,17 +1,33 @@
-from typing import Optional
+from typing import Any, Optional
 
 import numba as nb
 import numpy as np
+import numpy.typing as npt
 from ovito.data import DataCollection
 from ovito.pipeline import Pipeline
 
-import lindemann
-from lindemann.trajectory import read
-
 
 @nb.njit(fastmath=True, parallel=False)
-def calculate_frame(positions, mean_distances, m2_distances, frame: int, num_atoms: int):
+def calculate_frame(
+    positions: npt.NDArray[np.float32],
+    mean_distances: npt.NDArray[np.float32],
+    m2_distances: npt.NDArray[np.float32],
+    frame: int,
+    num_atoms: int,
+) -> np.floating[Any]:
+    """
+    Calculates the Lindemann Index for a specific frame.
+
+    Args:
+        positions (npt.NDArray[np.float32]): Array of atomic positions for the current frame.
+        mean_distances (npt.NDArray[np.float32]): Array to store the mean distances between pairs of atoms.
+        m2_distances (npt.NDArray[np.float32]): Array to store the squared differences of distances between pairs of atoms.
+        frame (int): The current frame index.
+        num_atoms (int): The number of atoms.
 
+    Returns:
+        float: The Lindemann index for the current frame.
+    """
     index = 0
     frame_count = frame + 1
     for i in range(num_atoms):
@@ -30,7 +46,23 @@ def calculate_frame(positions, mean_distances, m2_distances, frame: int, num_ato
     return np.mean(np.sqrt(m2_distances / frame_count) / mean_distances)
 
 
-def calculate(pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None):
+def calculate(
+    pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None
+) -> npt.NDArray[np.float32]:
+    """
+    Calculates the Lindemann indices for a series of frames from an OVITO pipeline.
+
+    Args:
+        pipeline (Pipeline): The OVITO pipeline object.
+        data (DataCollection): The data collection object from OVITO.
+        nframes (Optional[int]): The number of frames to process. If None, all frames are processed.
+
+    Returns:
+        npt.NDArray[np.float32]: Array of Lindemann indices for each frame.
+
+    Raises:
+        ValueError: If the requested number of frames exceeds the available frames in the pipeline.
+    """
     num_particle = data.particles.count
     num_frame = pipeline.source.num_frames
     if nframes is None:

lindemann/index/online_trj.py

@@ -1,16 +1,35 @@
-from typing import Optional
+from typing import Any, Optional
 
 import numba as nb
 import numpy as np
+import numpy.typing as npt
 from ovito.data import DataCollection
 from ovito.pipeline import Pipeline
 
 from lindemann.trajectory import read
 
 
 @nb.njit(fastmath=True, parallel=False)
-def calculate_frame(positions, mean_distances, m2_distances, frame: int, num_atoms: int):
+def calculate_frame(
+    positions: npt.NDArray[np.float32],
+    mean_distances: npt.NDArray[np.float32],
+    m2_distances: npt.NDArray[np.float32],
+    frame: int,
+    num_atoms: int,
+) -> None:
+    """
+    Updates the mean and variance of distances between pairs of atoms for a specific frame.
 
+    Args:
+        positions (npt.NDArray[np.float32]): Array of atomic positions for the current frame.
+        mean_distances (npt.NDArray[np.float32]): Array to store the mean distances between pairs of atoms.
+        m2_distances (npt.NDArray[np.float32]): Array to store the squared differences of distances between pairs of atoms.
+        frame (int): The current frame index.
+        num_atoms (int): The number of atoms.
+
+    Returns:
+        None
+    """
     index = 0
     frame_count = frame + 1
     for i in range(num_atoms):
@@ -28,7 +47,23 @@ def calculate_frame(positions, mean_distances, m2_distances, frame: int, num_ato
             index += 1
 
 
-def calculate(pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None):
+def calculate(
+    pipeline: Pipeline, data: DataCollection, nframes: Optional[int] = None
+) -> np.floating[Any]:
+    """
+    Calculates the overall Lindemann index for a series of frames from an OVITO pipeline.
+
+    Args:
+        pipeline (Pipeline): The OVITO pipeline object.
+        data (DataCollection): The data collection object from OVITO.
+        nframes (Optional[int]): The number of frames to process. If None, all frames are processed.
+
+    Returns:
+        float: The overall Lindemann index.
+
+    Raises:
+        ValueError: If the requested number of frames exceeds the available frames in the pipeline.
+    """
     num_particle = data.particles.count
     num_frame = pipeline.source.num_frames
     if nframes is None:
@@ -46,5 +81,4 @@ def calculate(pipeline: Pipeline, data: DataCollection, nframes: Optional[int] =
             data.particles["Position"].array, mean_distances, m2_distances, frame, num_particle
         )
 
-    linde = np.mean(np.sqrt(m2_distances / nframes) / mean_distances)
-    return linde
+    return np.mean(np.sqrt(m2_distances / nframes) / mean_distances)

lindemann/index/per_atoms.py

@@ -3,14 +3,18 @@
 import numpy.typing as npt
 
 
-@nb.njit(fastmath=True, error_model="numpy")  # type: ignore # , cache=True) #(parallel=True)
+@nb.njit(fastmath=True, parallel=False)
 def calculate(frames: npt.NDArray[np.float32]) -> npt.NDArray[np.float32]:
-    """Calculate the contribution of each atom to the lindemann index over the frames
+    """
+    Calculate the contribution of each atom to the Lindemann index over the frames.
 
     Args:
-        frames: numpy array of shape(frames,atoms)
+        frames (npt.NDArray[np.float32]): A numpy array of shape (frames, atoms, 3) containing the atomic positions
+                                          over multiple frames.
+
     Returns:
-        npt.NDArray[np.float32]: Returns 1D array with the progression of the lindeman index per frame of shape(frames, atoms)
+        npt.NDArray[np.float32]: A 2D array of shape (frames, atoms) containing the progression of the Lindemann index
+                                 per frame.
     """
     len_frames, natoms, _ = frames.shape
 

lindemann/index/per_frames.py

@@ -1,9 +1,19 @@
 import numba as nb
 import numpy as np
+import numpy.typing as npt
 
 
 @nb.njit(fastmath=True, parallel=False)
-def calculate(positions):
+def calculate(positions: npt.NDArray[np.float32]):
+    """
+    Calculates the Lindemann index for each frame of atomic positions.
+
+    Args:
+        positions (npt.NDArray[np.float32]): Array of atomic positions with shape (num_frames, num_atoms, 3).
+
+    Returns:
+        npt.NDArray[np.float32]: Array of Lindemann indices for each frame.
+    """
     num_frames, num_atoms, _ = positions.shape
     num_distances = num_atoms * (num_atoms - 1) // 2
 

lindemann/index/per_trj.py

@@ -1,9 +1,21 @@
+from typing import Any
+
 import numba as nb
 import numpy as np
+import numpy.typing as npt
 
 
 @nb.njit(fastmath=True)
-def calculate(positions):
+def calculate(positions: npt.NDArray[np.float32]) -> np.floating[Any]:
+    """
+    Calculates the overall Lindemann index for a series of atomic positions over multiple frames.
+
+    Args:
+        positions (npt.NDArray[np.float32]): Array of atomic positions with shape (num_frames, num_atoms, 3).
+
+    Returns:
+        np.floating[np.Any]: The overall Lindemann index.
+    """
     num_frames, num_atoms, _ = positions.shape
     num_distances = num_atoms * (num_atoms - 1) // 2
 

lindemann/trajectory/plt_plot.py

@@ -6,7 +6,7 @@
 mpl.use("Agg")
 
 
-def lindemann_vs_frames(indices: npt.NDArray[np.float64]) -> str:
+def lindemann_vs_frames(indices: npt.NDArray[np.float32]) -> str:
     plt.figure(1)
     plt.title("Lindemann index per frame")
     plt.xlabel("Frames")