Repository: glue-viz/glue-astronomy

Issue:
Spectrum1D translator cannot handle spectral_axis in 3D cube
I am not sure what is the proper fix. I don't want to spend time hacking it only to be told it is wrong. Please advise. Thanks!

```python
import numpy as np
from astropy import units as u
from specutils import Spectrum1D

from glue_astronomy.translators.spectrum1d import Specutils1DHandler

sh = Specutils1DHandler()
spec = Spectrum1D(spectral_axis=np.arange(10) * u.pix, flux=np.ones((2, 2, 10)) * u.Jy)
obj = sh.to_data(spec)
```

```
.../glue_astronomy/translators/spectrum1d.py in to_data(self, obj)
    127         # Spectrum1D does it automatically on initialization.
    128         if len(obj.flux.shape) == 3:
--> 129             data = Data(coords=obj.wcs.swapaxes(-1, 0))
    130             data['flux'] = np.swapaxes(obj.flux, -1, 0)
    131             data.get_component('flux').units = str(obj.unit)

AttributeError: 'SpectralGWCS' object has no attribute 'swapaxes'
```

This is because when I pass `spectral_axis` into `Spectrum1D`, it creates a GWCS using https://github.com/astropy/specutils/blob/08de408d18d709a857b67e10ca2752d548eb3100/specutils/utils/wcs_utils.py#L185 and GWCS does not seem to understand `swapaxes`.

cc @astrofrog @rosteen 

This blocks spacetelescope/jdaviz#1040

Also see spacetelescope/gwcs#397

