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`. The time axis natively supports `astropy.time.Time` objects through a custom Matplotlib units converter (`ConciseAstropyConverter`), allowing use of `~matplotlib.axes.Axes.set_xlim` with `Time` objects.

radiospectra/mixins.py

@@ -1,7 +1,13 @@
+from contextlib import contextmanager
+
+import matplotlib.units as munits
+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,6 +43,64 @@ def _set_axis_converter(axis, converter):
             axis.converter = converter
 
 
+class ConciseAstropyConverter(munits.ConversionInterface):
+    """
+    Matplotlib unit converter for `astropy.time.Time` that uses concise date formatting.
+
+    Notes
+    -----
+    This converter turns times into Matplotlib internal date floats. It does
+    not handle leap seconds or nanosecond precision for the entire age of
+    the universe.
+    """
+
+    @staticmethod
+    def axisinfo(unit, axis):
+        locator = AutoDateLocator()
+        formatter = ConciseDateFormatter(locator)
+        return munits.AxisInfo(majloc=locator, majfmt=formatter, label=f"Time ({unit})")
+
+    @staticmethod
+    def convert(value, unit, axis):
+        if isinstance(value, Time):
+            return value.plot_date
+
+        if isinstance(value, (list, tuple, np.ndarray)):
+            # If it's already a Time object (could be array-valued), use its plot_date
+            if isinstance(value, Time):
+                return value.plot_date
+            # Otherwise iterate over the sequence
+            converted = [v.plot_date if isinstance(v, Time) else v for v in np.asarray(value, dtype=object).flat]
+
+            # Matplotlib dates are floats; let DateConverter handle any remaining datetimes
+            converted_val = np.asarray(converted).reshape(np.shape(value))
+            return DateConverter.convert(converted_val, unit, axis)
+
+        return DateConverter.convert(value, unit, axis)
+
+    @staticmethod
+    def default_units(x, axis):
+        if isinstance(x, Time):
+            return x.scale.upper()
+        return None
+
+
+@contextmanager
+def concise_time_support():
+    """
+    Context manager to enable concise time formatting for `astropy.time.Time` objects.
+    """
+    original_converter = munits.registry.get(Time)
+    munits.registry[Time] = ConciseAstropyConverter()
+    try:
+        yield
+    finally:
+        if original_converter is None:
+            del munits.registry[Time]
+        else:
+            munits.registry[Time] = original_converter
+
+
 class PcolormeshPlotMixin:
     """
     Class provides plotting functions using `~pcolormesh`.
@@ -59,9 +123,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 +136,18 @@ def plot(self, axes=None, **kwargs):
 
         axes.set_title(title)
 
-        with time_support(), quantity_support():
+        with concise_time_support(), 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)
-
             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])
-            fig.autofmt_xdate()
 
         # Set current axes/image if pyplot is being used (makes colorbar work)
         for i in plt.get_fignums():
@@ -110,22 +167,21 @@ def plotim(self, fig=None, axes=None, **kwargs):
         if axes is None:
             fig, axes = plt.subplots()
 
-        with time_support(), quantity_support():
+        with concise_time_support(), 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)
-
             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")
             im = NonUniformImage(axes, interpolation="none", **kwargs)
-            im.set_data(axes.convert_xunits(self.times), frequencies, self.data)
+            # NonUniformImage does not use the axis converter itself,
+            # so we manually convert the explicit input times down to floats.
+            times_numeric = axes.xaxis.convert_units(self.times)
+            im.set_data(times_numeric, frequencies, self.data)
             axes.add_image(im)
             axes.set_xlim(self.times[0], self.times[-1])
             axes.set_ylim(frequencies[0], frequencies[-1])

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()