docs: 📝 add | update docstrings

Author: JaerongA

element_array_ephys/plotting/probe_level.py

@@ -4,8 +4,16 @@
 import seaborn as sns
 
 
-def plot_raster(units, spike_times) -> matplotlib.figure.Figure:
+def plot_raster(units: np.ndarray, spike_times: np.ndarray) -> matplotlib.figure.Figure:
+    """Population raster plots.
 
+    Args:
+        units (np.ndarray): Recorded units.
+        spike_times (np.ndarray): Spike timestamps in seconds.
+
+    Returns:
+        matplotlib.figure.Figure: matplotlib figure object.
+    """
     units = np.arange(1, len(units) + 1)
     x = np.hstack(spike_times)
     y = np.hstack([np.full_like(s, u) for u, s in zip(units, spike_times)])
@@ -26,6 +34,16 @@ def plot_raster(units, spike_times) -> matplotlib.figure.Figure:
 def plot_driftmap(
     spike_times: np.ndarray, spike_depths: np.ndarray, colormap="gist_heat_r"
 ) -> matplotlib.figure.Figure:
+    """Plot unit activity drift map all units recorded in a given shank of a probe.
+
+    Args:
+        spike_times (np.ndarray): Spike timestamps in seconds.
+        spike_depths (np.ndarray): The depth of the electrode where the spike was found in μm.
+        colormap (str, optional): Colormap. Defaults to "gist_heat_r".
+
+    Returns:
+        matplotlib.figure.Figure: matplotlib figure object.
+    """
 
     spike_times = np.hstack(spike_times)
     spike_depths = np.hstack(spike_depths)
@@ -81,4 +99,4 @@ def plot_driftmap(
     ax_spkcount.set_ylim(depth_edges[0], depth_edges[-1])
     sns.despine()
 
-    return fig
+    return fig

element_array_ephys/plotting/unit_level.py

@@ -99,7 +99,7 @@ def plot_depth_waveforms(
     """Plot waveforms
 
     Args:
-        ephys (Module): Imported ephys module.
+        ephys (Module): Imported ephys module object.
         unit_key (dict[str, Any]): Key dictionary from ephys.CuratedClustering.Unit table.
         y_range (float, optional): Vertical range to show waveforms relative to the peak waveform in μm. Defaults to 60.