diff --git a/changelog/163.feature.rst b/changelog/163.feature.rst new file mode 100644 index 0000000..778321f --- /dev/null +++ b/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`). diff --git a/radiospectra/mixins.py b/radiospectra/mixins.py index 488c7cf..7dddcfe 100644 --- a/radiospectra/mixins.py +++ b/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) diff --git a/radiospectra/spectrogram/tests/test_spectrogrambase.py b/radiospectra/spectrogram/tests/test_spectrogrambase.py index ddd3a36..de69289 100644 --- a/radiospectra/spectrogram/tests/test_spectrogrambase.py +++ b/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()