diff --git a/changelog/163.feature.rst b/changelog/163.feature.rst new file mode 100644 index 00000000..8ff3d286 --- /dev/null +++ b/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. diff --git a/radiospectra/mixins.py b/radiospectra/mixins.py index 488c7cf2..337af2c4 100644 --- a/radiospectra/mixins.py +++ b/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]) diff --git a/radiospectra/spectrogram/tests/test_spectrogrambase.py b/radiospectra/spectrogram/tests/test_spectrogrambase.py index ddd3a36b..de692891 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()