Added image parsers to all classes that take images

Quantity arrays are now acceptable inputs. After parsing, the
attribute is now of type  instead of  as before.

specreduce/background.py

@@ -4,8 +4,9 @@
 from dataclasses import dataclass, field
 
 import numpy as np
-from astropy.nddata import NDData
+from astropy.nddata import NDData, VarianceUncertainty
 from astropy import units as u
+from specutils import Spectrum1D
 
 from specreduce.extract import _ap_weight_image, _to_spectrum1d_pixels
 from specreduce.tracing import Trace, FlatTrace
@@ -54,6 +55,41 @@ class Background:
     disp_axis: int = 1
     crossdisp_axis: int = 0
 
+    def _parse_image(self):
+        """
+        Convert all accepted image types to a consistently formatted Spectrum1D.
+        """
+
+        if isinstance(self.image, np.ndarray):
+            img = self.image
+        elif isinstance(self.image, u.quantity.Quantity):
+            img = self.image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = self.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(self.image, 'mask', None) is not None:
+            mask = self.image.mask
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        if getattr(self.image, 'uncertainty', None) is not None:
+            uncertainty = self.image.uncertainty
+        else:
+            uncertainty = VarianceUncertainty(np.ones(img.shape))
+
+        unit = getattr(self.image, 'unit', u.Unit('DN'))  # or u.Unit()?
+
+        spectral_axis = getattr(self.image, 'spectral_axis',
+                                (np.arange(img.shape[self.disp_axis])
+                                 if hasattr(self, 'disp_axis')
+                                 else np.arange(img.shape[1])) * u.pix)
+
+        self.image = Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                                uncertainty=uncertainty, mask=mask)
+
     def __post_init__(self):
         """
         Determine the background from an image for subtraction.
@@ -86,12 +122,9 @@ def _to_trace(trace):
                     raise ValueError('trace_object.trace_pos must be >= 1')
             return trace
 
-        if isinstance(self.image, NDData):
-            # NOTE: should the NDData structure instead be preserved?
-            # (NDData includes Spectrum1D under its umbrella)
-            self.image = self.image.data
+        self._parse_image()
 
-        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)
             if (np.any(trace.trace.data >= self.image.shape[self.crossdisp_axis]) or
@@ -116,9 +149,10 @@ 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=0)
+            self.bkg_array = np.average(self.image.data,
+                                        weights=self.bkg_wimage, axis=0)
         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=0)
         else:

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
@@ -149,6 +149,41 @@ class BoxcarExtract(SpecreduceOperation):
     def spectrum(self):
         return self.__call__()
 
+    def _parse_image(self):
+        """
+        Convert all accepted image types to a consistently formatted Spectrum1D.
+        """
+
+        if isinstance(self.image, np.ndarray):
+            img = self.image
+        elif isinstance(self.image, u.quantity.Quantity):
+            img = self.image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = self.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(self.image, 'mask', None) is not None:
+            mask = self.image.mask
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        if getattr(self.image, 'uncertainty', None) is not None:
+            uncertainty = self.image.uncertainty
+        else:
+            uncertainty = VarianceUncertainty(np.ones(img.shape))
+
+        unit = getattr(self.image, 'unit', u.Unit('DN'))  # or u.Unit()?
+
+        spectral_axis = getattr(self.image, 'spectral_axis',
+                                (np.arange(img.shape[self.disp_axis])
+                                 if hasattr(self, 'disp_axis')
+                                 else np.arange(img.shape[1])) * u.pix)
+
+        self.image = Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                                uncertainty=uncertainty, mask=mask)
+
     def __call__(self, image=None, trace_object=None, width=None,
                  disp_axis=None, crossdisp_axis=None):
         """
@@ -267,6 +302,86 @@ class HorneExtract(SpecreduceOperation):
     def spectrum(self):
         return self.__call__()
 
