Add convenience coordinate methods to NDCube

changelog/709.feature.rst

@@ -0,0 +1 @@
+Add convenience methods, `~ndcube.NDCube.celestial`, `~ndcube.NDCube.time`, `~ndcube.NDCube.spectral`, `~ndcube.NDCube.stokes`, to return coordinates for array indices from relevant axes for common coordinate types.

ndcube/ndcube.py

@@ -1247,6 +1247,117 @@
             raise ValueError("All axes are of length 1, therefore we will not squeeze NDCube to become a scalar. Use `axis=` keyword to specify a subset of axes to squeeze.")
         return self[tuple(item)]
 
+    @property
+    def celestial(self):
+        """
+        Returns celestial coordinates for all array indices in relevant axes.
+
+        Celestial world axis physical type name must contain 'pos.'.
+
+        Parameters
+        ----------
+        wcs : `astropy.wcs.wcsapi.BaseHighLevelWCS`
+           The WCS object with which to calculate the coordinates. Must be
+           self.wcs, self.extra_coords, or self.combined_wcs
+
+        Returns
+        -------
+        : `astropy.coordinates.SkyCoord`
+        """
+        return self._get_coords_by_word("pos")
+
+    @property
+    def time(self):
+        """
+        Returns time coordinates for all array indices in relevant axes.
+
+        Time world axis physical type name must contain 'time'.
+
+        Parameters
+        ----------
+        wcs : `astropy.wcs.wcsapi.BaseHighLevelWCS`
+           The WCS object with which to calculate the coordinates. Must be
+           self.wcs, self.extra_coords, or self.combined_wcs
+
+        Returns
+        -------
+        : `astropy.time.Time`
+        """
+        return self._get_coords_by_word("time")
+
+    @property
+    def spectral(self):
+        """
+        Returns spectral coordinates from WCS.
+
+        Spectral world axis physical type name must contain 'em.'.
+
+        Parameters
+        ----------
+        wcs : `astropy.wcs.wcsapi.BaseHighLevelWCS`
+           The WCS object with which to calculate the coordinates. Must be
+           self.wcs, self.extra_coords, or self.combined_wcs
+
+        Returns
+        -------
+        : `astropy.coordinates.SpectralCoord` or `astropy.units.Quantity`
+        """
+        return self._get_coords_by_word("em")
+
+    @property
+    def stokes(self):
+        """
+        Returns stokes polarization for all array indices in relevant axes.
+
+        Stokes world axis physical type name must contain 'stokes'.
+
+        Parameters
+        ----------
+        wcs : `astropy.wcs.wcsapi.BaseHighLevelWCS`
+           The WCS object with which to calculate the coordinates. Must be
+           self.wcs, self.extra_coords, or self.combined_wcs
+
+        Returns
+        -------
+        :
+        """
+        return self._get_coords_by_word("stokes")
+
+    def _get_coords_by_word(self, coord_words):
+            """
+            Returns coordinates from a WCS corresponding to a world axis physical type.
+            """
+            if isinstance(coord_words, str):
+                coord_words = [coord_words]
+            # Check if words are in physical types of main WCS and extra coords.
+            wcs_phys_types = []
+            ec_phys_types = []
+            for wcs, physical_types in zip([self.wcs, self.extra_coords.wcs], [wcs_phys_types, ec_phys_types]):
+                if wcs is not None:
+                    for name in wcs.world_axis_physical_types:
+                        words = set(name.replace(":", ".").split("."))
+                        if any(word in words for word in coord_words):
+                            physical_types.append(name)
+            # Determine type of WCS to use based on whether the words were
+            # found in WCS and/or extra coords.
+            if len(wcs_phys_types) > 0:
+               if len(ec_phys_types) > 0:
+                   wcs = self.combined_wcs
+                   physical_types = wcs_phys_types + ec_phys_types
+               else:
+                   wcs = self.wcs
+                   physical_types = wcs_phys_types
+            elif len(ec_phys_types) > 0:
+                wcs = self.extra_coords
+                physical_types = ec_phys_types
+            else:
+                return None
+            # Calculate coordinates.
+            coords = self.axis_world_coords(*physical_types, wcs=wcs)
+            if len(coords) == 1:
+                coords = coords[0]
+            return coords
+
 
 def _create_masked_array_for_rebinning(data, mask, operation_ignores_mask):
     m = None if (mask is None or mask is False or operation_ignores_mask) else mask

ndcube/tests/test_ndcube.py

@@ -1204,3 +1204,30 @@ def test_ndcube_quantity(ndcube_2d_ln_lt_units):
     cube = ndcube_2d_ln_lt_units
     expected = u.Quantity(cube.data, cube.unit)
     np.testing.assert_array_equal(cube.quantity, expected)
+
+
+def test_celestial(ndcube_3d_ln_lt_l_ec_time):
+    cube = ndcube_3d_ln_lt_l_ec_time
+    output = cube.celestial
+    expected = cube.axis_world_coords("lon")[0]
+    assert (output == expected).all()
+
+
+def test_time(ndcube_3d_ln_lt_l_ec_time):
+    cube = ndcube_3d_ln_lt_l_ec_time
+    output = cube.time
+    expected = cube.axis_world_coords("time", wcs=cube.extra_coords)[0]
+    assert (output == expected).all()
+
+
+def test_spectral(ndcube_3d_ln_lt_l_ec_time):
+    cube = ndcube_3d_ln_lt_l_ec_time
+    output = cube.spectral
+    expected = cube.axis_world_coords("em")[0]
+    assert (output == expected).all()
+
+
+def test_no_stokes(ndcube_3d_ln_lt_l_ec_time):
+    cube = ndcube_3d_ln_lt_l_ec_time
+    output = cube.stokes
+    assert output is None