Merge pull request #144 from ojustino/accept-spectrum1d

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.