Add coordinate lookup methods to GenericSpectrogram

changelog/129.feature.rst

@@ -0,0 +1,3 @@
+Add ``frequency_at_index``, ``time_at_index``, ``index_at_frequency``, and ``index_at_time``
+methods to `~radiospectra.spectrogram.GenericSpectrogram` for converting between
+array indices and physical coordinates.

radiospectra/spectrogram/spectrogrambase.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from radiospectra.exceptions import SpectraMetaValidationError
 from radiospectra.mixins import NonUniformImagePlotMixin, PcolormeshPlotMixin
 
@@ -101,6 +103,66 @@ def _validate_meta(self):
         if err_message:
             raise SpectraMetaValidationError("\n".join(err_message))
 
+    def frequency_at_index(self, index):
+        """
+        Return the frequency at a given index along the frequency axis.
+
+        Parameters
+        ----------
+        index : int
+            Index along the frequency axis.
+
+        Returns
+        -------
+        `~astropy.units.Quantity`
+        """
+        return self.frequencies[index]
+
+    def time_at_index(self, index):
+        """
+        Return the time at a given index along the time axis.
+
+        Parameters
+        ----------
+        index : int
+            Index along the time axis.
+
+        Returns
+        -------
+        `~astropy.time.Time`
+        """
+        return self.times[index]
+
+    def index_at_frequency(self, frequency):
+        """
+        Return the index along the frequency axis closest to a given frequency.
+
+        Parameters
+        ----------
+        frequency : `~astropy.units.Quantity`
+            The frequency to search for.
+
+        Returns
+        -------
+        int
+        """
+        return int(np.argmin(np.abs(self.frequencies - frequency)))
+
+    def index_at_time(self, time):
+        """
+        Return the index along the time axis closest to a given time.
+
+        Parameters
+        ----------
+        time : `~astropy.time.Time`
+            The time to search for.
+
+        Returns
+        -------
+        int
+        """
+        return int(np.argmin(np.abs(self.times - time)))
+
     def __repr__(self):
         return (
             f"<{self.__class__.__name__} {self.observatory}, {self.instrument}, {self.detector}"

radiospectra/spectrogram/tests/test_spectrogrambase.py

@@ -135,3 +135,34 @@ def test_plotim_uses_time_support_for_datetime_conversion(make_spectrogram):
     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_frequency_at_index(make_spectrogram):
+    spec = make_spectrogram(np.array([10, 20, 30, 40]) * u.MHz)
+
+    assert spec.frequency_at_index(0) == 10 * u.MHz
+    assert spec.frequency_at_index(2) == 30 * u.MHz
+
+
+def test_time_at_index(make_spectrogram):
+    spec = make_spectrogram(np.array([10, 20, 30, 40]) * u.MHz)
+
+    assert spec.time_at_index(0) == spec.times[0]
+    assert spec.time_at_index(3) == spec.times[3]
+
+
+def test_index_at_frequency(make_spectrogram):
+    spec = make_spectrogram(np.array([10, 20, 30, 40]) * u.MHz)
+
+    # 21 MHz is closer to 20 than 30, so index 1
+    assert spec.index_at_frequency(21 * u.MHz) == 1
+    # 39 MHz is closest to 40 MHz, index 3
+    assert spec.index_at_frequency(39 * u.MHz) == 3
+
+
+def test_index_at_time(make_spectrogram):
+    spec = make_spectrogram(np.array([10, 20, 30, 40]) * u.MHz)
+
+    # 10 seconds after times[1] is still closer to times[1] than times[2]
+    t = spec.times[1] + 10 * u.s
+    assert spec.index_at_time(t) == 1