Better handled Spectrum1D images across classes

CHANGES.rst

@@ -13,7 +13,10 @@ API Changes
 - Renamed KosmosTrace as FitTrace, a conglomerate class for traces that are fit
 to images instead of predetermined [#128]
 - The default number of bins for FitTrace is now its associated image's number
-of dispersion pixels instead of 20. Its default peak_method is now 'max'. [#128]
+of dispersion pixels instead of 20. Its default peak_method is now 'max' [#128]
+- All operations now accept Spectrum1D and Quantity-type images. All accepted
+image types are now processed internally as Spectrum1D objects [#144]
+- All operations' ``image`` attributes are now coerced Spectrum1D objects [#144]
 
 Bug Fixes
 ^^^^^^^^^

setup.cfg

@@ -14,11 +14,11 @@ github_project = astropy/specreduce
 [options]
 zip_safe = False
 packages = find:
-python_requires = >=3.7
+python_requires = >=3.8
 setup_requires = setuptools_scm
 install_requires =
     astropy
-    specutils
+    specutils>=1.9.1
     synphot
     matplotlib
     photutils

specreduce/background.py

@@ -5,16 +5,19 @@
 
 import numpy as np
 from astropy.nddata import NDData
+from astropy.utils.decorators import deprecated_attribute
 from astropy import units as u
+from specutils import Spectrum1D
 
-from specreduce.extract import _ap_weight_image, _to_spectrum1d_pixels
+from specreduce.core import _ImageParser
+from specreduce.extract import _ap_weight_image
 from specreduce.tracing import Trace, FlatTrace
 
 __all__ = ['Background']
 
 
 @dataclass
-class Background:
+class Background(_ImageParser):
     """
     Determine the background from an image for subtraction.
 
@@ -27,7 +30,7 @@ class Background:
 
     Parameters
     ----------
-    image : `~astropy.nddata.NDData` or array-like
+    image : `~astropy.nddata.NDData`-like or array-like
         image with 2-D spectral image data
     traces : List
         list of trace objects (or integers to define FlatTraces) to
@@ -54,13 +57,16 @@ class Background:
     disp_axis: int = 1
     crossdisp_axis: int = 0
 
+    # TO-DO: update bkg_array with Spectrum1D alternative (is bkg_image enough?)
+    bkg_array = deprecated_attribute('bkg_array', '1.3')
+
     def __post_init__(self):
         """
         Determine the background from an image for subtraction.
 
         Parameters
         ----------
-        image : `~astropy.nddata.NDData` or array-like
+        image : `~astropy.nddata.NDData`-like or array-like
             image with 2-D spectral image data
         traces : List
             list of trace objects (or integers to define FlatTraces) to
@@ -86,17 +92,18 @@ def _to_trace(trace):
                     raise ValueError('trace_object.trace_pos must be >= 1')
             return trace
 
+        self.image = self._parse_image(self.image)
+
         if self.width < 0:
             raise ValueError("width must be positive")
-
         if self.width == 0:
-            self.bkg_array = np.zeros(self.image.shape[self.disp_axis])
+            self._bkg_array = np.zeros(self.image.shape[self.disp_axis])
             return
 
         if isinstance(self.traces, Trace):
             self.traces = [self.traces]
 
-        bkg_wimage = np.zeros_like(self.image, dtype=np.float64)
+        bkg_wimage = np.zeros_like(self.image.data, dtype=np.float64)
         for trace in self.traces:
             trace = _to_trace(trace)
             windows_max = trace.trace.data.max() + self.width/2
@@ -127,12 +134,13 @@ def _to_trace(trace):
         self.bkg_wimage = bkg_wimage
 
         if self.statistic == 'average':
-            self.bkg_array = np.average(self.image, weights=self.bkg_wimage,
-                                        axis=self.crossdisp_axis)
+            self._bkg_array = np.average(self.image.data,
+                                         weights=self.bkg_wimage,
+                                         axis=self.crossdisp_axis)
         elif self.statistic == 'median':
-            med_image = self.image.copy()
+            med_image = self.image.data.copy()
             med_image[np.where(self.bkg_wimage) == 0] = np.nan
-            self.bkg_array = np.nanmedian(med_image, axis=self.crossdisp_axis)
+            self._bkg_array = np.nanmedian(med_image, axis=self.crossdisp_axis)
         else:
             raise ValueError("statistic must be 'average' or 'median'")
 
@@ -150,9 +158,11 @@ def two_sided(cls, image, trace_object, separation, **kwargs):
 
         Parameters
         ----------
-        image : nddata-compatible image
-            image with 2-D spectral image data
-        trace_object: Trace
+        image : `~astropy.nddata.NDData`-like or array-like
+            Image with 2-D spectral image data. Assumes cross-dispersion
+            (spatial) direction is axis 0 and dispersion (wavelength)
+            direction is axis 1.
+        trace_object: `~specreduce.tracing.Trace`
             estimated trace of the spectrum to center the background traces
         separation: float
             separation from ``trace_object`` for the background regions
@@ -167,6 +177,7 @@ def two_sided(cls, image, trace_object, separation, **kwargs):
         crossdisp_axis : int
             cross-dispersion axis
         """
+        image = cls._parse_image(cls, image)
         kwargs['traces'] = [trace_object-separation, trace_object+separation]
         return cls(image=image, **kwargs)
 
@@ -183,9 +194,11 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
 
         Parameters
         ----------
-        image : nddata-compatible image
-            image with 2-D spectral image data
-        trace_object: Trace
+        image : `~astropy.nddata.NDData`-like or array-like
+            Image with 2-D spectral image data. Assumes cross-dispersion
+            (spatial) direction is axis 0 and dispersion (wavelength)
+            direction is axis 1.
+        trace_object: `~specreduce.tracing.Trace`
             estimated trace of the spectrum to center the background traces
         separation: float
             separation from ``trace_object`` for the background, positive will be
@@ -201,6 +214,7 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
         crossdisp_axis : int
             cross-dispersion axis
         """
+        image = cls._parse_image(cls, image)
         kwargs['traces'] = [trace_object+separation]
         return cls(image=image, **kwargs)
 
@@ -210,28 +224,32 @@ def bkg_image(self, image=None):
 
         Parameters
         ----------
-        image : nddata-compatible image or None
-            image with 2-D spectral image data.  If None, will extract
-            the background from ``image`` used to initialize the class.
+        image : `~astropy.nddata.NDData`-like or array-like, optional
+            Image with 2-D spectral image data. Assumes cross-dispersion
+            (spatial) direction is axis 0 and dispersion (wavelength)
+            direction is axis 1. If None, will extract the background
+            from ``image`` used to initialize the class. [default: None]
 
         Returns
         -------
-        array with same shape as ``image``.
+        Spectrum1D object with same shape as ``image``.
         """
-        if image is None:
-            image = self.image
-
-        return np.tile(self.bkg_array, (image.shape[0], 1))
+        image = self._parse_image(image)
+        return Spectrum1D(np.tile(self._bkg_array,
+                                  (image.shape[0], 1)) * image.unit,
+                          spectral_axis=image.spectral_axis)
 
     def bkg_spectrum(self, image=None):
         """
         Expose the 1D spectrum of the background.
 
         Parameters
         ----------
-        image : nddata-compatible image or None
-            image with 2-D spectral image data.  If None, will extract
-            the background from ``image`` used to initialize the class.
+        image : `~astropy.nddata.NDData`-like or array-like, optional
+            Image with 2-D spectral image data. Assumes cross-dispersion
+            (spatial) direction is axis 0 and dispersion (wavelength)
+            direction is axis 1. If None, will extract the background
+            from ``image`` used to initialize the class. [default: None]
 
         Returns
         -------
@@ -240,10 +258,15 @@ def bkg_spectrum(self, image=None):
             units as the input image (or u.DN if none were provided) and
             the spectral axis expressed in pixel units.
         """
-        bkg_image = self.bkg_image(image=image)
+        bkg_image = self.bkg_image(image)
 
-        ext1d = np.sum(bkg_image, axis=self.crossdisp_axis)
-        return _to_spectrum1d_pixels(ext1d * getattr(image, 'unit', u.DN))
+        try:
+            return bkg_image.collapse(np.sum, axis=self.crossdisp_axis)
+        except u.UnitTypeError:
+            # can't collapse with a spectral axis in pixels because
+            # SpectralCoord only allows frequency/wavelength equivalent units...
+            ext1d = np.sum(bkg_image.flux, axis=self.crossdisp_axis)
+            return Spectrum1D(ext1d, bkg_image.spectral_axis)
 
     def sub_image(self, image=None):
         """
@@ -259,14 +282,16 @@ def sub_image(self, image=None):
         -------
         array with same shape as ``image``
         """
-        if image is None:
-            image = self.image
+        image = self._parse_image(image)
 
-        if isinstance(image, NDData):
-            # https://docs.astropy.org/en/stable/nddata/mixins/ndarithmetic.html
-            return image.subtract(self.bkg_image(image)*image.unit)
-        else:
-            return image - self.bkg_image(image)
+        # a compare_wcs argument is needed for Spectrum1D.subtract() in order to
+        # avoid a TypeError from SpectralCoord when image's spectral axis is in
+        # pixels. it is not needed when image's spectral axis has physical units
+        kwargs = ({'compare_wcs': None} if image.spectral_axis.unit == u.pix
+                  else {})
+
+        # https://docs.astropy.org/en/stable/nddata/mixins/ndarithmetic.html
+        return image.subtract(self.bkg_image(image), **kwargs)
 
     def sub_spectrum(self, image=None):
         """
@@ -287,8 +312,13 @@ def sub_spectrum(self, image=None):
         """
         sub_image = self.sub_image(image=image)
 
-        ext1d = np.sum(sub_image, axis=self.crossdisp_axis)
-        return _to_spectrum1d_pixels(ext1d * getattr(image, 'unit', u.DN))
+        try:
+            return sub_image.collapse(np.sum, axis=self.crossdisp_axis)
+        except u.UnitTypeError:
+            # can't collapse with a spectral axis in pixels because
+            # SpectralCoord only allows frequency/wavelength equivalent units...
+            ext1d = np.sum(sub_image.flux, axis=self.crossdisp_axis)
+            return Spectrum1D(ext1d, spectral_axis=sub_image.spectral_axis)
 
     def __rsub__(self, image):
         """

specreduce/core.py

@@ -1,13 +1,86 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import inspect
+import numpy as np
+
+from astropy import units as u
+from astropy.nddata import VarianceUncertainty
 from dataclasses import dataclass
+from specutils import Spectrum1D
 
 __all__ = ['SpecreduceOperation']
 
 
+class _ImageParser:
+    """
+    Coerces images from accepted formats to Spectrum1D objects for
+    internal use in specreduce's operation classes.
+
+    Fills any and all of uncertainty, mask, units, and spectral axis
+    that are missing in the provided image with generic values.
+    Accepted image types are:
+
+        - `~specutils.spectra.spectrum1d.Spectrum1D` (preferred)
+        - `~astropy.nddata.ccddata.CCDData`
+        - `~astropy.nddata.ndddata.NDDData`
+        - `~astropy.units.quantity.Quantity`
+        - `~numpy.ndarray`
+    """
+    def _parse_image(self, image, disp_axis=1):
+        """
+        Convert all accepted image types to a consistently formatted
+        Spectrum1D object.
+
+        Parameters
+        ----------
+        image : `~astropy.nddata.NDData`-like or array-like, required
+            The image to be parsed. If None, defaults to class' own
+            image attribute.
+        disp_axis : int, optional
+            The index of the image's dispersion axis. Should not be
+            changed until operations can handle variable image
+            orientations. [default: 1]
+        """
+
+        # would be nice to handle (cross)disp_axis consistently across
+        # operations (public attribute? private attribute? argument only?) so
+        # it can be called from self instead of via kwargs...
+
+        if image is None:
+            # useful for Background's instance methods
+            return self.image
+
+        if isinstance(image, np.ndarray):
+            img = image
+        elif isinstance(image, u.quantity.Quantity):
+            img = image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = image.data
+
+        # mask and uncertainty are set as None when they aren't specified upon
+        # creating a Spectrum1D object, so we must check whether these
+        # attributes are absent *and* whether they are present but set as None
+        if getattr(image, 'mask', None) is not None:
+            mask = image.mask
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        if getattr(image, 'uncertainty', None) is not None:
+            uncertainty = image.uncertainty
+        else:
+            uncertainty = VarianceUncertainty(np.ones(img.shape))
+
+        unit = getattr(image, 'unit', u.Unit('DN'))
+
+        spectral_axis = getattr(image, 'spectral_axis',
+                                np.arange(img.shape[disp_axis]) * u.pix)
+
+        return Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                          uncertainty=uncertainty, mask=mask)
+
+
 @dataclass
-class SpecreduceOperation:
+class SpecreduceOperation(_ImageParser):
     """
     An operation to perform as part of a spectroscopic reduction pipeline.
 

specreduce/extract.py

@@ -7,7 +7,7 @@
 
 from astropy import units as u
 from astropy.modeling import Model, models, fitting
-from astropy.nddata import NDData
+from astropy.nddata import NDData, VarianceUncertainty
 
 from specreduce.core import SpecreduceOperation
 from specreduce.tracing import Trace, FlatTrace
@@ -117,12 +117,6 @@ def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
     return wimage
 
 
-def _to_spectrum1d_pixels(fluxes):
-    # TODO: add wavelength units, uncertainty and mask to spectrum1D object
-    return Spectrum1D(spectral_axis=np.arange(len(fluxes)) * u.pixel,
-                      flux=fluxes)
-
-
 @dataclass
 class BoxcarExtract(SpecreduceOperation):
     """
@@ -137,15 +131,15 @@ class BoxcarExtract(SpecreduceOperation):
 
     Parameters
     ----------
-    image : nddata-compatible image
+    image : `~astropy.nddata.NDData`-like or array-like, required
         image with 2-D spectral image data
-    trace_object : Trace
+    trace_object : Trace, required
         trace object
-    width : float
+    width : float, optional
         width of extraction aperture in pixels
-    disp_axis : int
+    disp_axis : int, optional
         dispersion axis
-    crossdisp_axis : int
+    crossdisp_axis : int, optional
         cross-dispersion axis
 
     Returns
@@ -171,15 +165,15 @@ def __call__(self, image=None, trace_object=None, width=None,
 
         Parameters
         ----------
-        image : nddata-compatible image
+        image : `~astropy.nddata.NDData`-like or array-like, required
             image with 2-D spectral image data
-        trace_object : Trace
+        trace_object : Trace, required
             trace object
-        width : float
+        width : float, optional
             width of extraction aperture in pixels [default: 5]
-        disp_axis : int
+        disp_axis : int, optional
             dispersion axis [default: 1]
-        crossdisp_axis : int
+        crossdisp_axis : int, optional
             cross-dispersion axis [default: 0]
 
 
@@ -195,6 +189,9 @@ def __call__(self, image=None, trace_object=None, width=None,
         disp_axis = disp_axis if disp_axis is not None else self.disp_axis
         crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
 
+        # handle image processing based on its type
+        self.image = self._parse_image(image)
+
         # TODO: this check can be removed if/when implemented as a check in FlatTrace
         if isinstance(trace_object, FlatTrace):
             if trace_object.trace_pos < 1:
@@ -204,16 +201,17 @@ def __call__(self, image=None, trace_object=None, width=None,
             raise ValueError("width must be positive")
 
         # weight image to use for extraction
-        wimage = _ap_weight_image(
+        wimg = _ap_weight_image(
             trace_object,
             width,
             disp_axis,
             crossdisp_axis,
-            image.shape)
+            self.image.shape)
 
         # extract
-        ext1d = np.sum(image * wimage, axis=crossdisp_axis)
-        return _to_spectrum1d_pixels(ext1d * getattr(image, 'unit', u.DN))
+        ext1d = np.sum(self.image.data * wimg, axis=crossdisp_axis)
+        return Spectrum1D(ext1d * self.image.unit,
+                          spectral_axis=self.image.spectral_axis)
 
 
 @dataclass
@@ -225,7 +223,7 @@ class HorneExtract(SpecreduceOperation):
     Parameters
     ----------
 
-    image : `~astropy.nddata.NDData` or array-like, required
+    image : `~astropy.nddata.NDData`-like or array-like, required
         The input 2D spectrum from which to extract a source. An
         NDData object must specify uncertainty and a mask. An array
         requires use of the ``variance``, ``mask``, & ``unit`` arguments.
@@ -260,7 +258,7 @@ class HorneExtract(SpecreduceOperation):
     unit : `~astropy.units.Unit` or str, optional
         (Only used if ``image`` is not an NDData object.)
         The associated unit for the data in ``image``. If blank,
-        fluxes are interpreted as unitless. [default: None]
+        fluxes are interpreted in DN. [default: None]
 
     """
     image: NDData
@@ -277,6 +275,120 @@ class HorneExtract(SpecreduceOperation):
     def spectrum(self):
         return self.__call__()
 
+    def _parse_image(self, image,
+                     variance=None, mask=None, unit=None, disp_axis=1):
+        """
+        Convert all accepted image types to a consistently formatted
+        Spectrum1D object.
+
+        HorneExtract needs its own version of this method because it is
+        more stringent in its requirements for input images. The extra
+        arguments are needed to handle cases where these parameters were
+        specified as arguments and those where they came as attributes
+        of the image object.
+
+        Parameters
+        ----------
+        image : `~astropy.nddata.NDData`-like or array-like, required
+            The image to be parsed. If None, defaults to class' own
+            image attribute.
+        variance : `~numpy.ndarray`, optional
+            (Only used if ``image`` is not an NDData object.)
+            The associated variances for each pixel in the image. Must
+            have the same dimensions as ``image``. If all zeros, the variance
+            will be ignored and treated as all ones.  If any zeros, those
+            elements will be excluded via masking.  If any negative values,
+            an error will be raised.
+        mask : `~numpy.ndarray`, optional
+            (Only used if ``image`` is not an NDData object.)
+            Whether to mask each pixel in the image. Must have the same
+            dimensions as ``image``. If blank, all non-NaN pixels are
+            unmasked.
+        unit : `~astropy.units.Unit` or str, optional
+            (Only used if ``image`` is not an NDData object.)
+            The associated unit for the data in ``image``. If blank,
+            fluxes are interpreted in DN.
+        disp_axis : int, optional
+            The index of the image's dispersion axis. Should not be
+            changed until operations can handle variable image
+            orientations. [default: 1]
+        """
+
+        if isinstance(image, np.ndarray):
+            img = image
+        elif isinstance(image, u.quantity.Quantity):
+            img = image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = image.data
+
+        # mask is set as None when not specified upon creating a Spectrum1D
+        # object, so we must check whether it is absent *and* whether it's
+        # present but set as None
+        if getattr(image, 'mask', None) is not None:
+            mask = image.mask
+        elif mask is not None:
+            pass
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        if img.shape != mask.shape:
+            raise ValueError('image and mask shapes must match.')
+
+        # Process uncertainties, converting to variances when able and throwing
+        # an error when uncertainties are missing or less easily converted
+        if (hasattr(image, 'uncertainty')
+                and image.uncertainty is not None):
+            if image.uncertainty.uncertainty_type == 'var':
+                variance = image.uncertainty.array
+            elif image.uncertainty.uncertainty_type == 'std':
+                warnings.warn("image NDData object's uncertainty "
+                              "interpreted as standard deviation. if "
+                              "incorrect, use VarianceUncertainty when "
+                              "assigning image object's uncertainty.")
+                variance = image.uncertainty.array**2
+            elif image.uncertainty.uncertainty_type == 'ivar':
+                variance = 1 / image.uncertainty.array
+            else:
+                # other options are InverseUncertainty and UnknownUncertainty
+                raise ValueError("image NDData object has unexpected "
+                                 "uncertainty type. instead, try "
+                                 "VarianceUncertainty or StdDevUncertainty.")
+        elif (hasattr(image, 'uncertainty')
+              and image.uncertainty is None):
+            # ignore variance arg to focus on updating NDData object
+            raise ValueError('image NDData object lacks uncertainty')
+        else:
+            if variance is None:
+                raise ValueError("if image is a numpy or Quantity array, a "
+                                 "variance must be specified. consider "
+                                 "wrapping it into one object by instead "
+                                 "passing an NDData image.")
+            elif image.shape != variance.shape:
+                raise ValueError("image and variance shapes must match")
+
+        if np.any(variance < 0):
+            raise ValueError("variance must be fully positive")
+        if np.all(variance == 0):
+            # technically would result in infinities, but since they're all
+            # zeros, we can override ones to simulate an unweighted case
+            variance = np.ones_like(variance)
+        if np.any(variance == 0):
+            # exclude such elements by editing the input mask
+            mask[variance == 0] = True
+            # replace the variances to avoid a divide by zero warning
+            variance[variance == 0] = np.nan
+
+        variance = VarianceUncertainty(variance)
+
+        unit = getattr(image, 'unit',
+                       u.Unit(unit) if unit is not None else u.Unit('DN'))
+
+        spectral_axis = getattr(image, 'spectral_axis',
+                                np.arange(img.shape[disp_axis]) * u.pix)
+
+        return Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                          uncertainty=variance, mask=mask)
+
     def __call__(self, image=None, trace_object=None,
                  disp_axis=None, crossdisp_axis=None,
                  bkgrd_prof=None,
@@ -288,7 +400,7 @@ def __call__(self, image=None, trace_object=None,
         Parameters
         ----------
 
-        image : `~astropy.nddata.NDData` or array-like, required
+        image : `~astropy.nddata.NDData`-like or array-like, required
             The input 2D spectrum from which to extract a source. An
             NDData object must specify uncertainty and a mask. An array
             requires use of the ``variance``, ``mask``, & ``unit`` arguments.
@@ -322,7 +434,7 @@ def __call__(self, image=None, trace_object=None,
         unit : `~astropy.units.Unit` or str, optional
             (Only used if ``image`` is not an NDData object.)
             The associated unit for the data in ``image``. If blank,
-            fluxes are interpreted as unitless.
+            fluxes are interpreted in DN.
 
 
         Returns
@@ -339,70 +451,16 @@ def __call__(self, image=None, trace_object=None,
         mask = mask if mask is not None else self.mask
         unit = unit if unit is not None else self.unit
 
-        # handle image and associated data based on image's type
-        if isinstance(image, NDData):
-            img = np.ma.array(image.data, mask=image.mask)
-            unit = image.unit if image.unit is not None else u.Unit()
-
-            if image.uncertainty is not None:
-                # prioritize NDData's uncertainty over variance argument
-                if image.uncertainty.uncertainty_type == 'var':
-                    variance = image.uncertainty.array
-                elif image.uncertainty.uncertainty_type == 'std':
-                    # NOTE: CCDData defaults uncertainties given as pure arrays
-                    # to std and logs a warning saying so upon object creation.
-                    # should we remind users again here?
-                    warnings.warn("image NDData object's uncertainty "
-                                  "interpreted as standard deviation. if "
-                                  "incorrect, use VarianceUncertainty when "
-                                  "assigning image object's uncertainty.")
-                    variance = image.uncertainty.array**2
-                elif image.uncertainty.uncertainty_type == 'ivar':
-                    variance = 1 / image.uncertainty.array
-                else:
-                    # other options are InverseVariance and UnknownVariance
-                    raise ValueError("image NDData object has unexpected "
-                                     "uncertainty type. instead, try "
-                                     "VarianceUncertainty or StdDevUncertainty.")
-            else:
-                # ignore variance arg to focus on updating NDData object
-                raise ValueError('image NDData object lacks uncertainty')
+        # parse image and replace optional arguments with updated values
+        self.image = self._parse_image(image, variance, mask, unit, disp_axis)
+        variance = self.image.uncertainty.array
+        unit = self.image.unit
 
-        else:
-            if variance is None:
-                raise ValueError('if image is a numpy array, a variance must '
-                                 'be specified. consider wrapping it into one '
-                                 'object by instead passing an NDData image.')
-            elif image.shape != variance.shape:
-                raise ValueError('image and variance shapes must match')
-
-            # check optional arguments, filling them in if absent
-            if mask is None:
-                mask = np.ma.masked_invalid(image).mask
-            elif image.shape != mask.shape:
-                raise ValueError('image and mask shapes must match.')
-
-            if isinstance(unit, str):
-                unit = u.Unit(unit)
-            else:
-                unit = unit if unit is not None else u.Unit()
-
-            # create image
-            img = np.ma.array(image, mask=mask)
-
-        if np.all(variance == 0):
-            # technically would result in infinities, but since they're all zeros
-            # we can just do the unweighted case by overriding with all ones
-            variance = np.ones_like(variance)
-
-        if np.any(variance < 0):
-            raise ValueError("variance must be fully positive")
-
-        if np.any(variance == 0):
-            # exclude these elements by editing the input mask
-            img.mask[variance == 0] = True
-            # replace the variances to avoid a divide by zero warning
-            variance[variance == 0] = np.nan
+        # mask any previously uncaught invalid values
+        or_mask = np.logical_or(mask,
+                                np.ma.masked_invalid(self.image.data).mask)
+        img = np.ma.masked_array(self.image.data, or_mask)
+        mask = img.mask
 
         # co-add signal in each image column
         ncols = img.shape[crossdisp_axis]
@@ -451,7 +509,8 @@ def __call__(self, image=None, trace_object=None,
         extraction = result * norms
 
         # convert the extraction to a Spectrum1D object
-        return _to_spectrum1d_pixels(extraction * unit)
+        return Spectrum1D(extraction * unit,
+                          spectral_axis=self.image.spectral_axis)
 
 
 @dataclass

specreduce/tests/test_background.py

@@ -2,7 +2,7 @@
 import numpy as np
 
 import astropy.units as u
-from astropy.nddata import CCDData
+from astropy.nddata import VarianceUncertainty
 from specutils import Spectrum1D
 
 from specreduce.background import Background
@@ -16,7 +16,11 @@
 image = np.ones(shape=(30, 10))
 for j in range(image.shape[0]):
     image[j, ::] *= j
-image = CCDData(image, unit=u.Jy)
+image = Spectrum1D(image * u.DN,
+                   uncertainty=VarianceUncertainty(np.ones_like(image)))
+image_um = Spectrum1D(image.flux,
+                      spectral_axis=np.arange(image.data.shape[1]) * u.um,
+                      uncertainty=VarianceUncertainty(np.ones_like(image.data)))
 
 
 def test_background():
@@ -28,12 +32,21 @@ def test_background():
     trace = FlatTrace(image, trace_pos)
     bkg_sep = 5
     bkg_width = 2
-    # all the following should be equivalent:
+
+    # all the following should be equivalent, whether image's spectral axis
+    # is in pixels or physical units:
     bg1 = Background(image, [trace-bkg_sep, trace+bkg_sep], width=bkg_width)
     bg2 = Background.two_sided(image, trace, bkg_sep, width=bkg_width)
     bg3 = Background.two_sided(image, trace_pos, bkg_sep, width=bkg_width)
-    assert np.allclose(bg1.bkg_array, bg2.bkg_array)
-    assert np.allclose(bg1.bkg_array, bg3.bkg_array)
+    assert np.allclose(bg1.bkg_image().flux, bg2.bkg_image().flux)
+    assert np.allclose(bg1.bkg_image().flux, bg3.bkg_image().flux)
+
+    bg4 = Background(image_um, [trace-bkg_sep, trace+bkg_sep], width=bkg_width)
+    bg5 = Background.two_sided(image_um, trace, bkg_sep, width=bkg_width)
+    bg6 = Background.two_sided(image_um, trace_pos, bkg_sep, width=bkg_width)
+    assert np.allclose(bg1.bkg_image().flux, bg4.bkg_image().flux)
+    assert np.allclose(bg1.bkg_image().flux, bg5.bkg_image().flux)
+    assert np.allclose(bg1.bkg_image().flux, bg6.bkg_image().flux)
 
     # test that creating a one_sided background works
     Background.one_sided(image, trace, bkg_sep, width=bkg_width)
@@ -45,16 +58,23 @@ def test_background():
     sub1 = image - bg1
     sub2 = bg1.sub_image(image)
     sub3 = bg1.sub_image()
-    assert np.allclose(sub1, sub2)
-    assert np.allclose(sub1, sub3)
+    assert np.allclose(sub1.flux, sub2.flux)
+    assert np.allclose(sub2.flux, sub3.flux)
+
+    sub4 = image_um - bg4
+    sub5 = bg4.sub_image(image_um)
+    sub6 = bg4.sub_image()
+    assert np.allclose(sub1.flux, sub4.flux)
+    assert np.allclose(sub4.flux, sub5.flux)
+    assert np.allclose(sub5.flux, sub6.flux)
 
     bkg_spec = bg1.bkg_spectrum()
     assert isinstance(bkg_spec, Spectrum1D)
     sub_spec = bg1.sub_spectrum()
     assert isinstance(sub_spec, Spectrum1D)
     # test that width==0 results in no background
     bg = Background.two_sided(image, trace, bkg_sep, width=0)
-    assert np.all(bg.bkg_image() == 0)
+    assert np.all(bg.bkg_image().flux == 0)
 
 
 def test_warnings_errors():

specreduce/tests/test_extract.py

@@ -2,7 +2,7 @@
 import pytest
 
 import astropy.units as u
-from astropy.nddata import CCDData
+from astropy.nddata import CCDData, VarianceUncertainty, UnknownUncertainty
 
 from specreduce.extract import BoxcarExtract, HorneExtract, OptimalExtract
 from specreduce.tracing import FlatTrace, ArrayTrace
@@ -79,7 +79,7 @@ def test_boxcar_array_trace():
     assert np.allclose(spectrum.flux.value, np.full_like(spectrum.flux.value, 75.))
 
 
-def test_horne_array_validation():
+def test_horne_image_validation():
     #
     # Test HorneExtract scenarios specific to its use with an image of
     # type `~numpy.ndarray` (instead of the default `~astropy.nddata.NDData`).
@@ -88,47 +88,64 @@ def test_horne_array_validation():
     extract = OptimalExtract(image.data, trace)  # equivalent to HorneExtract
 
     # an array-type image must come with a variance argument
-    with pytest.raises(ValueError, match=r'.*array.*variance.*specified.*'):
+    with pytest.raises(ValueError, match=r'.*variance must be specified.*'):
         ext = extract()
 
+    # an NDData-type image can't have an empty uncertainty attribute
+    with pytest.raises(ValueError, match=r'.*NDData object lacks uncertainty'):
+        ext = extract(image=image)
+
+    # an NDData-type image's uncertainty must be of type VarianceUncertainty
+    # or type StdDevUncertainty
+    with pytest.raises(ValueError, match=r'.*unexpected uncertainty type.*'):
+        err = UnknownUncertainty(np.ones_like(image))
+        image.uncertainty = err
+        ext = extract(image=image)
+
     # an array-type image must have the same dimensions as its variance argument
     with pytest.raises(ValueError, match=r'.*shapes must match.*'):
+        # remember variance, mask, and unit args are only checked if image
+        # object doesn't have those attributes (e.g., numpy and Quantity arrays)
         err = np.ones_like(image[0])
-        ext = extract(variance=err)
+        ext = extract(image=image.data, variance=err)
 
     # an array-type image must have the same dimensions as its mask argument
     with pytest.raises(ValueError, match=r'.*shapes must match.*'):
         err = np.ones_like(image)
         mask = np.zeros_like(image[0])
-        ext = extract(variance=err, mask=mask)
+        ext = extract(image=image.data, variance=err, mask=mask)
 
     # an array-type image given without mask and unit arguments is fine
-    # and produces a unitless result
+    # and produces an extraction with flux in DN and spectral axis in pixels
     err = np.ones_like(image)
-    ext = extract(variance=err)
-    assert ext.unit == u.Unit()
+    ext = extract(image=image.data, variance=err, mask=None, unit=None)
+    assert ext.unit == u.Unit('DN')
+    assert np.all(ext.spectral_axis
+                  == np.arange(image.shape[extract.disp_axis]) * u.pix)
 
 
 def test_horne_variance_errors():
     trace = FlatTrace(image, 3.0)
 
-    # all zeros are treated as non-weighted (give non-zero fluxes)
-    err = np.zeros_like(image)
-    mask = np.zeros_like(image)
-    extract = HorneExtract(image.data, trace, variance=err, mask=mask, unit=u.Jy)
+    # all zeros are treated as non-weighted (i.e., as non-zero fluxes)
+    image.uncertainty = VarianceUncertainty(np.zeros_like(image))
+    image.mask = np.zeros_like(image)
+    extract = HorneExtract(image, trace)
     ext = extract.spectrum
     assert not np.all(ext == 0)
 
-    # single zero value adjusts mask (does not raise error)
+    # single zero value adjusts mask and does not raise error
     err = np.ones_like(image)
-    err[0] = 0
-    mask = np.zeros_like(image)
-    ext = extract(variance=err, mask=mask, unit=u.Jy)
-    assert not np.all(ext == 0)
+    err[0][0] = 0
+    image.uncertainty = VarianceUncertainty(err)
+    ext = extract(image)
+    assert not np.all(ext == 1)
 
     # single negative value raises error
-    err = np.ones_like(image)
-    err[0] = -1
-    mask = np.zeros_like(image)
+    err = image.uncertainty.array
+    err[0][0] = -1
     with pytest.raises(ValueError, match='variance must be fully positive'):
-        ext = extract(variance=err, mask=mask, unit=u.Jy)
+        # remember variance, mask, and unit args are only checked if image
+        # object doesn't have those attributes (e.g., numpy and Quantity arrays)
+        ext = extract(image=image.data, variance=err,
+                      mask=image.mask, unit=u.Jy)

specreduce/tests/test_image_parsing.py

@@ -0,0 +1,104 @@
+import numpy as np
+
+from astropy import units as u
+from astropy.io import fits
+from astropy.nddata import CCDData, NDData, VarianceUncertainty
+from astropy.utils.data import download_file
+
+from specreduce.extract import HorneExtract
+from specreduce.tracing import FlatTrace
+from specutils import Spectrum1D, SpectralAxis
+
+# fetch test image
+fn = download_file('https://stsci.box.com/shared/static/exnkul627fcuhy5akf2gswytud5tazmw.fits',
+                   cache=True)
+
+# duplicate image in all accepted formats
+# (one Spectrum1D variant has a physical spectral axis; the other is in pixels)
+img = fits.getdata(fn).T
+flux = img * u.MJy / u.sr
+sax = SpectralAxis(np.linspace(14.377, 3.677, flux.shape[-1]) * u.um)
+unc = VarianceUncertainty(np.random.rand(*flux.shape))
+
+all_images = {}
+all_images['arr'] = img
+all_images['s1d'] = Spectrum1D(flux, spectral_axis=sax, uncertainty=unc)
+all_images['s1d_pix'] = Spectrum1D(flux, uncertainty=unc)
+all_images['ccd'] = CCDData(img, uncertainty=unc, unit=flux.unit)
+all_images['ndd'] = NDData(img, uncertainty=unc, unit=flux.unit)
+all_images['qnt'] = img * flux.unit
+
+# save default values used for spectral axis and uncertainty when they are not
+# available from the image object or provided by the user
+sax_def = np.arange(img.shape[1]) * u.pix
+unc_def = np.ones_like(img)
+
+
+# (for use inside tests)
+def compare_images(key, collection, compare='s1d'):
+    # was input converted to Spectrum1D?
+    assert isinstance(collection[key], Spectrum1D), (f"image '{key}' not "
+                                                     "of type Spectrum1D")
+
+    # do key's fluxes match its comparison's fluxes?
+    assert np.allclose(collection[key].data,
+                       collection[compare].data), (f"images '{key}' and "
+                                                   f"'{compare}' have unequal "
+                                                   "flux values")
+
+    # if the image came with a spectral axis, was it kept? if not, was the
+    # default spectral axis in pixels applied?
+    sax_provided = hasattr(all_images[key], 'spectral_axis')
+    assert np.allclose(collection[key].spectral_axis,
+                       (all_images[key].spectral_axis if sax_provided
+                        else sax_def)), (f"spectral axis of image '{key}' does "
+                                         f"not match {'input' if sax_provided else 'default'}")
+
+    # if the image came with an uncertainty, was it kept? if not, was the
+    # default uncertainty created?
+    unc_provided = hasattr(all_images[key], 'uncertainty')
+    assert np.allclose(collection[key].uncertainty.array,
+                       (all_images[key].uncertainty.array if unc_provided
+                        else unc_def)), (f"uncertainty of image '{key}' does "
+                                         f"not match {'input' if unc_provided else 'default'}")
+
+    # were masks created despite none being given? (all indices should be False)
+    assert (getattr(collection[key], 'mask', None)
+            is not None), f"no mask was created for image '{key}'"
+    assert np.all(collection[key].mask == 0), ("mask not all False "
+                                               f"for image '{key}'")
+
+
+# test consistency of general image parser results
+def test_parse_general():
+    all_images_parsed = {k: FlatTrace._parse_image(object, im)
+                         for k, im in all_images.items()}
+
+    for key in all_images_parsed.keys():
+        compare_images(key, all_images_parsed)
+
+
+# use verified general image parser results to check HorneExtract's image parser
+def test_parse_horne():
+    # HorneExtract's parser is more stringent than the general one, hence the
+    # separate test. Given proper inputs, both should produce the same results.
+    images_collection = {k: {} for k in all_images.keys()}
+
+    for key, col in images_collection.items():
+        img = all_images[key]
+        col['general'] = FlatTrace._parse_image(object, img)
+
+        if hasattr(all_images[key], 'uncertainty'):
+            defaults = {}
+        else:
+            # save default values of attributes used in general parser when
+            # they are not available from the image object. HorneExtract always
+            # requires a variance, so it's chosen here to be on equal footing
+            # with the general case
+            defaults = {'variance': unc_def,
+                        'mask': np.ma.masked_invalid(img).mask,
+                        'unit': getattr(img, 'unit', u.DN)}
+
+        col[key] = HorneExtract._parse_image(object, img, **defaults)
+
+        compare_images(key, col, compare='general')

specreduce/tracing.py

@@ -5,11 +5,13 @@
 import warnings
 
 from astropy.modeling import Model, fitting, models
-from astropy.nddata import CCDData, NDData
+from astropy.nddata import NDData
 from astropy.stats import gaussian_sigma_to_fwhm
 from astropy.utils.decorators import deprecated
 import numpy as np
 
+from specreduce.core import _ImageParser
+
 __all__ = ['Trace', 'FlatTrace', 'ArrayTrace', 'FitTrace']
 
 
@@ -20,15 +22,15 @@ class Trace:
 
     Parameters
     ----------
-    image : `~astropy.nddata.CCDData`
+    image : `~astropy.nddata.NDData`-like or array-like, required
         Image to be traced
 
     Properties
     ----------
     shape : tuple
         Shape of the array describing the trace
     """
-    image: CCDData
+    image: NDData
 
     def __post_init__(self):
         self.trace_pos = self.image.shape[0] / 2
@@ -79,7 +81,7 @@ def __sub__(self, delta):
 
 
 @dataclass
-class FlatTrace(Trace):
+class FlatTrace(Trace, _ImageParser):
     """
     Trace that is constant along the axis being traced.
 
@@ -95,6 +97,8 @@ class FlatTrace(Trace):
     trace_pos: float
 
     def __post_init__(self):
+        self.image = self._parse_image(self.image)
+
         self.set_position(self.trace_pos)
 
     def set_position(self, trace_pos):
@@ -107,12 +111,12 @@ def set_position(self, trace_pos):
             Position of the trace
         """
         self.trace_pos = trace_pos
-        self.trace = np.ones_like(self.image[0]) * self.trace_pos
+        self.trace = np.ones_like(self.image.data[0]) * self.trace_pos
         self._bound_trace()
 
 
 @dataclass
-class ArrayTrace(Trace):
+class ArrayTrace(Trace, _ImageParser):
     """
     Define a trace given an array of trace positions.
 
@@ -124,6 +128,8 @@ class ArrayTrace(Trace):
     trace: np.ndarray
 
     def __post_init__(self):
+        self.image = self._parse_image(self.image)
+
         nx = self.image.shape[1]
         nt = len(self.trace)
         if nt != nx:
@@ -139,7 +145,7 @@ def __post_init__(self):
 
 
 @dataclass
-class FitTrace(Trace):
+class FitTrace(Trace, _ImageParser):
     """
     Trace the spectrum aperture in an image.
 
@@ -156,7 +162,7 @@ class FitTrace(Trace):
 
     Parameters
     ----------
-    image : `~astropy.nddata.NDData` or array-like, required
+    image : `~astropy.nddata.NDData`-like or array-like, required
         The image over which to run the trace. Assumes cross-dispersion
         (spatial) direction is axis 0 and dispersion (wavelength)
         direction is axis 1.
@@ -203,12 +209,12 @@ class FitTrace(Trace):
     _disp_axis = 1
 
     def __post_init__(self):
-        # handle multiple image types and mask uncaught invalid values
-        if isinstance(self.image, NDData):
-            img = np.ma.masked_invalid(np.ma.masked_array(self.image.data,
-                                                          mask=self.image.mask))
-        else:
-            img = np.ma.masked_invalid(self.image)
+        self.image = self._parse_image(self.image)
+
+        # mask any previously uncaught invalid values
+        or_mask = np.logical_or(self.image.mask,
+                                np.ma.masked_invalid(self.image.data).mask)
+        img = np.ma.masked_array(self.image.data, or_mask)
 
         # validate arguments
         valid_peak_methods = ('gaussian', 'centroid', 'max')