[WIP] Improve Raster Plot (#1017)

* add templates for layer-specific PSD and update the Drive-Dipole-Spectrogram (3x1) tempate to point to the layer-specific PSDs

* add layer to legend

* update plot_spikes_raster to 1) have Cell ID y-axis label, 2) accept show_legend argument that will optionally hide the legend, and 3) accept marker_size argument to change the spike marker sizes

---------

Co-authored-by: Austin E. Soplata <[email protected]>

Author: dylansdaniels

hnn_core/cell_response.py

@@ -293,8 +293,15 @@ def mean_rates(self, tstart, tstop, gid_ranges, mean_type='all'):
 
         return spike_rates
 
-    def plot_spikes_raster(self, trial_idx=None, ax=None, show=True,
-                           colors=None):
+    def plot_spikes_raster(
+            self,
+            trial_idx=None,
+            ax=None,
+            show=True,
+            colors=None,
+            show_legend=True,
+            marker_size=1.0,
+            ):
         """Plot the aggregate spiking activity according to cell type.
 
         Parameters
@@ -314,8 +321,13 @@ def plot_spikes_raster(self, trial_idx=None, ax=None, show=True,
             The matplotlib figure object.
         """
         return plot_spikes_raster(
-            cell_response=self, trial_idx=trial_idx, ax=ax, show=show,
-            colors=colors)
+            cell_response=self,
+            trial_idx=trial_idx,
+            ax=ax, show=show,
+            colors=colors,
+            show_legend=show_legend,
+            marker_size=marker_size,
+        )
 
     def plot_spikes_hist(self, trial_idx=None, ax=None, spike_types=None,
                          color=None, invert_spike_types=None, show=True,

hnn_core/viz.py

@@ -573,9 +573,16 @@ def plot_spikes_hist(cell_response, trial_idx=None, ax=None, spike_types=None,
     return ax.get_figure()
 
 
-def plot_spikes_raster(cell_response, trial_idx=None, ax=None, show=True,
-                       cell_types=None, colors=None,
-                       ):
+def plot_spikes_raster(
+        cell_response,
+        trial_idx=None,
+        ax=None,
+        show=True,
+        cell_types=None,
+        colors=None,
+        show_legend=True,
+        marker_size=1.0,
+        ):
     """Plot the aggregate spiking activity according to cell type.
 
     Parameters
@@ -588,10 +595,16 @@ def plot_spikes_raster(cell_response, trial_idx=None, ax=None, show=True,
         An axis object from matplotlib. If None, a new figure is created.
     show : bool
         If True, show the figure.
-    cell_types: list of str
+    cell_types : list of str
         List of cell types to plot
-    colors: list of str | None
+    colors : list of str | None
         Optional custom colors to plot. Default will use the color cycler.
+    show_legend : bool
+        If True, show the legend with colors for cell types
+    marker_size : float
+        Optional marker size to use when plotting spikes. Uses
+        "linelengths" argument of ax.eventplot, which accepts positive
+        numeric values only
 
     Returns
     -------
@@ -649,6 +662,26 @@ def plot_spikes_raster(cell_response, trial_idx=None, ax=None, show=True,
                                  f"Got {colors.keys()}")
             cell_colors.update(colors)
 
+    # validate show_legend argument
+    _validate_type(
+        show_legend,
+        bool,
+        'show_legend',
+        'bool'
+    )
+
+    # validate marker_size
+    _validate_type(
+        marker_size,
+        (float, int),
+        'marker_size',
+        'float or int'
+    )
+
+    # if marker_size is <= 0, set it to the default value of 1.0
+    if marker_size <= 0:
+        marker_size = 1.0
+
     # Extract desired trials
     spike_times = np.concatenate(
         np.array(cell_response._spike_times, dtype=object)[trial_idx])
@@ -672,19 +705,33 @@ def plot_spikes_raster(cell_response, trial_idx=None, ax=None, show=True,
 
         if cell_type_times:
             events.append(
-                ax.eventplot(cell_type_times, lineoffsets=cell_type_ypos,
-                             color=color,
-                             label=cell_type, linelengths=1))
+                ax.eventplot(
+                    cell_type_times,
+                    lineoffsets=cell_type_ypos,
+                    color=color,
+                    label=cell_type,
+                    linelengths=marker_size
+                )
+            )
         else:
             # Blank plot for no spiking
             events.append(
                 ax.eventplot([-1], lineoffsets=[-1],
                              color=color,
                              label=cell_type, linelengths=1))
 
+    # set legend
     ax.legend(handles=[e[0] for e in events], loc=1)
+    if not show_legend:
+        ax.get_legend().remove()
+
+    # set axis labels
     ax.set_xlabel('Time (ms)')
-    ax.get_yaxis().set_visible(False)
+    ax.set_ylabel('Cell ID')
+
+    # hide y-axis ticks and tick labels
+    ax.set_yticklabels([])
+    ax.tick_params(axis='y', length=0)
 
     if len(cell_response.times) > 0:
         ax.set_xlim(left=0, right=cell_response.times[-1])