Fix #160: Improve default time axis formatting using ConciseDateFormatter

changelog/163.feature.rst

@@ -0,0 +1 @@
+Improved default time axis formatting on spectrogram plots by using `matplotlib.dates.ConciseDateFormatter` with a custom matplotlib units converter for `astropy.time.Time`, preserving support for both `Time` and `datetime` values when setting axis limits (e.g., via `~matplotlib.axes.Axes.set_xlim`).

radiospectra/mixins.py

@@ -1,7 +1,10 @@
+import numpy as np
 from matplotlib import pyplot as plt
+from matplotlib.dates import AutoDateLocator, ConciseDateFormatter, DateConverter
 from matplotlib.image import NonUniformImage
 
-from astropy.visualization import quantity_support, time_support
+from astropy.time import Time
+from astropy.visualization import quantity_support
 
 
 def _get_axis_converter(axis):
@@ -37,7 +40,62 @@ def _set_axis_converter(axis, converter):
             axis.converter = converter
 
 
-class PcolormeshPlotMixin:
+class _TimeDateConverter(DateConverter):
+    """
+    Converter that supports both matplotlib date values and astropy Time.
+    """
+
+    @staticmethod
+    def _is_time_sequence(value):
+        return (
+            isinstance(value, (list, tuple, np.ndarray))
+            and np.size(value)
+            and all(isinstance(v, Time) for v in np.asarray(value, dtype=object).flat)
+        )
+
+    @staticmethod
+    def default_units(x, axis):
+        if isinstance(x, Time) or _TimeDateConverter._is_time_sequence(x):
+            return None
+        return DateConverter.default_units(x, axis)
+
+    @staticmethod
+    def convert(value, unit, axis):
+        if isinstance(value, Time):
+            return value.plot_date
+
+        if _TimeDateConverter._is_time_sequence(value):
+            converted = [v.plot_date for v in np.asarray(value, dtype=object).flat]
+            if isinstance(value, np.ndarray):
+                return np.asarray(converted, dtype=float).reshape(value.shape)
+            return converted
+
+        return DateConverter.convert(value, unit, axis)
+
+
+_TIME_DATE_CONVERTER = _TimeDateConverter()
+
+
+class _TimeAxisMixin:
+    def _set_time_converter(self, axes):
+        """
+        Ensure the x-axis supports both `~astropy.time.Time` and datetime inputs.
+        """
+        if not getattr(axes.xaxis, "_converter_is_explicit", False):
+            _set_axis_converter(axes.xaxis, _TIME_DATE_CONVERTER)
+
+    def _setup_time_axis(self, axes):
+        """
+        Apply `~matplotlib.dates.ConciseDateFormatter` to the x-axis.
+        """
+        locator = AutoDateLocator()
+        formatter = ConciseDateFormatter(locator)
+        axes.xaxis.set_major_locator(locator)
+        axes.xaxis.set_major_formatter(formatter)
+        axes.set_xlabel(f"Time ({self.times.scale})")
+
+
+class PcolormeshPlotMixin(_TimeAxisMixin):
     """
     Class provides plotting functions using `~pcolormesh`.
     """
