Use NDAstroData.unit

astrodata/core.py

@@ -321,6 +321,15 @@ def nddata(self):
         """
         return self._nddata[0] if self.is_single else self._nddata
 
+    @nddata.setter
+    @assign_only_single_slice
+    def nddata(self, new_nddata):
+        self.data = new_nddata.data
+        self.unit = new_nddata.unit
+        self.uncertainty = new_nddata.uncertainty
+        self.mask = new_nddata.mask
+        self.wcs = new_nddata.wcs
+
     def table(self):
         # FIXME: do we need this in addition to .tables ?
         return self._tables.copy()
@@ -444,6 +453,20 @@ def wcs(self):
     def wcs(self, value):
         self.nddata.wcs = value
 
+    @property
+    @returns_list
+    def unit(self):
+        """
+        A list of `astropy.units.Unit` objects (or a single object if this
+        is a single slice) attached to the science data, for each extension.
+        """
+        return [nd.unit for nd in self._nddata]
+
+    @unit.setter
+    @assign_only_single_slice
+    def unit(self, value):
+        self.nddata.unit = value
+
     def __iter__(self):
         if self.is_single:
             yield self
@@ -849,7 +872,8 @@ def _process_pixel_plane(self, pixim, name=None, top_level=False,
 
         return nd
 
-    def _append_array(self, data, name=None, header=None, add_to=None):
+    def _append_array(self, data, name=None, header=None, add_to=None,
+                      unit=None):
         if name in {'DQ', 'VAR'}:
             raise ValueError(f"'{name}' need to be associated to a "
                              f"'{DEFAULT_EXTENSION}' one")
@@ -866,21 +890,22 @@ def _append_array(self, data, name=None, header=None, add_to=None):
             hdu = fits.ImageHDU(data, header=header)
             hdu.header['EXTNAME'] = hname
             ret = self._append_imagehdu(hdu, name=hname, header=None,
-                                        add_to=None)
+                                        add_to=None, unit=unit)
         else:
             ret = add_to.meta['other'][name] = data
 
         return ret
 
-    def _append_imagehdu(self, hdu, name, header, add_to):
+    def _append_imagehdu(self, hdu, name, header, add_to, unit):
         if name in {'DQ', 'VAR'} or add_to is not None:
-            return self._append_array(hdu.data, name=name, add_to=add_to)
+            return self._append_array(hdu.data, name=name, add_to=add_to,
+                                      unit=unit)
         else:
             nd = self._process_pixel_plane(hdu, name=name, top_level=True,
                                            custom_header=header)
-            return self._append_nddata(nd, name, add_to=None)
+            return self._append_nddata(nd, name, add_to=None, unit=unit)
 
-    def _append_raw_nddata(self, raw_nddata, name, header, add_to):
+    def _append_raw_nddata(self, raw_nddata, name, header, add_to, unit):
         # We want to make sure that the instance we add is whatever we specify
         # as NDDataObject, instead of the random one that the user may pass
         top_level = add_to is None
@@ -889,9 +914,10 @@ def _append_raw_nddata(self, raw_nddata, name, header, add_to):
         processed_nddata = self._process_pixel_plane(raw_nddata,
                                                      top_level=top_level,
                                                      custom_header=header)
-        return self._append_nddata(processed_nddata, name=name, add_to=add_to)
+        return self._append_nddata(processed_nddata, name=name, add_to=add_to,
+                                   unit=unit)
 
-    def _append_nddata(self, new_nddata, name, add_to):
+    def _append_nddata(self, new_nddata, name, add_to, unit):
         # NOTE: This method is only used by others that have constructed NDData
         # according to our internal format. We don't accept new headers at this
         # point, and that's why it's missing from the signature.  'name' is
@@ -900,6 +926,9 @@ def _append_nddata(self, new_nddata, name, add_to):
             raise TypeError("You can only append NDData derived instances "
                             "at the top level")
 
+        if new_nddata.unit is None:
+            new_nddata.unit = unit
+
         hd = new_nddata.meta['header']
         hname = hd.get('EXTNAME', DEFAULT_EXTENSION)
         if hname == DEFAULT_EXTENSION:
@@ -910,7 +939,7 @@ def _append_nddata(self, new_nddata, name, add_to):
 
         return new_nddata
 
-    def _append_table(self, new_table, name, header, add_to):
+    def _append_table(self, new_table, name, header, add_to, unit):
         tb = _process_table(new_table, name, header)
         hname = tb.meta['header'].get('EXTNAME')
 
@@ -942,7 +971,7 @@ def find_next_num(tables):
             add_to.meta['other'][hname] = tb
         return tb
 
-    def _append_astrodata(self, ad, name, header, add_to):
+    def _append_astrodata(self, ad, name, header, add_to, unit):
         if not ad.is_single:
             raise ValueError("Cannot append AstroData instances that are "
                              "not single slices")
@@ -954,9 +983,10 @@ def _append_astrodata(self, ad, name, header, add_to):
         if header is not None:
             new_nddata.meta['header'] = deepcopy(header)
 
-        return self._append_nddata(new_nddata, name=None, add_to=None)
+        return self._append_nddata(new_nddata, name=None, add_to=None,
+                                   unit=unit)
 
-    def _append(self, ext, name=None, header=None, add_to=None):
+    def _append(self, ext, name=None, header=None, add_to=None, unit='adu'):
         """
         Internal method to dispatch to the type specific methods. This is
         called either by ``.append`` to append on top-level objects only or
