Add operations to return WCS coordinates from argmin/argmax

CHANGES.rst

@@ -1,5 +1,8 @@
 0.6 (unreleased)
 ----------------
+- Add ``argmax_world`` and ``argmin_world`` to return the argmin/max position
+  in WCS coordinates. This is ONLY defined for independent WCS axes
+  (e.g., spectral) #680
 - Bugfix: subcube producing spatial offsets for large images #666
 - Switch to using standalone casa-formats-io package for reading in CASA
   .image files (this was split out from spectral-cube). #684

docs/moments.rst

@@ -22,25 +22,25 @@ axis::
           moment is computed as the variance. For the actual formulas used for
           the moments, please see :class:`~spectral_cube.SpectralCube.moment`.
           For linewidth maps, see the `Linewidth maps`_ section.
-          
+
 You may also want to convert the unit of the datacube into a velocity one before
-you can obtain a genuine velocity map via a 1st moment map. So first it will be necessary to 
+you can obtain a genuine velocity map via a 1st moment map. So first it will be necessary to
 apply the :class:`~spectral_cube.SpectralCube.with_spectral_unit` method from this package with the proper attribute settings::
 
     >>> nii_cube = cube.with_spectral_unit(u.km/u.s,
                                            velocity_convention='optical',
                                            rest_value=6584*u.AA)  # doctest: +SKIP
 
-Note that the ``rest_value`` in the above code refers to the wavelength of the targeted line 
-in the 1D spectrum corresponding to the 3rd dimension. Also, since not all velocity values are relevant, 
-next we will use the :class:`~spectral_cube.SpectralCube.spectral_slab` method to slice out the chunk of 
+Note that the ``rest_value`` in the above code refers to the wavelength of the targeted line
+in the 1D spectrum corresponding to the 3rd dimension. Also, since not all velocity values are relevant,
+next we will use the :class:`~spectral_cube.SpectralCube.spectral_slab` method to slice out the chunk of
 the cube that actually contains the line::
 
     >>> nii_cube = cube.with_spectral_unit(u.km/u.s,
                                            velocity_convention='optical',
                                            rest_value=6584*u.AA)  # doctest: +SKIP
     >>> nii_subcube = nii_cube.spectral_slab(-60*u.km/u.s,-20*u.km/u.s)  # doctest: +SKIP
-    
+
 Finally, we can now generate the 1st moment map containing the expected velocity structure::
 
     >>> moment_1 = nii_subcube.moment(order=1)  # doctest: +SKIP
@@ -73,7 +73,7 @@ will create the quicklook grayscale image and save it to a png all in one go.
 
 Moment map equations
 ^^^^^^^^^^^^^^^^^^^^
- 
+
 The moments are defined below, using :math:`v` for the spectral (velocity,
 frequency, wavelength, or energy) axis and :math:`I_v` as the intensity,
 or otherwise measured flux, value in a given spectral channel.
@@ -85,7 +85,7 @@ The equation for the 0th moment is:
 The equation for the 1st moment is:
 
 .. math:: M_1 = \frac{\int v I_v  dv}{\int I_v dv} = \frac{\int v I_v dv}{M_0}
-   
+
 Higher-order moments (:math:`N\geq2`) are defined as:
 
 .. math:: M_N = \frac{\int I_v (v - M_1)^N dv}{M_0}
@@ -112,14 +112,51 @@ with either of these two commands::
     >>> fwhm_map = cube.linewidth_fwhm()  # doctest: +SKIP
 
 ``~spectral_cube.SpectralCube.linewidth_sigma`` computes a sigma linewidth map
-along the spectral axis, where sigma is the width of a Gaussian, while 
+along the spectral axis, where sigma is the width of a Gaussian, while
 ``~spectral_cube.SpectralCube.linewidth_fwhm`` computes a FWHM
 linewidth map along the same spectral axis.
 
 The linewidth maps are related to the second moment by
 
 .. math:: \sigma = \sqrt{M_2} \\
