Add and improve docstrings across core modules

Expanded and standardized docstrings for classes and functions in classifier, core, decoder, environments, discrete/continuous state transitions, and likelihood modules. Docstrings now include parameter and return value descriptions, improving code readability and API documentation for trajectory decoding and classification.

Author: edeno

replay_trajectory_classification/classifier.py

@@ -1,5 +1,7 @@
-"""State space models that classify trajectories as well as decode the
-trajectory from population spiking
+"""State space models that classify trajectories and decode from population spiking.
+
+This module contains classes for trajectory classification using neural population
+activity from both sorted spikes and clusterless data.
 """
 
 from __future__ import annotations
@@ -77,7 +79,24 @@
 
 
 class _ClassifierBase(BaseEstimator, ABC):
-    """Base class for classifier objects."""
+    """Base class for trajectory classifier objects.
+
+    Parameters
+    ----------
+    environments : list[Environment], optional
+        List of spatial environments to classify trajectories within.
+    observation_models : ObservationModel, optional
+        Observation models for neural data, by default None.
+    continuous_transition_types : list[list[...]], optional
+        Continuous state transition models for each state.
+    discrete_transition_type : DiagonalDiscrete | RandomDiscrete | UniformDiscrete | UserDefinedDiscrete, optional
+        Discrete state transition model, by default DiagonalDiscrete(0.968).
+    initial_conditions_type : UniformInitialConditions | UniformOneEnvironmentInitialConditions, optional
+        Initial conditions model, by default UniformInitialConditions().
+    infer_track_interior : bool, optional
+        Whether to infer track interior from position data, by default True.
+
+    """
 
     def __init__(
         self,
@@ -121,13 +140,15 @@ def __init__(
     def fit_environments(
         self, position: NDArray[np.float64], environment_labels: Optional[NDArray[np.int64]] = None
     ) -> None:
-        """Fits the Environment class on the position data to get information about the spatial environment.
+        """Fit the Environment class on position data to extract spatial information.
 
         Parameters
         ----------
         position : NDArray[np.float64], shape (n_time, n_position_dims)
+            Position coordinates over time.
         environment_labels : NDArray[np.int64], optional, shape (n_time,)
-            Labels for each time points about which environment it corresponds to, by default None
+            Labels indicating which environment each time point corresponds to,
+            by default None.
 
         """
         for environment in self.environments:
@@ -144,7 +165,7 @@ def fit_environments(
         )
 
     def fit_initial_conditions(self):
