docstrings and typing

neuro_py/behavior/cheeseboard.py

@@ -3,6 +3,37 @@
 
 
 def plot_grid_with_circle_and_random_dots():
+    """
+    Plots a 15x15 grid of dots within a circle, highlights 3 randomly chosen dots 
+    within the circle, and draws a grey box at the bottom.
+
+    The function generates a grid of points within a circle of a specified radius 
+    and randomly selects three points from within the circle. These points are 
+    colored red and slightly enlarged. Additionally, a grey box is drawn at the 
+    bottom of the plot.
+
+    Parameters
+    ----------
+    None
+
+    Returns
+    -------
+    None
+
+    Notes
+    -----
+    - The grid is plotted on a 15x15 layout, with points that fall within the 
+      circle of radius 6.8 being displayed.
+    - The randomly selected points must be at least 4 grid units apart.
+    - A grey rectangular box is drawn near the bottom of the plot for aesthetic 
+      purposes.
+
+    Examples
+    --------
+    >>> plot_grid_with_circle_and_random_dots()
+    # This will display a plot of a circle containing a grid of dots with 
+    # 3 randomly chosen dots highlighted in red.
+    """
     # Create a 15x15 grid of dots within the circle
     x = np.linspace(-7, 7, 17)
     y = np.linspace(-7, 7, 17)

neuro_py/behavior/kinematics.py

@@ -1,16 +1,85 @@
-import numpy as np
-
 from typing import Union
 
+import numpy as np
+
 
 def get_velocity(
     position: np.ndarray, time: Union[np.ndarray, None] = None
 ) -> np.ndarray:
+    """
+    Computes the velocity from position data.
+
+    If time is not provided, it assumes a constant time step between position
+    samples. The velocity is calculated as the gradient of the position with
+    respect to time along the first axis.
+
+    Parameters
+    ----------
+    position : np.ndarray
+        An array of position data. This can be 1D (for single-dimensional positions)
+        or 2D (for multi-dimensional positions, e.g., x and y coordinates over time).
+    time : Union[np.ndarray, None], optional
+        An array of time values corresponding to the position data. If None,
+        the function assumes a constant time step. Default is None.
+
+    Returns
+    -------
+    np.ndarray
+        An array of velocity values, where each velocity is the rate of change of
+        position with respect to time.
+
+    Examples
+    --------
+    >>> position = np.array([0, 1, 4, 9, 16])
+    >>> get_velocity(position)
+    array([1., 2., 4., 6., 7.])
+
+    >>> position = np.array([[0, 0], [1, 1], [4, 4], [9, 9], [16, 16]])
+    >>> time = np.array([0, 1, 2, 3, 4])
+    >>> get_velocity(position, time)
+    array([[1., 1.],
+        [2., 2.],
+        [4., 4.],
+        [6., 6.],
+        [7., 7.]])
+    """
     if time is None:
         time = np.arange(position.shape[0])
     return np.gradient(position, time, axis=0)
 
 
 def get_speed(position: np.ndarray, time: Union[np.ndarray, None] = None) -> np.ndarray:
+    """
+    Computes the speed from position data.
+
+    Speed is the magnitude of the velocity vector at each time point. If time is
+    not provided, it assumes a constant time step between position samples.
+
+    Parameters
+    ----------
+    position : np.ndarray
+        An array of position data. This can be 1D (for single-dimensional positions)
+        or 2D (for multi-dimensional positions, e.g., x and y coordinates over time).
+    time : Union[np.ndarray, None], optional
+        An array of time values corresponding to the position data. If None,
+        the function assumes a constant time step. Default is None.
+
+    Returns
+    -------
+    np.ndarray
+        An array of speed values, where each speed is the magnitude of the velocity
+        at the corresponding time point.
+
+    Examples
+    --------
+    >>> position = np.array([0, 1, 4, 9, 16])
+    >>> get_speed(position)
+    array([1.41421356, 2.82842712, 5.65685425, 8.48528137, 9.89949494])
+
+    >>> position = np.array([[0, 0], [1, 1], [4, 4], [9, 9], [16, 16]])
+    >>> time = np.array([0, 1, 2, 3, 4])
+    >>> get_speed(position, time)
+    array([1.41421356, 2.82842712, 5.65685425, 8.48528137, 9.89949494])
+    """
     velocity = get_velocity(position, time=time)
     return np.sqrt(np.sum(velocity**2, axis=1))