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 comparison 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):
+ '''
+ Convert a 2D plane of pixel positions 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,76 @@ def argmin(self, axis=None, how='auto', **kwargs):
  reduce=False, projection=False,
  how=how, axis=axis, **kwargs)
 
+ def _argmaxmin_world(self, axis, method, **kwargs):
+ '''
+ Return the spatial or spectral index of the maximum or minimum value.
+ Use `argmax_world` and `argmin_world` directly.
+ '''
+
+ operation_name = '{}_world'.format(method)
+
+ if wcs_utils.is_pixel_axis_to_wcs_correlated(self.wcs, axis):
+ raise WCSCelestialError("{} requires the celestial axes"
+ " to be aligned along image axes."
+ .format(operation_name))
+
+ if method == 'argmin':
+ arg_pixel_plane = self.argmin(axis=axis, **kwargs)
+ elif method == 'argmax':
+ arg_pixel_plane = self.argmax(axis=axis, **kwargs)
+ else:
+ raise ValueError("`method` must be 'argmin' or 'argmax'")
+
+ # Convert to WCS coordinates.
+ out = cube_utils.world_take_along_axis(self, arg_pixel_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 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`.
+ '''
+
+ return self._argmaxmin_world(axis, 'argmax', **kwargs)
+
+ @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`.
+ '''
+
+ return self._argmaxmin_world(axis, 'argmin', **kwargs)
+
  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