+    def _parse_image(self, variance=None, mask=None, unit=None):
+        """
+        Convert all accepted image types to a consistently formatted Spectrum1D.
+        Takes some extra arguments exactly as they come from self.__call__() to
+        handle cases where users specify them as arguments instead of as
+        attributes of their image object.
+        """
+
+        if isinstance(self.image, np.ndarray):
+            img = self.image
+        elif isinstance(self.image, u.quantity.Quantity):
+            img = self.image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = self.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(self.image, 'mask', None) is not None:
+            mask = self.image.mask
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        # Process uncertainties, converting to variances when able and throwing
+        # an error when uncertainties are missing or less easily converted
+        if (hasattr(self.image, 'uncertainty')
+                and self.image.uncertainty is not None):
+            if self.image.uncertainty.uncertainty_type == 'var':
+                variance = self.image.uncertainty.array
+            elif self.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 = self.image.uncertainty.array**2
+            elif self.image.uncertainty.uncertainty_type == 'ivar':
+                variance = 1 / self.image.uncertainty.array
+            else:
+                # other options are InverseVariance and UnknownVariance
+                raise ValueError("image NDData object has unexpected "
+                                 "uncertainty type. instead, try "
+                                 "VarianceUncertainty or StdDevUncertainty.")
+        elif (hasattr(self.image, 'uncertainty')
+              and self.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 self.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(self.image, 'unit',
+                       u.Unit(self.unit) if self.unit is not None else u.Unit())
+
+        spectral_axis = getattr(self.image, 'spectral_axis',
+                                (np.arange(img.shape[self.disp_axis])
+                                 if hasattr(self, 'disp_axis')
+                                 else np.arange(img.shape[1])) * u.pix)
+
+        self.image = 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,
@@ -329,71 +444,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):
-            # (NDData includes Spectrum1D under its umbrella)
-            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')
-
-        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)
+        # parse image and replace optional arguments with updated values
+        self._parse_image(variance, mask, unit)
+        variance = self.image.uncertainty.array
+        unit = self.image.unit
 
-        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]

specreduce/tracing.py

@@ -5,8 +5,9 @@
 import warnings
 
 from astropy.modeling import fitting, models
-from astropy.nddata import NDData
+from astropy.nddata import NDData, VarianceUncertainty
 from astropy.stats import gaussian_sigma_to_fwhm
+from astropy import units as u
 from scipy.interpolate import UnivariateSpline
 from specutils import Spectrum1D
 import numpy as np
@@ -39,9 +40,39 @@ def __getitem__(self, i):
         return self.trace[i]
 
     def _parse_image(self):
-        if isinstance(self.image, Spectrum1D):
-            # NOTE: should the Spectrum1D structure instead be preserved?
-            self.image = self.image.data
+        """
+        Convert all accepted image types to a consistently formatted Spectrum1D.
+        """
+
+        if isinstance(self.image, np.ndarray):
+            img = self.image
+        elif isinstance(self.image, u.quantity.Quantity):
+            img = self.image.value
+        else:  # NDData, including CCDData and Spectrum1D
+            img = self.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(self.image, 'mask', None) is not None:
+            mask = self.image.mask
+        else:
+            mask = np.ma.masked_invalid(img).mask
+
+        if getattr(self.image, 'uncertainty', None) is not None:
+            uncertainty = self.image.uncertainty
+        else:
+            uncertainty = VarianceUncertainty(np.ones(img.shape))
+
+        unit = getattr(self.image, 'unit', u.Unit('DN'))  # or u.Unit()?
+
+        spectral_axis = getattr(self.image, 'spectral_axis',
+                                (np.arange(img.shape[self._disp_axis])
+                                 if hasattr(self, '_disp_axis')
+                                 else np.arange(img.shape[1])) * u.pix)
+
+        self.image = Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                                uncertainty=uncertainty, mask=mask)
 
     @property
     def shape(self):
@@ -115,7 +146,7 @@ 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()
 
 
@@ -212,12 +243,10 @@ class KosmosTrace(Trace):
     def __post_init__(self):
         super()._parse_image()
 
-        # 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)
+        # 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')