Added translators for NDDataArray and StdDevUncertainty

glue_astronomy/translators/__init__.py

@@ -1,4 +1,4 @@
-from . import ccddata  # noqa
+from . import nddata  # noqa
 from . import regions  # noqa
 from . import spectral_cube  # noqa
 from . import spectrum1d  # noqa

glue_astronomy/translators/nddata.py

@@ -0,0 +1,172 @@
+from astropy.wcs import WCS
+from astropy.wcs.wcsapi import BaseHighLevelWCS
+from astropy.nddata import CCDData, NDData, NDDataArray
+from astropy.nddata.nduncertainty import StdDevUncertainty
+from astropy import units as u
+
+from glue.config import data_translator
+from glue.core import Data, Subset
+from glue.core.coordinates import Coordinates
+
+
+def _get_attribute(attribute, data):
+    if isinstance(attribute, str):
+        attribute = data.id[attribute]
+    elif len(data.main_components) == 0:
+        raise ValueError('Data object has no attributes.')
+    elif attribute is None:
+        if len(data.main_components) == 1:
+            attribute = data.main_components[0]
+        else:
+            raise ValueError("Data object has more than one attribute, so "
+                             "you will need to specify which one to use as "
+                             "the flux for the spectrum using the "
+                             "attribute= keyword argument.")
+    return attribute
+
+
+def _get_value_and_mask(subset_state, data, values):
+    if subset_state is None:
+        mask = None
+    else:
+        mask = data.get_mask(subset_state=subset_state)
+        values = values.copy()
+        # Flip mask to match astropy.ndddata formalism
+        mask = ~mask
+    return values, mask
+
+
+def _get_data_and_subset_state(data_or_subset):
+    if isinstance(data_or_subset, Subset):
+        data = data_or_subset.data
+        subset_state = data_or_subset.subset_state
+    else:
+        data = data_or_subset
+        subset_state = None
+    return data, subset_state
+
+
+@data_translator(NDDataArray)
+class NDDataArrayHandler:
+
+    def to_data(self, obj):
+        data = Data(coords=obj.wcs)
+        data['data'] = obj.data
+        data.get_component('data').units = str(obj.unit)
+        data.meta.update(obj.meta)
+        return data
+
+    def to_object(self, data_or_subset, attribute=None):
+        """
+        Convert a glue Data object to a NDDataArray object.
+
+        Parameters
+        ----------
+        data_or_subset : `glue.core.data.Data` or `glue.core.subset.Subset`
+            The data to convert to a NDDataArray object
+        attribute : `glue.core.component_id.ComponentID`
+            The attribute to use for the NDDataArray data
+        """
+
+        data, subset_state = _get_data_and_subset_state(data_or_subset)
+
+        if isinstance(data.coords, WCS) or isinstance(data.coords, BaseHighLevelWCS):
+            wcs = data.coords
+        elif type(data.coords) is Coordinates or data.coords is None:
+            wcs = None
+        else:
+            raise TypeError('data.coords should be an instance of Coordinates or WCS')
+
+        attribute = _get_attribute(attribute, data)
+        component = data.get_component(attribute)
+        values = data.get_data(attribute)
+        values, mask = _get_value_and_mask(subset_state, data, values)
+
+        result = NDDataArray(
+            values, unit=component.units, mask=mask, wcs=wcs, meta=data.meta
+        )
+
+        return result
+
+
+@data_translator(CCDData)
+class CCDDataHandler(NDDataArrayHandler):
+
+    def to_object(self, data_or_subset, attribute=None):
+        """
+        Convert a glue Data object to a CCDData object.
+
+        Parameters
+        ----------
+        data_or_subset : `glue.core.data.Data` or `glue.core.subset.Subset`
+            The data to convert to a CCDData object
+        attribute : `glue.core.component_id.ComponentID`
+            The attribute to use for the CCDData data
+        """
+
+        data, subset_state = _get_data_and_subset_state(data_or_subset)
+
+        if isinstance(data.coords, WCS):
+            has_fitswcs = True
+            wcs = data.coords
+        elif isinstance(data.coords, BaseHighLevelWCS):
+            has_fitswcs = False
+            wcs = data.coords
+        elif type(data.coords) is Coordinates or data.coords is None:
+            has_fitswcs = True  # For backward compatibility
+            wcs = None
+        else:
+            raise TypeError('data.coords should be an instance of Coordinates or WCS')
+
+        attribute = _get_attribute(attribute, data)
+        component = data.get_component(attribute)
+
+        if data.ndim != 2:
+            raise ValueError("Only 2-dimensional datasets can be converted to CCDData")
+
+        values = data.get_data(attribute)
+        values, mask = _get_value_and_mask(subset_state, data, values)
+        values = values * u.Unit(component.units)
+
+        if has_fitswcs:
+            result = CCDData(values, mask=mask, wcs=wcs, meta=data.meta)
+        else:
+            # https://github.com/astropy/astropy/issues/11727
+            result = NDData(values, mask=mask, wcs=wcs, meta=data.meta)
+
+        return result
+
+
+@data_translator(StdDevUncertainty)
+class StdDevUncertaintyHandler:
+
+    def to_data(self, obj):
+        data = Data()
+        data['data'] = obj.array
+        data.get_component('data').units = str(obj.unit)
+        return data
+
+    def to_object(self, data_or_subset, attribute=None):
+        """
+        Convert a glue Data object to a StdDevUncertainty object.
+
+        Parameters
+        ----------
+        data_or_subset : `glue.core.data.Data` or `glue.core.subset.Subset`
+            The data to convert to a StdDevUncertainty object
+        attribute : `glue.core.component_id.ComponentID`
+            The attribute to use for the StdDevUncertainty data
+        """
+
+        if isinstance(data_or_subset, Subset):
+            data = data_or_subset.data
+        else:
+            data = data_or_subset
+
+        attribute = _get_attribute(attribute, data)
+        component = data.get_component(attribute)
+        values = data.get_data(attribute)
+
+        result = StdDevUncertainty(values, unit=component.units)
+
+        return result