-        """Constructs the initial probability for the state and each spatial bin."""
+        """Construct the initial probability distribution for states and spatial bins."""
         logger.info("Fitting initial conditions...")
         environment_names_to_state = [
             obs.environment_name for obs in self.observation_models
@@ -178,7 +199,7 @@ def fit_continuous_state_transition(
         encoding_group_labels: Optional[NDArray[np.int64]] = None,
         environment_labels: Optional[NDArray[np.int64]] = None,
     ) -> None:
-        """Constructs the transition matrices for the continuous states.
+        """Construct the transition matrices for the continuous states.
 
         Parameters
         ----------
@@ -240,7 +261,7 @@ def fit_continuous_state_transition(
                 ] = st
 
     def fit_discrete_state_transition(self):
-        """Constructs the transition matrix for the discrete states."""
+        """Construct the transition matrix for the discrete states."""
         logger.info("Fitting discrete state transition")
         n_states = len(self.continuous_transition_types)
         self.discrete_state_transition_ = (

replay_trajectory_classification/continuous_state_transitions.py

@@ -18,10 +18,12 @@ def _normalize_row_probability(x: NDArray[np.float64]) -> NDArray[np.float64]:
     Parameters
     ----------
     x : NDArray[np.float64], shape (n_rows, n_cols)
+        Input matrix to normalize.
 
     Returns
     -------
     normalized_x : NDArray[np.float64], shape (n_rows, n_cols)
+        Row-normalized matrix where each row sums to 1.
 
     """
     x /= x.sum(axis=1, keepdims=True)
@@ -32,19 +34,19 @@ def _normalize_row_probability(x: NDArray[np.float64]) -> NDArray[np.float64]:
 def estimate_movement_var(
     position: NDArray[np.float64], sampling_frequency: int = 1
 ) -> NDArray[np.float64]:
-    """Estimates the movement variance based on position.
+    """Estimate the movement variance based on position data.
 
     Parameters
     ----------
-    position : NDArray[np.float64], shape (n_time, n_position_dim)
-        Position of the animal
+    position : NDArray[np.float64], shape (n_time, n_position_dims)
+        Position coordinates of the animal over time.
     sampling_frequency : int, optional
-        Number of samples per second.
+        Number of samples per second, by default 1.
 
     Returns
     -------
-    movement_var : NDArray[np.float64], shape (n_position_dim,)
-        Variance of the movement.
+    movement_var : NDArray[np.float64], shape (n_position_dims,)
+        Covariance matrix of movement scaled by sampling frequency.
 
     """
     position = atleast_2d(position)

replay_trajectory_classification/core.py

@@ -1,4 +1,8 @@
-"""Core algorithms for decoding."""
+"""Core algorithms for decoding.
+
+This module contains the fundamental algorithms for Bayesian decoding including
+causal and acausal state estimation and classification.
+"""
 
 from __future__ import annotations
 
@@ -46,18 +50,17 @@ def get_centers(bin_edges: NDArray[np.float64]) -> NDArray[np.float64]:
 
 @njit(parallel=True, error_model="numpy")
 def normalize_to_probability(distribution: NDArray[np.float64]) -> NDArray[np.float64]:
-    """Ensure the distribution integrates to 1 so that it is a probability
-    distribution.
+    """Ensure the distribution integrates to 1 so that it is a probability distribution.
 
     Parameters
     ----------
     distribution : NDArray[np.float64]
-        Probability distribution to normalize
+        Probability distribution values to normalize.
 
     Returns
     -------
-    NDArray[np.float64]
-        Normalized probability distribution that sums to 1
+    normalized_distribution : NDArray[np.float64]
+        Normalized probability distribution that sums to 1.
 
     """
     return distribution / np.nansum(distribution)
@@ -73,13 +76,18 @@ def _causal_decode(
     Parameters
     ----------
     initial_conditions : NDArray[np.float64], shape (n_bins,)
+        Initial probability distribution over state bins.
     state_transition : NDArray[np.float64], shape (n_bins, n_bins)
+        Transition probability matrix between state bins.
     likelihood : NDArray[np.float64], shape (n_time, n_bins)
+        Likelihood values for each time point and state bin.
 
     Returns
     -------
     posterior : NDArray[np.float64], shape (n_time, n_bins)
+        Posterior probability distribution over time and state bins.
     log_data_likelihood : float
+        Log-likelihood of the observed data.
 
     """
 
@@ -109,11 +117,14 @@ def _acausal_decode(
     Parameters
     ----------
     causal_posterior : NDArray[np.float64], shape (n_time, n_bins, 1)
+        Causal (forward-pass) posterior probabilities.
     state_transition : NDArray[np.float64], shape (n_bins, n_bins)
+        Transition probability matrix between state bins.
 
-    Return
-    ------
+    Returns
+    -------
     acausal_posterior : NDArray[np.float64], shape (n_time, n_bins, 1)
+        Acausal (forward-backward) posterior probabilities.
 
     """
     acausal_posterior = np.zeros_like(causal_posterior)
@@ -150,14 +161,20 @@ def _causal_classify(
     Parameters
     ----------
     initial_conditions : NDArray[np.float64], shape (n_states, n_bins, 1)
+        Initial probability distribution for each state and spatial bin.
     continuous_state_transition : NDArray[np.float64], shape (n_states, n_states, n_bins, n_bins)
+        Continuous-space transition probabilities between states and bins.
     discrete_state_transition : NDArray[np.float64], shape (n_states, n_states)
+        Discrete state transition probabilities.
     likelihood : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+        Likelihood values for each time point, state, and spatial bin.
 
     Returns
     -------
     causal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+        Causal posterior probabilities over time, states, and spatial bins.
     log_data_likelihood : float
+        Log-likelihood of the observed data.
 
     """
     n_time, n_states, n_bins, _ = likelihood.shape
@@ -196,12 +213,16 @@ def _acausal_classify(
     Parameters
     ----------
     causal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+        Causal (forward-pass) posterior probabilities.
     continuous_state_transition : NDArray[np.float64], shape (n_states, n_states, n_bins, n_bins)
+        Continuous-space transition probabilities between states and bins.
     discrete_state_transition : NDArray[np.float64], shape (n_states, n_states)
+        Discrete state transition probabilities.
 
-    Return
-    ------
+    Returns
+    -------
     acausal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+        Acausal (forward-backward) posterior probabilities.
 
     """
     acausal_posterior = np.zeros_like(causal_posterior)
@@ -242,11 +263,14 @@ def scaled_likelihood(log_likelihood: NDArray[np.float64], axis: int = 1) -> NDA
     Parameters
     ----------
     log_likelihood : NDArray[np.float64], shape (n_time, n_bins)
-    axis : int
+        Log-likelihood values to be scaled.
+    axis : int, optional
+        Axis along which to find the maximum, by default 1.
 
     Returns
     -------
     scaled_log_likelihood : NDArray[np.float64], shape (n_time, n_bins)
+        Likelihood values scaled so that the maximum is 1.
 
     """
     max_log_likelihood = np.nanmax(log_likelihood, axis=axis, keepdims=True)
@@ -269,11 +293,14 @@ def mask(value: NDArray[np.float64], is_track_interior: NDArray[np.bool_]) -> ND
     Parameters
     ----------
     value : NDArray[np.float64], shape (..., n_bins)
+        Input values to be masked.
     is_track_interior : NDArray[np.bool_], shape (n_bins,)
+        Boolean array indicating which bins are part of the track interior.
 
     Returns
     -------
-    masked_value : NDArray[np.float64]
+    masked_value : NDArray[np.float64], shape (..., n_bins)
+        Values with non-track bins set to NaN.
 
     """
     try:
@@ -288,24 +315,27 @@ def check_converged(
     previous_log_likelihood: NDArray[np.float64],
     tolerance: float = 1e-4,
 ) -> Tuple[bool, bool]:
-    """We have converged if the slope of the log-likelihood function falls below 'tolerance',
+    """Check if log-likelihood has converged.
 
-    i.e., |f(t) - f(t-1)| / avg < tolerance,
-    where avg = (|f(t)| + |f(t-1)|)/2 and f(t) is log lik at iteration t.
+    We have converged if the slope of the log-likelihood function falls below
+    'tolerance', i.e., |f(t) - f(t-1)| / avg < tolerance, where
+    avg = (|f(t)| + |f(t-1)|)/2 and f(t) is log likelihood at iteration t.
 
     Parameters
     ----------
     log_likelihood : NDArray[np.float64]
-        Current log likelihood
+        Current log likelihood values.
     previous_log_likelihood : NDArray[np.float64]
-        Previous log likelihood
+        Previous log likelihood values.
     tolerance : float, optional
-        threshold for similarity, by default 1e-4
+        Threshold for similarity, by default 1e-4.
 
     Returns
     -------
     is_converged : bool
+        Whether the log-likelihood has converged.
     is_increasing : bool
+        Whether the log-likelihood is increasing.
 
     """
     delta_log_likelihood = abs(log_likelihood - previous_log_likelihood)
@@ -334,13 +364,18 @@ def _causal_decode_gpu(
         Parameters
         ----------
         initial_conditions : NDArray[np.float64], shape (n_bins,)
+            Initial probability distribution over state bins.
         state_transition : NDArray[np.float64], shape (n_bins, n_bins)
+            Transition probability matrix between state bins.
         likelihood : NDArray[np.float64], shape (n_time, n_bins)
+            Likelihood values for each time point and state bin.
 
         Returns
         -------
         posterior : NDArray[np.float64], shape (n_time, n_bins)
+            Posterior probability distribution over time and state bins.
         log_data_likelihood : float
+            Log-likelihood of the observed data.
 
         """
 
@@ -373,11 +408,14 @@ def _acausal_decode_gpu(
         Parameters
         ----------
         causal_posterior : NDArray[np.float64], shape (n_time, n_bins, 1)
+            Causal (forward-pass) posterior probabilities.
         state_transition : NDArray[np.float64], shape (n_bins, n_bins)
+            Transition probability matrix between state bins.
 
-        Return
-        ------
+        Returns
+        -------
         acausal_posterior : NDArray[np.float64], shape (n_time, n_bins, 1)
+            Acausal (forward-backward) posterior probabilities.
 
         """
         causal_posterior = cp.asarray(causal_posterior, dtype=cp.float32)
@@ -414,14 +452,20 @@ def _causal_classify_gpu(
         Parameters
         ----------
         initial_conditions : NDArray[np.float64], shape (n_states, n_bins, 1)
+            Initial probability distribution for each state and spatial bin.
         continuous_state_transition : NDArray[np.float64], shape (n_states, n_states, n_bins, n_bins)
+            Continuous-space transition probabilities between states and bins.
         discrete_state_transition : NDArray[np.float64], shape (n_states, n_states)
+            Discrete state transition probabilities.
         likelihood : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+            Likelihood values for each time point, state, and spatial bin.
 
         Returns
         -------
         causal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+            Causal posterior probabilities over time, states, and spatial bins.
         log_data_likelihood : float
+            Log-likelihood of the observed data.
 
         """
         initial_conditions = cp.asarray(initial_conditions, dtype=cp.float32)
@@ -467,12 +511,16 @@ def _acausal_classify_gpu(
         Parameters
         ----------
         causal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+            Causal (forward-pass) posterior probabilities.
         continuous_state_transition : NDArray[np.float64], shape (n_states, n_states, n_bins, n_bins)
+            Continuous-space transition probabilities between states and bins.
         discrete_state_transition : NDArray[np.float64], shape (n_states, n_states)
+            Discrete state transition probabilities.
 
-        Return
-        ------
+        Returns
+        -------
         acausal_posterior : NDArray[np.float64], shape (n_time, n_states, n_bins, 1)
+            Acausal (forward-backward) posterior probabilities.
 
         """
         causal_posterior = cp.asarray(causal_posterior, dtype=cp.float32)