Merge pull request #509 from astrofrog/spectral-coordinates

Added a SpectralCoordinates class for glue plugin

specviz/third_party/glue/tests/test_utils.py

@@ -1,17 +1,23 @@
 import pytest
 import numpy as np
+from numpy.testing import assert_allclose
 
 from astropy import units as u
 from astropy.wcs import WCS
 from astropy.tests.helper import assert_quantity_allclose
 
-pytest.importorskip("glue")
+pytest.importorskip("glue")  # noqa
 
 from glue.core import Data
 from glue.core.component import Component
 from glue.core.coordinates import WCSCoordinates
 
-from ..utils import glue_data_has_spectral_axis, glue_data_to_spectrum1d
+from ..utils import glue_data_has_spectral_axis, glue_data_to_spectrum1d, SpectralCoordinates
+
+
+def test_conversion_utils_notspec():
+    data = Data(label='not spectrum')
+    assert not glue_data_has_spectral_axis(data)
 
 
 def test_conversion_utils_1d():
@@ -24,14 +30,49 @@ def test_conversion_utils_1d():
     # Set up glue Coordinates object
     coords = WCSCoordinates(wcs=wcs)
 
-    data1 = Data(label='spectrum', coords=coords)
-    data1.add_component(Component(np.array([3.4, 2.3, -1.1, 0.3]), units='Jy'), 'x')
-
-    data2 = Data(label='not spectrum')
+    data = Data(label='spectrum', coords=coords)
+    data.add_component(Component(np.array([3.4, 2.3, -1.1, 0.3]), units='Jy'), 'x')
 
-    assert glue_data_has_spectral_axis(data1)
-    assert not glue_data_has_spectral_axis(data2)
+    assert glue_data_has_spectral_axis(data)
 
-    spec = glue_data_to_spectrum1d(data1, data1.id['x'])
+    spec = glue_data_to_spectrum1d(data, data.id['x'])
     assert_quantity_allclose(spec.spectral_axis, [1, 2, 3, 4] * u.m / u.s)
     assert_quantity_allclose(spec.flux, [3.4, 2.3, -1.1, 0.3] * u.Jy)
+
+
+def test_conversion_utils_3d():
+
+    # Set up simple spectral WCS
+    wcs = WCS(naxis=3)
+    wcs.wcs.ctype = ['RA---TAN', 'DEC--TAN', 'VELO-LSR']
+    wcs.wcs.set()
+
+    # Set up glue Coordinates object
+    coords = WCSCoordinates(wcs=wcs)
+
+    data = Data(label='spectral-cube', coords=coords)
+    data.add_component(Component(np.ones((3, 4, 5)), units='Jy'), 'x')
+
+    assert glue_data_has_spectral_axis(data)
+
+    spec = glue_data_to_spectrum1d(data, data.id['x'], statistic='sum')
+    assert_quantity_allclose(spec.spectral_axis, [1, 2, 3] * u.m / u.s)
+    assert_quantity_allclose(spec.flux, [20, 20, 20] * u.Jy)
+
+
+def test_conversion_utils_spectral_coordinates():
+
+    # Set up glue Coordinates object
+    coords = SpectralCoordinates([1, 4, 10] * u.micron)
+
+    data = Data(label='spectrum1d', coords=coords)
+    data.add_component(Component(np.array([3, 4, 5]), units='Jy'), 'x')
+
+    assert_allclose(data.coords.pixel2world([0, 0.5, 1, 1.5, 2]),
+                    [[1, 2.5, 4, 7, 10]])
+
+    assert glue_data_has_spectral_axis(data)
+
+    spec = glue_data_to_spectrum1d(data, data.id['x'])
+    assert_quantity_allclose(spec.spectral_axis, [1, 4, 10] * u.micron)
+    assert_quantity_allclose(spec.flux, [3, 4, 5] * u.Jy)

specviz/third_party/glue/utils.py

@@ -1,13 +1,41 @@
+import numpy as np
+
 from astropy import units as u
 from astropy.wcs import WCSSUB_SPECTRAL
 
 from specutils import Spectrum1D
 from glue.core.subset import Subset
-from glue.core.coordinates import WCSCoordinates
+from glue.core.coordinates import Coordinates, WCSCoordinates
 
 __all__ = ['glue_data_has_spectral_axis', 'glue_data_to_spectrum1d']
 
 
+class SpectralCoordinates(Coordinates):
+    """
+    This is a sub-class of Coordinates that is intended for 1-d spectral axes
+    given by a :class:`~astropy.units.Quantity` array.
+    """
+
+    def __init__(self, values):
+        self._index = np.arange(len(values))
+        self._values = values
+
+    @property
+    def spectral_axis(self):
+        return self._values
+
+    def world2pixel(self, *world):
+        return tuple(np.interp(world, self._values.value, self._index,
+                               left=np.nan, right=np.nan))
+
+    def pixel2world(self, *pixel):
+        return tuple(np.interp(pixel, self._index, self._values.value,
+                               left=np.nan, right=np.nan))
+
+    def dependent_axes(self, axis):
+        return (axis,)
+
+
 def glue_data_has_spectral_axis(data):
     """
     Check whether a glue Data object is a 1D spectrum.
@@ -21,6 +49,9 @@ def glue_data_has_spectral_axis(data):
     if isinstance(data, Subset):
         data = data.data
 
+    if isinstance(data.coords, SpectralCoordinates):
+        return True
+
     if not isinstance(data.coords, WCSCoordinates):
         return False
 
@@ -51,11 +82,24 @@ def glue_data_to_spectrum1d(data_or_subset, attribute, statistic='mean'):
         data = data_or_subset
         subset_state = None
 
-    # Find spectral axis
-    spec_axis = data.coords.wcs.naxis - 1 - data.coords.wcs.wcs.spec
+    if isinstance(data.coords, WCSCoordinates):
+
+        # Find spectral axis
+        spec_axis = data.coords.wcs.naxis - 1 - data.coords.wcs.wcs.spec
+
+        # Find non-spectral axes
+        axes = tuple(i for i in range(data.ndim) if i != spec_axis)
+
+        kwargs = {'wcs': data.coords.wcs.sub([WCSSUB_SPECTRAL])}
+
+    elif isinstance(data.coords, SpectralCoordinates):
+
+        kwargs = {'spectral_axis': data.coords.spectral_axis}
+
+    else:
+
+        raise TypeError('data.coords should be an instance of WCSCoordinates or SpectralCoordinates')
 
-    # Find non-spectral axes
-    axes = tuple(i for i in range(data.ndim) if i != spec_axis)
     component = data.get_component(attribute)
 
     # Collapse values to profile
@@ -71,6 +115,4 @@ def glue_data_to_spectrum1d(data_or_subset, attribute, statistic='mean'):
     else:
         values = values * u.Unit(component.units)
 
-    wcs_spec = data.coords.wcs.sub([WCSSUB_SPECTRAL])
-
-    return Spectrum1D(values, wcs=wcs_spec)
+    return Spectrum1D(values, **kwargs)