glue_astronomy/translators/tests/test_ccddata.py → glue_astronomy/translators/tests/test_nddata.py

@@ -3,7 +3,8 @@
 from numpy.testing import assert_allclose, assert_equal
 
 from astropy import units as u
-from astropy.nddata import CCDData
+from astropy.nddata import CCDData, NDDataArray
+from astropy.nddata.nduncertainty import StdDevUncertainty
 from astropy.wcs import WCS
 
 from glue.core import Data, DataCollection
@@ -106,6 +107,35 @@ def test_to_ccddata_default_attribute():
                                  'keyword argument.')
 
 
+@pytest.mark.parametrize(
+    'cls, kwargs, data_attr',
+    [(NDDataArray, {'wcs': WCS_CELESTIAL}, 'data'),
+     (StdDevUncertainty, {}, 'array')])
+def test_from_nddata(cls, kwargs, data_attr):
+    spec = cls([[2, 3], [4, 5]] * u.Jy, **kwargs)
+
+    data_collection = DataCollection()
+
+    data_collection['image'] = spec
+
+    data = data_collection['image']
+
+    assert isinstance(data, Data)
+    assert len(data.main_components) == 1
+    assert data.main_components[0].label == 'data'
+    assert_allclose(data['data'], [[2, 3], [4, 5]])
+    component = data.get_component('data')
+    assert component.units == 'Jy'
+
+    # Check round-tripping
+    image_new = data.get_object(cls, attribute='data')
+    assert isinstance(image_new, cls)
+    if hasattr(image_new, 'wcs'):
+        assert image_new.wcs is WCS_CELESTIAL
+    assert_allclose(getattr(image_new, data_attr), [[2, 3], [4, 5]])
+    assert image_new.unit is u.Jy
+
+
 @pytest.mark.parametrize('with_wcs', (False, True))
 def test_from_ccddata(with_wcs):
 
@@ -130,7 +160,7 @@ def test_from_ccddata(with_wcs):
     assert component.units == 'Jy'
 
     # Check round-tripping
-    image_new = data.get_object(attribute='data')
+    image_new = data.get_object(CCDData, attribute='data')
     assert isinstance(image_new, CCDData)
     assert image_new.wcs is (WCS_CELESTIAL if with_wcs else None)
     assert_allclose(image_new.data, [[2, 3], [4, 5]])
@@ -144,22 +174,25 @@ def test_meta_round_trip():
     meta = {'BUNIT': 'Jy/beam',
             'some_variable': 10}
 
-    spec = CCDData([[2, 3], [4, 5]] * u.Jy, wcs=wcs, meta=meta)
-
+    flux = [[2, 3], [4, 5]] * u.Jy
+    kwargs = dict(wcs=wcs, meta=meta)
+    classes = [CCDData, NDDataArray]
+    image_names = ['image_ccd', 'image_ndd']
     data_collection = DataCollection()
 
-    data_collection['image'] = spec
+    for cls, image_name in zip(classes, image_names):
+        data_collection[image_name] = cls(flux, **kwargs)
 
-    data = data_collection['image']
+        data = data_collection[image_name]
 
-    assert isinstance(data, Data)
-    assert len(data.meta) == 2
-    assert data.meta['BUNIT'] == 'Jy/beam'
-    assert data.meta['some_variable'] == 10
+        assert isinstance(data, Data)
+        assert len(data.meta) == 2
+        assert data.meta['BUNIT'] == 'Jy/beam'
+        assert data.meta['some_variable'] == 10
 
-    # Check round-tripping
-    image_new = data.get_object(attribute='data')
-    assert isinstance(image_new, CCDData)
-    assert len(image_new.meta) == 2
-    assert image_new.meta['BUNIT'] == 'Jy/beam'
-    assert image_new.meta['some_variable'] == 10
+        # Check round-tripping
+        image_new = data.get_object(cls, attribute='data')
+        assert isinstance(image_new, cls)
+        assert len(image_new.meta) == 2
+        assert image_new.meta['BUNIT'] == 'Jy/beam'
+        assert image_new.meta['some_variable'] == 10

tox.ini

@@ -1,5 +1,5 @@
 [tox]
-envlist = py{36,37,38,39,310}-{test,docs}-{,dev}
+envlist = py{36,37,38,39,310}-{test,docs}{,-dev}
 requires = pip >= 18.0
            setuptools >= 30.3.0