Discussion:
codecov[bot]: # [Codecov](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz) Report
> Merging [#68](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz) (5093e0d) into [main](https://codecov.io/gh/glue-viz/glue-astronomy/commit/7f3229d01bdc2bff64fdb1a83738447cf696c59d?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz) (7f3229d) will **decrease** coverage by `3.60%`.
> The diff coverage is `45.45%`.

```diff
@@            Coverage Diff             @@
##             main      #68      +/-   ##
==========================================
- Coverage   95.80%   92.20%   -3.61%     
==========================================
  Files          15       15              
  Lines        1145     1142       -3     
==========================================
- Hits         1097     1053      -44     
- Misses         48       89      +41     
```


| [Impacted Files](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz) | Coverage Δ | |
|---|---|---|
| [glue\_astronomy/translators/spectrum1d.py](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz#diff-Z2x1ZV9hc3Ryb25vbXkvdHJhbnNsYXRvcnMvc3BlY3RydW0xZC5weQ==) | `80.39% <45.45%> (-12.18%)` | :arrow_down: |
| [...lue\_astronomy/translators/tests/test\_spectrum1d.py](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz#diff-Z2x1ZV9hc3Ryb25vbXkvdHJhbnNsYXRvcnMvdGVzdHMvdGVzdF9zcGVjdHJ1bTFkLnB5) | `81.70% <0.00%> (-18.30%)` | :arrow_down: |
| [...onomy/io/spectral\_cube/tests/test\_spectral\_cube.py](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz#diff-Z2x1ZV9hc3Ryb25vbXkvaW8vc3BlY3RyYWxfY3ViZS90ZXN0cy90ZXN0X3NwZWN0cmFsX2N1YmUucHk=) | `100.00% <0.00%> (+14.28%)` | :arrow_up: |
| [glue\_astronomy/\_\_init\_\_.py](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz#diff-Z2x1ZV9hc3Ryb25vbXkvX19pbml0X18ucHk=) | `100.00% <0.00%> (+28.57%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz). Last update [7f3229d...5093e0d](https://codecov.io/gh/glue-viz/glue-astronomy/pull/68?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=glue-viz).

rosteen: I fixed a few remaining issues with the world<->pixel conversion and added test coverage for the 3D case, I think this is ready for final review now. I still have a decent bit of work on the related Jdaviz PR.

Patch:
diff --git a/CHANGES.rst b/CHANGES.rst
index e316e24..97d9b2c 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -1,3 +1,10 @@
+0.4.0 (unreleased)
+------------------
+
+- Updated ``Spectrum1D`` translator to generate dummy WCS when needed for any
+  dimensionality, and to preserve specutils axis order when translating
+  to Glue ``Data``. [#68]
+
 0.3.3 (2022-03-22)
 ------------------
 
diff --git a/glue_astronomy/translators/spectrum1d.py b/glue_astronomy/translators/spectrum1d.py
index 65875c1..9caeabb 100644
--- a/glue_astronomy/translators/spectrum1d.py
+++ b/glue_astronomy/translators/spectrum1d.py
@@ -38,35 +38,39 @@ class PaddedSpectrumWCS(BaseWCSWrapper, HighLevelWCSMixin):
     # datasets while glue always needs the dimensionality to match,
     # so this class pads the WCS so that it is n-dimensional.
 
-    # NOTE: for now this only handles padding the WCS into 2D WCS. Rather than
-    # generalize this we can just remove this class and use CompoundLowLevelWCS
-    # from NDCube once it is in a released version.
+    # NOTE: This class could be updated to use CompoundLowLevelWCS from NDCube.
 
-    def __init__(self, wcs):
+    def __init__(self, wcs, ndim):
         self.spectral_wcs = wcs
+        self.flux_ndim = ndim
+
+        if self.flux_ndim == 2:
+            self.spatial_keys = ['spatial']
+        else:
+            self.spatial_keys = [f"spatial{i}" for i in range(0, self.flux_ndim-1)]
 
     @property
     def pixel_n_dim(self):
-        return 2
+        return self.flux_ndim
 
     @property
     def world_n_dim(self):
-        return 2
+        return self.flux_ndim
 
     @property
     def world_axis_physical_types(self):
-        return [self.spectral_wcs.world_axis_physical_types[0], None]
+        return [self.spectral_wcs.world_axis_physical_types[0], *[None]*(self.flux_ndim-1)]
 
     @property
     def world_axis_units(self):
-        return (self.spectral_wcs.world_axis_units[0], None)
+        return (self.spectral_wcs.world_axis_units[0], *[None]*(self.flux_ndim-1))
 
     def pixel_to_world_values(self, *pixel_arrays):
         # The ravel and reshape are needed because of
         # https://github.com/astropy/astropy/issues/12154
         px = np.array(pixel_arrays[0])
         world_arrays = [self.spectral_wcs.pixel_to_world_values(px.ravel()).reshape(px.shape),
-                        pixel_arrays[1]]
+                        *pixel_arrays[1:]]
         return tuple(world_arrays)
 
     def world_to_pixel_values(self, *world_arrays):
@@ -74,23 +78,22 @@ class PaddedSpectrumWCS(BaseWCSWrapper, HighLevelWCSMixin):
         # https://github.com/astropy/astropy/issues/12154
         wx = np.array(world_arrays[0])
         pixel_arrays = [self.spectral_wcs.world_to_pixel_values(wx.ravel()).reshape(wx.shape),
-                        world_arrays[1]]
+                        *world_arrays[1:]]
         return tuple(pixel_arrays)
 
     @property
     def world_axis_object_components(self):
-        return [
-            self.spectral_wcs.world_axis_object_components[0],
-            ('spatial', 'value', 'value')
-        ]
+        return [self.spectral_wcs.world_axis_object_components[0],
+                *[(key, 'value', 'value') for key in self.spatial_keys]]
 
     @property
     def world_axis_object_classes(self):
         spectral_key = self.spectral_wcs.world_axis_object_components[0][0]
-        return {
-            spectral_key: self.spectral_wcs.world_axis_object_classes[spectral_key],
-            'spatial': (u.Quantity, (), {'unit': u.pixel})
-        }
+        obj_classes = {spectral_key: self.spectral_wcs.world_axis_object_classes[spectral_key]}
+        for key in self.spatial_keys:
+            obj_classes[key] = (u.Quantity, (), {'unit': u.pixel})
+
+        return obj_classes
 
     @property
     def pixel_shape(self):
@@ -102,15 +105,20 @@ class PaddedSpectrumWCS(BaseWCSWrapper, HighLevelWCSMixin):
 
     @property
     def pixel_axis_names(self):
-        return tuple([self.spectral_wcs.pixel_axis_names[0], 'spatial'])
+        return tuple([self.spectral_wcs.pixel_axis_names[0], *self.spatial_keys])
 
     @property
     def world_axis_names(self):
-        return (UCD_TO_SPECTRAL_NAME[self.spectral_wcs.world_axis_physical_types[0]], 'Offset')
+        if self.flux_ndim == 2:
+            names = ['Offset']
+        else:
+            names = [f"Offset{i}" for i in range(0, self.flux_ndim-1)]
+
+        return (UCD_TO_SPECTRAL_NAME[self.spectral_wcs.world_axis_physical_types[0]], *names)
 
     @property
     def axis_correlation_matrix(self):
-        return np.identity(2).astype('bool')
+        return np.identity(self.flux_ndim).astype('bool')
 
     @property
     def serialized_classes(self):
@@ -125,15 +133,17 @@ class Specutils1DHandler:
         # Glue expects spectral axis first for cubes (opposite of specutils).
         # Swap the spectral axis to first here. to_object doesn't need this because
         # Spectrum1D does it automatically on initialization.
-        if len(obj.flux.shape) == 3:
-            data = Data(coords=obj.wcs.swapaxes(-1, 0))
-            data['flux'] = np.swapaxes(obj.flux, -1, 0)
-            data.get_component('flux').units = str(obj.unit)
+        if obj.flux.ndim > 1:
+            if obj.wcs.world_n_dim == 1:
+                data = Data(coords=PaddedSpectrumWCS(obj.wcs, obj.flux.ndim))
+                data['flux'] = obj.flux
+            else:
+                data = Data(coords=obj.wcs)
+                data['flux'] = obj.flux
+                data.get_component('flux').units = str(obj.unit)
         else:
-            if obj.flux.ndim == 1 and obj.wcs.world_n_dim == 1 and isinstance(obj.wcs, GWCS):
+            if obj.wcs.world_n_dim == 1 and isinstance(obj.wcs, GWCS):
                 data = Data(coords=SpectralCoordinates(obj.spectral_axis))
-            elif obj.flux.ndim == 2 and obj.wcs.world_n_dim == 1:
-                data = Data(coords=PaddedSpectrumWCS(obj.wcs))
             else:
                 data = Data(coords=obj.wcs)
             data['flux'] = obj.flux
@@ -141,19 +151,13 @@ class Specutils1DHandler:
 
         # Include uncertainties if they exist
         if obj.uncertainty is not None:
-            if len(obj.flux.shape) == 3:
-                data['uncertainty'] = np.swapaxes(obj.uncertainty.quantity, -1, 0)
-            else:
-                data['uncertainty'] = obj.uncertainty.quantity
+            data['uncertainty'] = obj.uncertainty.quantity
             data.get_component('uncertainty').units = str(obj.uncertainty.unit)
             data.meta.update({'uncertainty_type': obj.uncertainty.uncertainty_type})
 
         # Include mask if it exists
         if obj.mask is not None:
-            if len(obj.flux.shape) == 3:
-                data['mask'] = np.swapaxes(obj.mask, -1, 0)
-            else:
-                data['mask'] = obj.mask
+            data['mask'] = obj.mask
 
         data.meta.update(obj.meta)
 
@@ -194,7 +198,7 @@ class Specutils1DHandler:
 
             if isinstance(data.coords, PaddedSpectrumWCS):
                 spec_axis = 0
-                axes = tuple(range(1, data.ndim))
+                axes = tuple(range(0, data.ndim-1))
                 kwargs = {'wcs': data.coords.spectral_wcs}
             elif isinstance(data.coords, WCS):
 
@@ -258,7 +262,7 @@ class Specutils1DHandler:
                     values = data.compute_statistic(statistic, attribute, axis=axes,
                                                     subset_state=subset_state)
                     if mask is not None:
-                        collapse_axes = tuple([x for x in range(1, data.ndim)])
+                        collapse_axes = tuple([x for x in range(0, data.ndim-1)])
                         mask = np.all(mask, collapse_axes)
                 else:
                     values = data.get_data(attribute)


Test patch:
diff --git a/glue_astronomy/translators/tests/test_spectrum1d.py b/glue_astronomy/translators/tests/test_spectrum1d.py
index d71f81a..ec8c7c7 100644
--- a/glue_astronomy/translators/tests/test_spectrum1d.py
+++ b/glue_astronomy/translators/tests/test_spectrum1d.py
@@ -166,13 +166,13 @@ def test_from_spectrum1d(mode):
     assert isinstance(data, Data)
     assert len(data.main_components) == 3
     assert data.main_components[0].label == 'flux'
-    assert_allclose(data['flux'], flux.value)
+    assert_allclose(data['flux'], spec.flux.value)
     component = data.get_component('flux')
     assert component.units == 'Jy'
 
     # Check uncertainty parsing within glue data object
     assert data.main_components[1].label == 'uncertainty'
-    assert_allclose(data['uncertainty'], uncertainty.array)
+    assert_allclose(data['uncertainty'], spec.uncertainty.array)
     component = data.get_component('uncertainty')
     assert component.units == 'Jy2'
 
@@ -193,8 +193,6 @@ def test_from_spectrum1d(mode):
     if mode == 'wcs3d':
         assert_quantity_allclose(spec_new.flux, np.ones((5, 4, 4))*u.Unit('Jy'))
         assert spec_new.uncertainty is not None
-        print(spec_new.uncertainty)
-        print(uncertainty)
         assert_quantity_allclose(spec_new.uncertainty.quantity,
                                  np.ones((5, 4, 4))*0.01*u.Jy**2)
     else:
@@ -203,7 +201,8 @@ def test_from_spectrum1d(mode):
         assert_quantity_allclose(spec_new.uncertainty.quantity, [0.1, 0.1, 0.1, 0.1] * u.Jy**2)
 
 
-def test_spectrum1d_2d_data():
+@pytest.mark.parametrize('spec_ndim', (2, 3))
+def test_spectrum1d_2d_data(spec_ndim):
 
     # This test makes sure that 2D spectra represented as Spectrum1D round-trip
     # Note that Spectrum1D will typically have a 1D spectral WCS even if the
@@ -217,11 +216,14 @@ def test_spectrum1d_2d_data():
     wcs.wcs.cdelt = [10]
     wcs.wcs.set()
 
-    flux = np.ones((3, 2)) * u.Unit('Jy')
+    if spec_ndim == 2:
+        flux = np.ones((3, 2)) * u.Unit('Jy')
+    elif spec_ndim == 3:
+        flux = np.ones((3, 3, 2)) * u.Unit('Jy')
 
     spec = Spectrum1D(flux, wcs=wcs, meta={'instrument': 'spamcam'})
 
-    assert spec.data.ndim == 2
+    assert spec.data.ndim == spec_ndim
     assert spec.wcs.naxis == 1
 
     data_collection = DataCollection()
@@ -235,27 +237,61 @@ def test_spectrum1d_2d_data():
     assert data.main_components[0].label == 'flux'
     assert_allclose(data['flux'], flux.value)
 
-    assert data.coords.pixel_n_dim == 2
-    assert data.coords.world_n_dim == 2
-    assert len(data.pixel_component_ids) == 2
-    assert len(data.world_component_ids) == 2
-
-    assert data.coordinate_components[0].label == 'Pixel Axis 0 [y]'
-    assert data.coordinate_components[1].label == 'Pixel Axis 1 [x]'
-    assert data.coordinate_components[2].label == 'Offset'
-    assert data.coordinate_components[3].label == 'Frequency'
-
-    assert_equal(data['Offset'], [[0, 0], [1, 1], [2, 2]])
-    assert_equal(data['Frequency'], [[10, 20], [10, 20], [10, 20]])
-
-    s, o = data.coords.pixel_to_world(1, 2)
-    assert isinstance(s, SpectralCoord)
-
-    # Check round-tripping of coordinates
-    with pytest.warns(AstropyUserWarning, match='No observer defined on WCS'):
-        px, py = data.coords.world_to_pixel(s, o)
-    assert_allclose(px, 1)
-    assert_allclose(py, 2)
+    assert data.coords.pixel_n_dim == spec_ndim
+    assert data.coords.world_n_dim == spec_ndim
+    assert len(data.pixel_component_ids) == spec_ndim
+    assert len(data.world_component_ids) == spec_ndim
+
+    if spec_ndim == 2:
+        assert data.coordinate_components[0].label == 'Pixel Axis 0 [y]'
+        assert data.coordinate_components[1].label == 'Pixel Axis 1 [x]'
+        assert data.coordinate_components[2].label == 'Offset'
+        assert data.coordinate_components[3].label == 'Frequency'
+        assert data.coords.pixel_axis_names == ('', 'spatial')
+
+        assert_equal(data['Offset'], [[0, 0], [1, 1], [2, 2]])
+        assert_equal(data['Frequency'], [[10, 20], [10, 20], [10, 20]])
+
+        s, o = data.coords.pixel_to_world(1, 2)
+        assert isinstance(s, SpectralCoord)
+
+        # Check round-tripping of coordinates
+        with pytest.warns(AstropyUserWarning, match='No observer defined on WCS'):
+            px, py = data.coords.world_to_pixel(s, o)
+        assert_allclose(px, 1)
+        assert_allclose(py, 2)
+
+    elif spec_ndim == 3:
+        assert data.coordinate_components[0].label == 'Pixel Axis 0 [z]'
+        assert data.coordinate_components[1].label == 'Pixel Axis 1 [y]'
+        assert data.coordinate_components[2].label == 'Pixel Axis 2 [x]'
+        assert data.coordinate_components[3].label == 'Offset1'
+        assert data.coordinate_components[4].label == 'Offset0'
+        assert data.coordinate_components[5].label == 'Frequency'
+        assert data.coords.pixel_axis_names == ('', 'spatial0', 'spatial1')
+
+        assert_equal(data['Offset1'], [[[0, 0], [0, 0], [0, 0]],
+                                       [[1, 1], [1, 1], [1, 1]],
+                                       [[2, 2], [2, 2], [2, 2]]])
+        assert_equal(data['Offset0'], [[[0, 0], [1, 1], [2, 2]],
+                                       [[0, 0], [1, 1], [2, 2]],
+                                       [[0, 0], [1, 1], [2, 2]]])
+        assert_equal(data['Frequency'], [[[10, 20], [10, 20], [10, 20]],
+                                         [[10, 20], [10, 20], [10, 20]],
+                                         [[10, 20], [10, 20], [10, 20]]])
+
+        s, o1, o2 = data.coords.pixel_to_world(1, 2, 0)
+        assert isinstance(s, SpectralCoord)
+
+        # Check round-tripping of coordinates
+        with pytest.warns(AstropyUserWarning, match='No observer defined on WCS'):
+            px, py, pz = data.coords.world_to_pixel(s, o1, o2)
+        assert_allclose(px, 1)
+        assert_allclose(py, 2)
+        assert_allclose(pz, 0)
+
+    assert data.coords.world_axis_units == ('Hz', *(None,)*(spec_ndim-1))
+    assert data.coords.world_axis_physical_types == ['em.freq', *(None,)*(spec_ndim-1)]
 
     # Check round-tripping of translation
     spec_new = data.get_object(statistic=None)