Fix shared x-axis limits for spectrogram plots

changelog/219.bugfix.rst

@@ -0,0 +1 @@
+Fixed plotting multiple spectrograms on shared x-axes so each plot expands the shared time range instead of replacing it.

radiospectra/mixins.py

@@ -39,6 +39,32 @@ def _set_axis_converter(axis, converter):
             axis.converter = converter
 
 
+def _shared_x_limits(axes):
+    """
+    Return the current shared x limits if any shared sibling has plotted data.
+    """
+    siblings = axes.get_shared_x_axes().get_siblings(axes)
+    if any(sibling.has_data() for sibling in siblings):
+        return axes.get_xlim()
+    return None
+
+
+def _set_x_limits(axes, x_values, previous_limits):
+    """
+    Set x limits, expanding previous shared-axis limits when present.
+    """
+    if previous_limits is None:
+        axes.set_xlim(x_values[0], x_values[-1])
+        return
+
+    converted_limits = axes.convert_xunits(x_values)
+    lower = min(previous_limits[0], previous_limits[1], converted_limits[0], converted_limits[-1])
+    upper = max(previous_limits[0], previous_limits[1], converted_limits[0], converted_limits[-1])
+    if previous_limits[0] > previous_limits[1]:
+        lower, upper = upper, lower
+    axes.set_xlim(lower, upper)
+
+
 class PcolormeshPlotMixin:
     """
     Class provides plotting functions using `~matplotlib.pyplot.pcolormesh`.
@@ -86,12 +112,14 @@ def plot(self, axes=None, **kwargs):
             if converter_x is not None and not getattr(axes.xaxis, "_converter_is_explicit", False):
                 _set_axis_converter(axes.xaxis, converter_x)
 
+            previous_x_limits = _shared_x_limits(axes)
+
             axes.plot(self.times[[0, -1]], self.frequencies[[0, -1]], linestyle="None", marker="None")
             if self.times.shape[0] == self.data.shape[0] and self.frequencies.shape[0] == self.data.shape[1]:
                 ret = axes.pcolormesh(self.times, self.frequencies, data, shading="auto", **kwargs)
             else:
                 ret = axes.pcolormesh(self.times, self.frequencies, data[:-1, :-1], shading="auto", **kwargs)
-            axes.set_xlim(self.times[0], self.times[-1])
+            _set_x_limits(axes, self.times[[0, -1]], previous_x_limits)
             fig.autofmt_xdate()
 
         # Set current axes/image if pyplot is being used (makes colorbar work)
@@ -132,9 +160,11 @@ def plotim(self, fig=None, axes=None, **kwargs):
             axes.yaxis.update_units(self.frequencies)
             frequencies = axes.yaxis.convert_units(self.frequencies)
 
+            previous_x_limits = _shared_x_limits(axes)
+
             axes.plot(self.times[[0, -1]], self.frequencies[[0, -1]], linestyle="None", marker="None")
             im = NonUniformImage(axes, interpolation="none", **kwargs)
             im.set_data(axes.convert_xunits(self.times), frequencies, self.data)
             axes.add_image(im)
-            axes.set_xlim(self.times[0], self.times[-1])
+            _set_x_limits(axes, self.times[[0, -1]], previous_x_limits)
             axes.set_ylim(frequencies[0], frequencies[-1])

radiospectra/spectrogram/tests/test_spectrogrambase.py

@@ -4,6 +4,7 @@
 import numpy as np
 
 import astropy.units as u
+from astropy.time import Time
 
 
 def test_plot_mixed_frequency_units_on_same_axes(make_spectrogram):
@@ -42,6 +43,27 @@ def test_plot_mixed_frequency_units_mhz_first(make_spectrogram):
     assert y_max >= 4, "y-axis should cover up to 4 MHz"
 
 
+def test_plot_shared_x_axes_keep_all_spectrogram_times(make_spectrogram):
+    """Shared x axes should not be limited to whichever spectrogram plots last."""
+    times1 = Time("2020-01-01T00:00:00", format="isot") + np.arange(4) * u.min
+    times2 = Time("2020-01-01T00:10:00", format="isot") + np.arange(4) * u.min
+    specs = [
+        make_spectrogram(np.linspace(10, 40, 4) * u.MHz, times=times1),
+        make_spectrogram(np.linspace(50, 80, 4) * u.MHz, times=times2),
+    ]
+
+    for order in ((0, 1), (1, 0)):
+        fig, axes = plt.subplots(2, 1, sharex=True)
+
+        for index in order:
+            specs[index].plot(axes=axes[index])
+
+        expected_limits = axes[0].convert_xunits(Time([times1[0], times2[-1]]))
+        np.testing.assert_allclose(axes[0].get_xlim(), expected_limits)
+        np.testing.assert_allclose(axes[1].get_xlim(), expected_limits)
+        plt.close(fig)
+
+
 def test_plotim(make_spectrogram):
     """Test NonUniformImagePlotMixin.plotim() executes without error."""
     rad_im = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)