Merge pull request #86 from bmorris3/better-nddata-support

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, UNCERT_REF
+
 
 def _get_attribute(attribute, data):
     if isinstance(attribute, str):
@@ -53,6 +55,9 @@ 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:
+            uncert = obj.uncertainty.represent_as(StdDevUncertainty)
+            data['uncertainty'] = uncert.array
         data.meta.update(obj.meta)
         return data
 
@@ -70,20 +75,42 @@ 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:
+            uncert_cls = UNCERT_REF[
+                data.meta.get('uncertainty_type', 'std')
+            ]
+            uncertainty = uncert_cls(
+                data.get_component('uncertainty').data
+            ).represent_as(StdDevUncertainty)
+        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

@@ -4,7 +4,9 @@
 
 from astropy import units as u
 from astropy.nddata import CCDData, NDDataArray
-from astropy.nddata.nduncertainty import StdDevUncertainty
+from astropy.nddata.nduncertainty import (
+    StdDevUncertainty, VarianceUncertainty, InverseVariance
+)
 from astropy.wcs import WCS
 
 from glue.core import Data, DataCollection
@@ -136,6 +138,30 @@ def test_from_nddata(cls, kwargs, data_attr):
     assert image_new.unit is u.Jy
 
 
+@pytest.mark.parametrize(
+    'uncertainty_type', (
+            StdDevUncertainty, VarianceUncertainty, InverseVariance
+    )
+)
+def test_nddata_uncertainty(uncertainty_type):
+    data = [[2, 3], [4, 5]] * u.Jy
+    uncertainty = uncertainty_type(np.ones_like(data.value))
+    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)
+    if isinstance(uncertainty, StdDevUncertainty):
+        assert_equal(spec_new.uncertainty.array, uncertainty.array)
+    else:
+        assert_equal(spec_new.uncertainty.array, uncertainty.represent_as(StdDevUncertainty).array)
+
+
 @pytest.mark.parametrize('with_wcs', (False, True))
 def test_from_ccddata(with_wcs):