@@ -972,12 +1002,14 @@ def _append(self, ext, name=None, header=None, add_to=None):
 
         for bases, method in dispatcher:
             if isinstance(ext, bases):
-                return method(ext, name=name, header=header, add_to=add_to)
+                return method(ext, name=name, header=header, add_to=add_to,
+                              unit=unit)
 
         # Assume that this is an array for a pixel plane
-        return self._append_array(ext, name=name, header=header, add_to=add_to)
+        return self._append_array(ext, name=name, header=header, add_to=add_to,
+                                  unit=unit)
 
-    def append(self, ext, name=None, header=None):
+    def append(self, ext, name=None, header=None, unit='adu'):
         """
         Adds a new top-level extension.
 
@@ -998,6 +1030,8 @@ def append(self, ext, name=None, header=None):
             It can consist in a combination of numbers and letters, with the
             restriction that the letters have to be all capital, and the first
             character cannot be a number ("[A-Z][A-Z0-9]*").
+        header : `astropy.io.fits.Header`
+            FITS Header to be associated with the NDData or Table object.
 
         Returns
         --------
@@ -1032,7 +1066,7 @@ def append(self, ext, name=None, header=None):
                           UserWarning)
             name = name.upper()
 
-        return self._append(ext, name=name, header=header)
+        return self._append(ext, name=name, header=header, unit=unit)
 
     @classmethod
     def read(cls, source, extname_parser=None):

astrodata/fits.py

@@ -30,6 +30,14 @@
 NO_DEFAULT = object()
 LOGGER = logging.getLogger(__name__)
 
+known_invalid_fits_unit_strings = {
+    'ELECTRONS/S': u.electron/u.s,
+    'ELECTRONS': u.electron,
+    'ELECTRON': u.electron,
+    'electrons': u.electron,
+    'electron': u.electron,
+}
+
 
 class FitsHeaderCollection:
     """Group access to a list of FITS Header-like objects.
@@ -391,7 +399,7 @@ def _prepare_hdulist(hdulist, default_extension='SCI', extname_parser=None):
     return HDUList(sorted(new_list, key=fits_ext_comp_key))
 
 