@@ -59,9 +117,7 @@ def plot(self, axes=None, **kwargs):
         """
 
         if axes is None:
-            fig, axes = plt.subplots()
-        else:
-            fig = axes.get_figure()
+            _, axes = plt.subplots()
 
         if hasattr(self.data, "value"):
             data = self.data.value
@@ -74,23 +130,23 @@ def plot(self, axes=None, **kwargs):
 
         axes.set_title(title)
 
-        with time_support(), quantity_support():
+        with quantity_support():
             # Pin existing converters to avoid warnings when re-plotting on shared axes.
             converter_y = _get_axis_converter(axes.yaxis)
             if converter_y is not None and not getattr(axes.yaxis, "_converter_is_explicit", False):
                 _set_axis_converter(axes.yaxis, converter_y)
 
-            converter_x = _get_axis_converter(axes.xaxis)
-            if converter_x is not None and not getattr(axes.xaxis, "_converter_is_explicit", False):
-                _set_axis_converter(axes.xaxis, converter_x)
+            self._set_time_converter(axes)
 
-            axes.plot(self.times[[0, -1]], self.frequencies[[0, -1]], linestyle="None", marker="None")
+            times_plot_date = self.times.plot_date
+            axes.plot(times_plot_date[[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)
+                ret = axes.pcolormesh(times_plot_date, 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])
-            fig.autofmt_xdate()
+                ret = axes.pcolormesh(times_plot_date, self.frequencies, data[:-1, :-1], shading="auto", **kwargs)
+            axes.set_xlim(times_plot_date[0], times_plot_date[-1])
+
+        self._setup_time_axis(axes)
 
         # Set current axes/image if pyplot is being used (makes colorbar work)
         for i in plt.get_fignums():
@@ -100,7 +156,7 @@ def plot(self, axes=None, **kwargs):
         return ret
 
 
-class NonUniformImagePlotMixin:
+class NonUniformImagePlotMixin(_TimeAxisMixin):
     """
     Class provides plotting functions using `NonUniformImage`.
     """
@@ -110,22 +166,23 @@ def plotim(self, fig=None, axes=None, **kwargs):
         if axes is None:
             fig, axes = plt.subplots()
 
-        with time_support(), quantity_support():
+        with quantity_support():
             # Pin existing converters to avoid warnings when re-plotting on shared axes.
             converter_y = _get_axis_converter(axes.yaxis)
             if converter_y is not None and not getattr(axes.yaxis, "_converter_is_explicit", False):
                 _set_axis_converter(axes.yaxis, converter_y)
 
-            converter_x = _get_axis_converter(axes.xaxis)
-            if converter_x is not None and not getattr(axes.xaxis, "_converter_is_explicit", False):
-                _set_axis_converter(axes.xaxis, converter_x)
+            self._set_time_converter(axes)
 
             axes.yaxis.update_units(self.frequencies)
             frequencies = axes.yaxis.convert_units(self.frequencies)
 
-            axes.plot(self.times[[0, -1]], self.frequencies[[0, -1]], linestyle="None", marker="None")
+            times_plot_date = self.times.plot_date
+            axes.plot(times_plot_date[[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)
+            im.set_data(times_plot_date, frequencies, self.data)
             axes.add_image(im)
-            axes.set_xlim(self.times[0], self.times[-1])
+            axes.set_xlim(times_plot_date[0], times_plot_date[-1])
             axes.set_ylim(frequencies[0], frequencies[-1])
+
+        self._setup_time_axis(axes)

radiospectra/spectrogram/tests/test_spectrogrambase.py

@@ -2,8 +2,10 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+from matplotlib.dates import ConciseDateFormatter
 
 import astropy.units as u
+from astropy.time import Time
 
 
 def test_plot_mixed_frequency_units_on_same_axes(make_spectrogram):
@@ -53,7 +55,9 @@ def test_plotim(make_spectrogram):
     plt.close("all")
 
     _, x_values, y_values, image = set_data.call_args.args
+    expected_times = rad_im.times.plot_date
     assert len(x_values) == len(rad_im.times)
+    np.testing.assert_allclose(x_values, expected_times)
     np.testing.assert_allclose(y_values, rad_im.frequencies.value)
     np.testing.assert_allclose(image, rad_im.data)
 
@@ -103,21 +107,21 @@ def test_plot_instrument_detector_differ(make_spectrogram):
     plt.close("all")
 
 
-def test_plot_uses_time_support_for_datetime_conversion(make_spectrogram):
-    """Plotting with non-UTC time scale should use time_support."""
+def test_plot_uses_plot_date_conversion(make_spectrogram):
+    """Plotting with non-UTC time scale should use matplotlib plot-date conversion."""
     spec = make_spectrogram(np.linspace(10, 40, 4) * u.MHz, scale="tt")
 
     mesh = spec.plot()
     x_limits = np.array(mesh.axes.get_xlim())
-    expected_tt_limits = mesh.axes.convert_xunits(spec.times[[0, -1]])
+    expected_tt_limits = spec.times.plot_date[[0, -1]]
 
     plt.close(mesh.axes.figure)
 
     np.testing.assert_allclose(x_limits, expected_tt_limits)
 
 
-def test_plotim_uses_time_support_for_datetime_conversion(make_spectrogram):
-    """plotim with non-UTC time scale should use time_support."""
+def test_plotim_uses_plot_date_conversion(make_spectrogram):
+    """plotim with non-UTC time scale should use matplotlib plot-date conversion."""
     spec = make_spectrogram(np.linspace(10, 40, 4) * u.MHz, scale="tt")
     fig, axes = plt.subplots()
 
@@ -130,8 +134,108 @@ def test_plotim_uses_time_support_for_datetime_conversion(make_spectrogram):
     plt.close(fig)
 
     _, x_values, y_values, image = set_data.call_args.args
-    expected_tt = axes.convert_xunits(spec.times)
+    expected_tt = spec.times.plot_date
 
     np.testing.assert_allclose(x_values, expected_tt)
     np.testing.assert_allclose(y_values, spec.frequencies.value)
     np.testing.assert_allclose(image, spec.data)
+
+
+def test_plot_accepts_astropy_time_xlim(make_spectrogram):
+    """plot() should still accept astropy Time values in set_xlim()."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    mesh = rad.plot()
+    mesh.axes.set_xlim(rad.times[0], rad.times[-1])
+    x_limits = np.array(mesh.axes.get_xlim())
+    plt.close(mesh.axes.figure)
+
+    np.testing.assert_allclose(x_limits, rad.times.plot_date[[0, -1]])
+
+
+def test_plot_accepts_datetime_xlim(make_spectrogram):
+    """plot() should accept datetime values in set_xlim()."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    mesh = rad.plot()
+    mesh.axes.set_xlim(rad.times.datetime[0], rad.times.datetime[-1])
+    x_limits = np.array(mesh.axes.get_xlim())
+    plt.close(mesh.axes.figure)
+
+    np.testing.assert_allclose(x_limits, rad.times.plot_date[[0, -1]])
+
+
+def test_plotim_accepts_astropy_time_xlim(make_spectrogram):
+    """plotim() should still accept astropy Time values in set_xlim()."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    fig, axes = plt.subplots()
+    with (
+        mock.patch("matplotlib.image.NonUniformImage.set_interpolation", autospec=True),
+        mock.patch("matplotlib.image.NonUniformImage.set_data", autospec=True),
+    ):
+        rad.plotim(axes=axes)
+    axes.set_xlim(rad.times[0], rad.times[-1])
+    x_limits = np.array(axes.get_xlim())
+    plt.close(fig)
+
+    np.testing.assert_allclose(x_limits, rad.times.plot_date[[0, -1]])
+
+
+def test_plotim_accepts_datetime_xlim(make_spectrogram):
+    """plotim() should accept datetime values in set_xlim()."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    fig, axes = plt.subplots()
+    with (
+        mock.patch("matplotlib.image.NonUniformImage.set_interpolation", autospec=True),
+        mock.patch("matplotlib.image.NonUniformImage.set_data", autospec=True),
+    ):
+        rad.plotim(axes=axes)
+    axes.set_xlim(rad.times.datetime[0], rad.times.datetime[-1])
+    x_limits = np.array(axes.get_xlim())
+    plt.close(fig)
+
+    np.testing.assert_allclose(x_limits, rad.times.plot_date[[0, -1]])
+
+
+def test_plot_handles_leap_seconds(make_spectrogram):
+    """plot() should not fail for times that include a leap second."""
+    leap_times = Time("2016-12-31T23:59:58", scale="utc") + np.arange(4) * u.s
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz, times=leap_times)
+    mesh = rad.plot()
+    x_limits = np.array(mesh.axes.get_xlim())
+    plt.close(mesh.axes.figure)
+
+    np.testing.assert_allclose(x_limits, leap_times.plot_date[[0, -1]])
+
+
+def test_plot_uses_concise_date_formatter(make_spectrogram):
+    """plot() should apply ConciseDateFormatter to the x-axis."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    mesh = rad.plot()
+    formatter = mesh.axes.xaxis.get_major_formatter()
+    plt.close("all")
+
+    assert isinstance(formatter, ConciseDateFormatter)
+
+
+def test_plotim_uses_concise_date_formatter(make_spectrogram):
+    """plotim() should apply ConciseDateFormatter to the x-axis."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    fig, axes = plt.subplots()
+    with (
+        mock.patch("matplotlib.image.NonUniformImage.set_interpolation", autospec=True),
+        mock.patch("matplotlib.image.NonUniformImage.set_data", autospec=True),
+    ):
+        rad.plotim(axes=axes)
+    formatter = axes.xaxis.get_major_formatter()
+    plt.close("all")
+
+    assert isinstance(formatter, ConciseDateFormatter)
+
+
+def test_plot_xlabel_includes_time_scale(make_spectrogram):
+    """plot() should set the x-label to include the time scale."""
+    rad = make_spectrogram(np.array([10, 20, 30, 40]) * u.kHz)
+    mesh = rad.plot()
+    xlabel = mesh.axes.get_xlabel()
+    plt.close("all")
+
+    assert "utc" in xlabel.lower()