-          FWHM = \sigma \sqrt{8 ln{2}} 
+          FWHM = \sigma \sqrt{8 ln{2}}
 
 These functions return :class:`~spectral_cube.lower_dimensional_structures.Projection` instances as for the
 `Moment maps`_.
+
+Additional 2D maps
+------------------
+
+Other common 2D views of a spectral cube include the maximum/minimum along a dimension and
+the location of the maximum or the minimum along that dimension.
+
+To produce a 2D map of the maximum along a cube dimension, use `~spectral_cube.SpectralCube.max`
+
+    >>> max_map = cube.max(axis=0)  # doctest: +SKIP
+
+Along the spectral axis, this will return the peak intensity map.
+
+Similarly we can use `~spectral_cube.SpectralCube.min` to make a 2D minimum map:
+
+    >>> min_map = cube.min(axis=0)  # doctest: +SKIP
+
+The `~spectral_cube.SpectralCube.argmax` and `~spectral_cube.SpectralCube.argmin` operations will
+return the pixel positions of the max/min along that axis:
+
+    >>> argmax_map = cube.argmax(axis=0)  # doctest: +SKIP
+    >>> argmin_map = cube.argmin(axis=0)  # doctest: +SKIP
+
+These maps are useful for identifying where signal is located within the spectral cube, however,
+it is more useful to return the WCS values of those pixels for comparisons with other data sets
+or for modeling. The `~spectral_cube.SpectralCube.argmax_world` and
+`~spectral_cube.SpectralCube.argmin_world` should be used in this case:
+
+    >>> world_argmax_map = cube.argmax_world(axis=0)  # doctest: +SKIP
+    >>> world_argmin_map = cube.argmin_world(axis=0)  # doctest: +SKIP
+
+Along the spectral axis, `~spectral_cube.SpectralCube.argmax_world` creates the often used
+"velocity at peak intensity," which may also be called the "peak velocity."
+
+.. note:: `cube.argmax_world` and `cube.argmin_world` are currently only defined along the spectral axis,
+          as the example above shows. This is because `argmax_world` and `argmin_world` operate along the
+          pixel axes, but they are not independent in WCS coordinates due to the curvature of the sky.

spectral_cube/cube_utils.py

@@ -12,9 +12,10 @@
 import numpy as np
 from astropy.wcs import (WCSSUB_SPECTRAL, WCSSUB_LONGITUDE, WCSSUB_LATITUDE)
 from . import wcs_utils
-from .utils import FITSWarning, AstropyUserWarning
+from .utils import FITSWarning, AstropyUserWarning, WCSCelestialError
 from astropy import log
 from astropy.io import fits
+from astropy.wcs.utils import is_proj_plane_distorted
 from astropy.io.fits import BinTableHDU, Column
 from astropy import units as u
 import itertools
@@ -486,3 +487,51 @@ def convert_bunit(bunit):
             unit = None
 
     return unit
+
+
+def world_take_along_axis(cube, position_plane, axis):
+    '''
+    Apply a 2D plane of pixel positions along collapsed along a third axis
+    to the equivalent WCS coordinates. For example, this will convert `argmax`
+    along the spectral axis to the equivalent spectral value (e.g., velocity at
+    peak intensity).
+
+    Parameters
+    ----------
+    cube : SpectralCube
+        A spectral cube.
+    position_plane : 2D numpy.ndarray
+        2D array of pixel positions along `axis`. For example, `position_plane` can
+        be the output of `argmax` or `argmin` along an axis.
+    axis : int
+        The axis that `position_plane` is collapsed along.
+
+    Returns
+    -------
+    out : astropy.units.Quantity
+        2D array of WCS coordinates.
+    '''
+
+    if wcs_utils.is_pixel_axis_to_wcs_correlated(cube.wcs, axis):
+        raise WCSCelestialError("world_take_along_axis requires the celestial axes"
+                                " to be aligned along image axes.")
+
+    # Get 1D slice along that axis.
+    world_slice = [0, 0]
+    world_slice.insert(axis, slice(None))
+
+    world_coords = cube.world[tuple(world_slice)][axis]
+
+    world_newaxis = [np.newaxis] * 2
+    world_newaxis.insert(axis, slice(None))
+    world_newaxis = tuple(world_newaxis)
+
+    plane_newaxis = [slice(None), slice(None)]
+    plane_newaxis.insert(axis, np.newaxis)
+    plane_newaxis = tuple(plane_newaxis)
+
+    out = np.take_along_axis(world_coords[world_newaxis],
+                             position_plane[plane_newaxis], axis=axis)
+    out = out.squeeze()
+
+    return out