-def read_fits(cls, source, extname_parser=None):
+def read_fits(cls, source, extname_parser=None, default_unit='adu'):
     """
     Takes either a string (with the path to a file) or an HDUList as input, and
     tries to return a populated AstroData (or descendant) instance.
@@ -477,12 +485,24 @@ def associated_extensions(ver):
                 not isinstance(parts['uncertainty'], FitsLazyLoadable)):
             parts['uncertainty'] = ADVarianceUncertainty(parts['uncertainty'])
 
+        bunit = header.get('BUNIT')
+        if bunit:
+            if bunit in known_invalid_fits_unit_strings:
+                bunit = known_invalid_fits_unit_strings[bunit]
+            try:
+                unit = u.Unit(bunit, format='fits')
+            except ValueError:
+                unit = u.Unit(default_unit)
+        else:
+            unit = u.Unit(default_unit)
+
         # Create the NDData object
         nd = NDDataObject(
             data=parts['data'],
             uncertainty=parts['uncertainty'],
             mask=parts['mask'],
             meta={'header': header},
+            unit=unit,
         )
 
         if parts['wcs'] is not None:
@@ -570,9 +590,17 @@ def ad_to_hdulist(ad):
                 if 'APPROXIMATE' not in wcs_dict.get('FITS-WCS', ''):
                     wcs = None  # There's no need to create a WCS extension
 
-        hdul.append(new_imagehdu(ext.data, header, 'SCI'))
+        data_hdu = new_imagehdu(ext.data, header, 'SCI')
+
+        if ext.unit is not None and ext.unit is not u.dimensionless_unscaled:
+            data_hdu.header['BUNIT'] = (ext.unit.to_string(),
+                                        "Physical units of the array values")
+
+        hdul.append(data_hdu)
+
         if ext.uncertainty is not None:
             hdul.append(new_imagehdu(ext.uncertainty.array, header, 'VAR'))
+
         if ext.mask is not None:
             hdul.append(new_imagehdu(ext.mask, header, 'DQ'))
 

astrodata/nddata.py

@@ -8,8 +8,8 @@
 from copy import deepcopy
 from functools import reduce
 
+import astropy.units as u
 import numpy as np
-
 from astropy.io.fits import ImageHDU
 from astropy.modeling import Model, models
 from astropy.nddata import (NDArithmeticMixin, NDData, NDSlicingMixin,
@@ -369,6 +369,67 @@ def variance(self, value):
         self.uncertainty = (ADVarianceUncertainty(value) if value is not None
                             else None)
 
+    # Override unit so that we can add a setter.
+    @property
+    def unit(self):
+        return self._unit
+
+    @unit.setter
+    def unit(self, value):
+        if value is None:
+            self._unit = None
+        else:
+            self._unit = u.Unit(value)
+
+    def convert_unit_to(self, unit, equivalencies=[]):
+        """
+        Returns a new `NDData` object whose values have been converted
+        to a new unit.
+
+        Copied from Astropy's NDDataArray.
+
+        Parameters
+        ----------
+        unit : `astropy.units.UnitBase` instance or str
+            The unit to convert to.
+
+        equivalencies : list of equivalence pairs, optional
+           A list of equivalence pairs to try if the units are not
+           directly convertible.  See :ref:`unit_equivalencies`.
+
+        Returns
+        -------
+        result : `~astropy.nddata.NDData`
+            The resulting dataset
+
+        Raises
+        ------
+        UnitsError
+            If units are inconsistent.
+
+        """
+        if self.unit is None:
+            raise ValueError("No unit specified on source data")
+        data = self.unit.to(unit, self.data, equivalencies=equivalencies)
+        if self.uncertainty is not None:
+            uncertainty_values = self.unit.to(unit, self.uncertainty.array,
+                                              equivalencies=equivalencies)
+            # should work for any uncertainty class
+            uncertainty = self.uncertainty.__class__(uncertainty_values)
+        else:
+            uncertainty = None
+        if self.mask is not None:
+            new_mask = self.mask.copy()
+        else:
+            new_mask = None
+        # Call __class__ in case we are dealing with an inherited type
+        result = self.__class__(data, uncertainty=uncertainty,
+                                mask=new_mask,
+                                wcs=self.wcs,
+                                meta=self.meta, unit=unit)
+
+        return result
+
     def set_section(self, section, input):
         """
         Sets only a section of the data. This method is meant to prevent