diff --git a/glue_astronomy/translators/nddata.py b/glue_astronomy/translators/nddata.py index cc684b6..fc7104c 100644 --- a/glue_astronomy/translators/nddata.py +++ b/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 diff --git a/glue_astronomy/translators/tests/test_nddata.py b/glue_astronomy/translators/tests/test_nddata.py index bbe9b56..e55a557 100644 --- a/glue_astronomy/translators/tests/test_nddata.py +++ b/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):