Adding support for uncertainty extraction to NDDataArray

glue_astronomy/translators/nddata.py

@@ -8,6 +8,8 @@
 from glue.core import Data, Subset
 from glue.core.coordinates import Coordinates
 
+from .spectrum1d import SpectralCoordinates
+
 
 def _get_attribute(attribute, data):
     if isinstance(attribute, str):
@@ -53,6 +55,10 @@ def to_data(self, obj):
         data = Data(coords=obj.wcs)
         data['data'] = obj.data
         data.get_component('data').units = str(obj.unit)
+        if obj.uncertainty is not None:
+            data['uncertainty'] = getattr(
+                obj.uncertainty, 'array', obj.uncertainty
+            )
         data.meta.update(obj.meta)
         return data
 
@@ -70,20 +76,39 @@ def to_object(self, data_or_subset, attribute=None):
 
         data, subset_state = _get_data_and_subset_state(data_or_subset)
 
-        if isinstance(data.coords, WCS) or isinstance(data.coords, BaseHighLevelWCS):
+        if isinstance(data.coords, (WCS, BaseHighLevelWCS, SpectralCoordinates)):
             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')
 
+        component_labels = [d.label for d in data.component_ids()]
+        if attribute is None:
+            for desired_label in ('data', 'flux'):
+                if desired_label in component_labels:
+                    attribute = desired_label
+                    break
+
         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)
 
+        if 'uncertainty' in component_labels:
+            uncertainty = StdDevUncertainty(
+                data.get_component('uncertainty').data
+            )
+        else:
+            uncertainty = None
+
         result = NDDataArray(
-            values, unit=component.units, mask=mask, wcs=wcs, meta=data.meta
+            values,
+            unit=component.units,
+            mask=mask,
+            wcs=wcs,
+            meta=data.meta,
+            uncertainty=uncertainty
         )
 
         return result

glue_astronomy/translators/tests/test_nddata.py

@@ -136,6 +136,22 @@ def test_from_nddata(cls, kwargs, data_attr):
     assert image_new.unit is u.Jy
 
 
+def test_nddata_uncertainty():
+    data = [[2, 3], [4, 5]] * u.Jy
+    uncertainty = StdDevUncertainty(np.ones_like(data))
+    spec = NDDataArray(
+        data=data,
+        uncertainty=uncertainty
+    )
+    data_collection = DataCollection()
+    data_collection['data'] = spec
+
+    data = data_collection['data']
+    spec_new = data.get_object(NDDataArray)
+    assert isinstance(spec_new.uncertainty, StdDevUncertainty)
+    assert_equal(spec_new.uncertainty.array, uncertainty.array)
+
+
 @pytest.mark.parametrize('with_wcs', (False, True))
 def test_from_ccddata(with_wcs):