Merge pull request #80 from hayesla/plotting_fix

Fixes #79

CHANGELOG.rst

@@ -4,8 +4,9 @@
 Breaking Changes
 ----------------
 
-- The old ``Spectrogram`` class has been removed
-- The new ``Spectrogram2`` class has been renamed to ``Spectrogram``.
+- The old ``Spectrogram`` class has been removed. (`#76 <https://github.com/sunpy/radiospectra/pull/76>`__)
+- The new ``Spectrogram2`` class has been renamed to ``Spectrogram``. (`#76 <https://github.com/sunpy/radiospectra/pull/76>`__)
+- Adding colorbar functionality to ``plot`` (`#80 <https://github.com/sunpy/radiospectra/pull/80>`__)
 
 0.4.0 (2022-05-24)
 ==================

radiospectra/mixins.py

@@ -16,6 +16,7 @@ def plot(self, axes=None, **kwargs):
 
         Returns
         -------
+        `matplotlib.collections.QuadMesh`
         """
         import matplotlib.dates as mdates
         from matplotlib import pyplot as plt
@@ -36,13 +37,19 @@ def plot(self, axes=None, **kwargs):
 
         axes.set_title(title)
         axes.plot(self.times.datetime[[0, -1]], self.frequencies[[0, -1]], linestyle="None", marker="None")
-        axes.pcolormesh(self.times.datetime, self.frequencies.value, data[:-1, :-1], shading="auto", **kwargs)
+        ret = axes.pcolormesh(self.times.datetime, self.frequencies.value, data[:-1, :-1], shading="auto", **kwargs)
         axes.set_xlim(self.times.datetime[0], self.times.datetime[-1])
         locator = mdates.AutoDateLocator(minticks=4, maxticks=8)
         formatter = mdates.ConciseDateFormatter(locator)
         axes.xaxis.set_major_locator(locator)
         axes.xaxis.set_major_formatter(formatter)
         fig.autofmt_xdate()
+        # Set current axes/image if pyplot is being used (makes colorbar work)
+        for i in plt.get_fignums():
+            if axes in plt.figure(i).axes:
+                plt.sca(axes)
+                plt.sci(ret)
+        return ret
 
 
 class NonUniformImagePlotMixin:

radiospectra/spectrogram/sources/eovsa.py

@@ -18,6 +18,7 @@ class EOVSASpectrogram(GenericSpectrogram):
     >>> spec  #doctest: +REMOTE_DATA
     <EOVSASpectrogram OWENS VALLEY, EOVSA, EOVSA 1105371.117591858 kHz - 17979686.737060547 kHz, 2021-05-07T13:48:20.999 to 2021-05-08T01:50:59.999>
     >>> spec.plot()  #doctest: +REMOTE_DATA
+    <matplotlib.collections.QuadMesh object at ...>
     """
 
     def __init__(self, data, meta, **kwargs):

radiospectra/spectrogram/sources/rstn.py

@@ -18,6 +18,7 @@ class RSTNSpectrogram(GenericSpectrogram):
     >>> spec  #doctest: +REMOTE_DATA
     <RSTNSpectrogram LEARMONTH, RSTN, RSTN 25000.0 kHz - 180000.0 kHz, 2017-09-06T22:31:51.000 to 2017-09-07T10:06:36.000>
     >>> spec.plot()  #doctest: +REMOTE_DATA
+    <matplotlib.collections.QuadMesh object at ...>
     """
 
     @classmethod

radiospectra/spectrogram/spectrogram_factory.py

@@ -466,7 +466,7 @@ def _read_idl_sav(file, instrument=None):
             }
             return data, meta
         else:
-            raise ValueError(f"Unrecognised IDL .sav file: {file}")
+            raise ValueError(f"Unrecognized IDL .sav file: {file}")
 
 
 Spectrogram = SpectrogramFactory(registry=GenericSpectrogram._registry, default_widget_type=GenericSpectrogram)