spectral_cube/spectral_cube.py

@@ -51,7 +51,8 @@
                     UnsupportedIterationStrategyWarning, WCSMismatchWarning,
                     NotImplementedWarning, SliceWarning, SmoothingWarning,
                     StokesWarning, ExperimentalImplementationWarning,
-                    BeamAverageWarning, NonFiniteBeamsWarning, BeamWarning)
+                    BeamAverageWarning, NonFiniteBeamsWarning, BeamWarning,
+                    WCSCelestialError)
 from .spectral_axis import (determine_vconv_from_ctype, get_rest_value_from_wcs,
                             doppler_beta, doppler_gamma, doppler_z)
 from .io.core import SpectralCubeRead, SpectralCubeWrite
@@ -812,6 +813,80 @@ def argmin(self, axis=None, how='auto', **kwargs):
                                          reduce=False, projection=False,
                                          how=how, axis=axis, **kwargs)
 
+    @warn_slow
+    def argmax_world(self, axis, **kwargs):
+        '''
+        Return the spatial or spectral index of the maximum value
+        along a line of sight.
+
+        Parameters
+        ----------
+        axis : int
+            The axis to return the peak location along. e.g., `axis=0`
+            will return the value of the spectral axis at the peak value.
+        kwargs : dict
+            Passed to `~SpectralCube.argmax`.
+        '''
+
+        if wcs_utils.is_pixel_axis_to_wcs_correlated(self.wcs, axis):
+            raise WCSCelestialError("argmax_world requires the celestial axes"
+                                    " to be aligned along image axes.")
+
+        argmax_plane = self.argmax(axis=axis, **kwargs)
+
+        # Convert to WCS coordinates.
+        out = cube_utils.world_take_along_axis(self, argmax_plane, axis)
+
+        # Compute whether the mask has any valid data along `axis`
+        collapsed_mask = self.mask.include().any(axis=axis)
+        out[~collapsed_mask] = np.NaN
+
+        # Return a Projection.
+        new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
+
+        meta = {'collapse_axis': axis}
+        meta.update(self._meta)
+
+        return Projection(out, copy=False, wcs=new_wcs, meta=meta,
+                          unit=out.unit, header=self._nowcs_header)
+
+    @warn_slow
+    def argmin_world(self, axis, **kwargs):
+        '''
+        Return the spatial or spectral index of the minimum value
+        along a line of sight.
+
+        Parameters
+        ----------
+        axis : int
+            The axis to return the peak location along. e.g., `axis=0`
+            will return the value of the spectral axis at the peak value.
+        kwargs : dict
+            Passed to `~SpectralCube.argmin`.
+        '''
+
+        if wcs_utils.is_pixel_axis_to_wcs_correlated(self.wcs, axis):
+            raise WCSCelestialError("argmin_world requires the celestial axes"
+                                    " to be aligned along image axes.")
+
+        argmin_plane = self.argmin(axis=axis, **kwargs)
+
+        # Convert to WCS coordinates.
+        out = cube_utils.world_take_along_axis(self, argmin_plane, axis)
+
+        # Compute whether the mask has any valid data along `axis`
+        collapsed_mask = self.mask.include().any(axis=axis)
+        out[~collapsed_mask] = np.NaN
+
+        # Return a Projection.
+        new_wcs = wcs_utils.drop_axis(self._wcs, np2wcs[axis])
+
+        meta = {'collapse_axis': axis}
+        meta.update(self._meta)
+
+        return Projection(out, copy=False, wcs=new_wcs, meta=meta,
+                          unit=out.unit, header=self._nowcs_header)
+
     def chunked(self, chunksize=1000):
         """
         Not Implemented.

spectral_cube/tests/test_spectral_cube.py

@@ -719,6 +719,45 @@ def test_transpose(self, method, data_adv, data_vad, use_dask):
                                 getattr(c2, method)(axis=axis, progressbar=True))
         self.c = self.d = None
 
+    @pytest.mark.parametrize('method', ('argmax_world', 'argmin_world'))
+    def test_transpose_arg_world(self, method, data_adv, data_vad, use_dask):
+        c1, d1 = cube_and_raw(data_adv, use_dask=use_dask)
+        c2, d2 = cube_and_raw(data_vad, use_dask=use_dask)
+
+        # The spectral axis should work in all of these test cases.
+        axis = 0
+        assert_allclose(getattr(c1, method)(axis=axis),
+                        getattr(c2, method)(axis=axis))
+        if not use_dask:
+            # check that all these accept progressbar kwargs
+            assert_allclose(getattr(c1, method)(axis=axis, progressbar=True),
+                            getattr(c2, method)(axis=axis, progressbar=True))
+
+        # But the spatial axes should fail since the pixel axes are correlated to
+        # the WCS celestial axes. Currently this will happen for ALL celestial axes.
+        for axis in [1, 2]:
+
+            with pytest.raises(utils.WCSCelestialError,
+                               match=re.escape(f"{method} requires the celestial axes")):
+
+                assert_allclose(getattr(c1, method)(axis=axis),
+                                getattr(c2, method)(axis=axis))
+
+        self.c = self.d = None
+
+    @pytest.mark.parametrize('method', ('argmax_world', 'argmin_world'))
+    def test_arg_world(self, method, data_adv, use_dask):
+        c1, d1 = cube_and_raw(data_adv, use_dask=use_dask)
+
+        # Pixel operation is same name with "_world" removed.
+        arg0_pixel = getattr(c1, method.split("_")[0])(axis=0)
+
+        arg0_world = np.take_along_axis(c1.spectral_axis[:, np.newaxis, np.newaxis],
+                                        arg0_pixel[np.newaxis, :, :], axis=0).squeeze()
+
+        assert_allclose(getattr(c1, method)(axis=0), arg0_world)
+
+        self.c = self.d = None
 
 class TestSlab(BaseTest):
 

spectral_cube/wcs_utils.py

@@ -483,6 +483,30 @@ def diagonal_wcs_to_cdelt(mywcs):
     return mywcs
 
 
+def is_pixel_axis_to_wcs_correlated(mywcs, axis):
+    """
+    Check if the chosen pixel axis correlates to other WCS axes. This tests
+    whether the pixel axis is correlated only to 1 WCS axis and can be
+    considered independent of the others.
+    """
+
+    axis_corr_matrix = mywcs.axis_correlation_matrix
+
+    # Map from numpy axis to WCS axis
+    wcs_axis = mywcs.world_n_dim - (axis + 1)
+
+    # Grab the row along the given spatial axis. This slice is along the WCS axes
+    wcs_axis_correlations = axis_corr_matrix[:, wcs_axis]
+
+    # The image axis should always be correlated to at least 1 WCS axis.
+    # i.e., the diagonal term is one in the matrix. Correlations with other axes will give
+    # a sum > 1
+    if wcs_axis_correlations.sum() > 1:
+        return True
+
+    return False
+
+
 def find_spatial_pixel_index(cube, xlo, xhi, ylo, yhi):
     '''
     Given low and high cuts, return the pixel coordinates for a rectangular