Add and expand module docstrings throughout package

Comprehensive docstrings were added or expanded for all major modules, providing overviews, usage context, and parameter documentation. This improves code readability and helps users understand the purpose and functionality of each component in the replay_trajectory_classification package.

Author: edeno

replay_trajectory_classification/__init__.py

@@ -1,3 +1,21 @@
+"""Replay trajectory classification package.
+
+A Python package for decoding spatial position from neural activity and
+classifying trajectory types in hippocampal replay events using state-space
+modeling approaches.
+
+The main classes are:
+
+- SortedSpikesClassifier/Decoder: For spike-sorted neural data
+- ClusterlessClassifier/Decoder: For clusterless (unsorted) neural data
+- Environment: Spatial environment representation with discrete grids
+
+Examples
+--------
+>>> from replay_trajectory_classification import SortedSpikesClassifier
+>>> classifier = SortedSpikesClassifier()
+"""
+
 # flake8: noqa
 from track_linearization import (
     get_linearized_position,

replay_trajectory_classification/clusterless_simulation.py

@@ -1,4 +1,8 @@
-"""Simulate clusterless spikes and associated spike waveform features."""
+"""Simulate clusterless spikes and associated spike waveform features.
+
+This module provides functions for simulating clusterless (unsorted) spike
+data with associated waveform features used in multiunit decoding approaches.
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/continuous_state_transitions.py

@@ -1,4 +1,9 @@
-"""Classes for constructing different types of movement models."""
+"""Classes for constructing different types of movement models.
+
+This module provides continuous state transition models that describe how
+animals move through space, including random walks, empirical movement,
+and uniform transitions.
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/decoder.py

@@ -395,6 +395,25 @@ def convert_results_to_xarray(
 
 
 class SortedSpikesDecoder(_DecoderBase):
+    """Decoder for sorted spike data.
+
+    This class provides spatial decoding functionality for sorted (clustered)
+    spike trains using Bayesian state-space models.
+
+    Parameters
+    ----------
+    environment : Environment, optional
+        Spatial environment representation, by default Environment().
+    transition_type : EmpiricalMovement | RandomWalk | RandomWalkDirection1 |
+        RandomWalkDirection2 | Uniform, optional
+        Movement model for state transitions, by default RandomWalk().
+    initial_conditions_type : UniformInitialConditions, optional
+        Initial state probability distribution, by default UniformInitialConditions().
+    infer_track_interior : bool, optional
+        Whether to infer valid track areas from position data, by default True.
+
+    """
+
     def __init__(
         self,
         environment: Environment = Environment(environment_name=""),
@@ -583,23 +602,25 @@ def predict(
 
 
 class ClusterlessDecoder(_DecoderBase):
-    """Classifies neural population representation of position from multiunit
-    spikes and waveforms.
+    """Decoder for clusterless (multiunit) spike data.
+
+    This class provides spatial decoding functionality for clusterless spike
+    data using waveform features and marked point processes.
 
     Parameters
     ----------
     environment : Environment, optional
-        The spatial environment to fit
+        The spatial environment to fit, by default Environment().
     transition_type : EmpiricalMovement | RandomWalk | RandomWalkDirection1 |
-        RandomWalkDirection2 | Uniform
-        The continuous state transition matrix
+        RandomWalkDirection2 | Uniform, optional
+        The continuous state transition matrix, by default RandomWalk().
     initial_conditions_type : UniformInitialConditions, optional
-        The initial conditions class instance
+        The initial conditions class instance, by default UniformInitialConditions().
     infer_track_interior : bool, optional
-        Whether to infer the spatial geometry of track from position
-    clusterless_algorithm : str
-        The type of clusterless algorithm. See _ClUSTERLESS_ALGORITHMS for keys
-    clusterless_algorithm_params : dict
+        Whether to infer the spatial geometry of track from position, by default True.
+    clusterless_algorithm : str, optional
+        The type of clusterless algorithm. See _CLUSTERLESS_ALGORITHMS for keys.
+    clusterless_algorithm_params : dict, optional
         Parameters for the clusterless algorithms.
 
     """

replay_trajectory_classification/discrete_state_transitions.py

@@ -1,4 +1,9 @@
-"""Classes to generate transitions between categories."""
+"""Classes for generating discrete state transitions.
+
+This module provides discrete state transition models that define
+transition probabilities between different trajectory categories
+or behavioral states.
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/initial_conditions.py

@@ -1,4 +1,8 @@
-"""Classes for constructing the initial conditions for the state space models."""
+"""Classes for constructing initial conditions for state space models.
+
+This module provides classes that define initial probability distributions
+over states and spatial bins at the start of decoding/classification.
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/likelihoods/__init__.py

@@ -1,4 +1,16 @@
-"""Functions to fit the position intensities and likelihoods for each data type"""
+"""Likelihood functions for different neural data types.
+
+This subpackage provides likelihood estimation functions for various types
+of neural data including sorted spikes, clusterless spikes, and calcium
+imaging data.
+
+The module supports:
+
+- Sorted spikes: GLM and KDE-based likelihood estimation
+- Clusterless data: Multiunit likelihood functions with mark information
+- Calcium imaging: Gamma likelihood for calcium activity traces
+- GPU acceleration: CUDA implementations for performance
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/likelihoods/calcium_likelihood.py

@@ -1,8 +1,14 @@
-"""Calculate a Gamma likelihood for calcium imaging activity traces.
+"""Calculate Gamma likelihood for calcium imaging activity traces.
+
+This module provides functions for computing Gamma-distributed likelihoods
+from calcium imaging data, which is commonly used for decoding from calcium
+fluorescence signals.
 
 References
 ----------
-[1] Farhoodi, S., Plitt, M.H., Giocomo, L., and Eden, U.T. (2020). Estimating Fluctuations in Neural Representations of Uncertain Environments. 20. .
+.. [1] Farhoodi, S., Plitt, M.H., Giocomo, L., and Eden, U.T. (2020).
+   Estimating Fluctuations in Neural Representations of Uncertain
+   Environments.
 
 """
 

replay_trajectory_classification/likelihoods/diffusion.py

@@ -1,6 +1,9 @@
-"""Calculate diffusion distances by simulating diffusion at
-each position bin. This can be used for more accurate 2D kernel density
-estimates with boundaries."""
+"""Calculate diffusion distances for accurate boundary-aware density estimates.
+
+This module simulates diffusion at each position bin to provide more accurate
+2D kernel density estimates that properly handle spatial boundaries in complex
+environments.
+"""
 
 from __future__ import annotations
 

replay_trajectory_classification/likelihoods/multiunit_likelihood.py

@@ -1,5 +1,9 @@
-"""Estimates a marked point process likelihood where the marks are
-features of the spike waveform. Features are float32."""
+"""Estimate marked point process likelihood for clusterless spike data.
+
+This module provides functions for computing likelihoods from clusterless
+(multiunit) spike data where marks are features of the spike waveforms.
+Features are represented as float32 arrays.
+"""
 
 from __future__ import annotations