[ENH] - Add colorbar to time-frequency plot

.github/workflows/build.yml

@@ -19,7 +19,7 @@ jobs:
       MODULE_NAME: neurodsp
     strategy:
       matrix:
-        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12"]
+        python-version: ["3.6", "3.7", "3.8", "3.9", "3.10", "3.11", "3.12", "3.13"]
 
     steps:
     - uses: actions/checkout@v4

neurodsp/plts/timefrequency.py

@@ -1,6 +1,7 @@
 """Plotting functions for neurodsp.timefrequency."""
 
 import numpy as np
+import matplotlib.pyplot as plt
 
 from neurodsp.plts.style import style_plot
 from neurodsp.plts.utils import check_ax, savefig
@@ -10,7 +11,8 @@
 
 @savefig
 @style_plot
-def plot_timefrequency(times, freqs, powers, x_ticks=5, y_ticks=5, ax=None, **kwargs):
+def plot_timefrequency(times, freqs, powers, x_ticks=5, y_ticks=5,
+                       colorbar=True, ax=None, **kwargs):
     """Plot a time-frequency representation of data.
 
     Parameters
@@ -26,6 +28,8 @@ def plot_timefrequency(times, freqs, powers, x_ticks=5, y_ticks=5, ax=None, **kw
         Defines the tick labels to add to the plot.
         If int, is the number of evenly sampled labels to add to the plot.
         If array_like, is a set of labels to add to the plot.
+    colorbar : bool, optional, default: True
+        Whether to add a colorbar to the plot.
     ax : matplotlib.Axes, optional
         Figure axes upon which to plot.
     **kwargs
@@ -70,3 +74,6 @@ def plot_timefrequency(times, freqs, powers, x_ticks=5, y_ticks=5, ax=None, **kw
     else:
         y_ticks_pos = [np.argmin(np.abs(freqs - val)) for val in y_ticks]
     ax.set(yticks=y_ticks_pos, yticklabels=y_ticks)
+
+    if colorbar:
+        plt.colorbar(ax.images[0])

setup.py

@@ -1,4 +1,4 @@
-"""NeuroDSP setup script."""
+"""Setup script for neurodsp."""
 
 import os
 from setuptools import setup, find_packages
@@ -51,11 +51,12 @@
         'Programming Language :: Python :: 3.10',
         'Programming Language :: Python :: 3.11',
         'Programming Language :: Python :: 3.12',
+        'Programming Language :: Python :: 3.13',
         ],
     platforms = 'any',
     project_urls = {
         'Documentation' : 'https://neurodsp-tools.github.io/',
         'Bug Reports' : 'https://github.com/neurodsp-tools/neurodsp/issues',
-        'Source' : 'https://github.com/neurodsp-tools/neurodsp'
+        'Source' : 'https://github.com/neurodsp-tools/neurodsp',
     },
 )