Fixed Spectrum1D parsing in Horne and its tests

Author: ojustino

specreduce/extract.py

@@ -216,7 +216,8 @@ def __call__(self, image=None, trace_object=None, width=None,
 
         # extract
         ext1d = np.sum(self.image.data * wimg, axis=crossdisp_axis)
-        return _to_spectrum1d_pixels(ext1d * self.image.unit)
+        return Spectrum1D(ext1d * self.image.unit,
+                          spectral_axis=self.image.spectral_axis)
 
 
 @dataclass
@@ -280,50 +281,86 @@ class HorneExtract(SpecreduceOperation):
     def spectrum(self):
         return self.__call__()
 
-    def _parse_image(self, variance=None, mask=None, unit=None):
+    def _parse_image(self, image,
+                     variance=None, mask=None, unit=None, disp_axis=1):
         """
-        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.
+        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 as unitless.
+        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(self.image, np.ndarray):
-            img = self.image
-        elif isinstance(self.image, u.quantity.Quantity):
-            img = self.image.value
+        if isinstance(image, np.ndarray):
+            img = image
+        elif isinstance(image, u.quantity.Quantity):
+            img = image.value
         else:  # NDData, including CCDData and Spectrum1D
-            img = self.image.data
+            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(self.image, 'mask', None) is not None:
-            mask = self.image.mask
+        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(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':
+        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 = self.image.uncertainty.array**2
-            elif self.image.uncertainty.uncertainty_type == 'ivar':
-                variance = 1 / self.image.uncertainty.array
+                variance = image.uncertainty.array**2
+            elif image.uncertainty.uncertainty_type == 'ivar':
+                variance = 1 / image.uncertainty.array
             else:
-                # other options are InverseVariance and UnknownVariance
+                # other options are InverseUncertainty and UnknownUncertainty
                 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):
+        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:
@@ -332,7 +369,7 @@ def _parse_image(self, variance=None, mask=None, unit=None):
                                  "variance must be specified. consider "
                                  "wrapping it into one object by instead "
                                  "passing an NDData image.")
-            elif self.image.shape != variance.shape:
+            elif image.shape != variance.shape:
                 raise ValueError("image and variance shapes must match")
 
         if np.any(variance < 0):
@@ -349,16 +386,14 @@ def _parse_image(self, variance=None, mask=None, unit=None):
 
         variance = VarianceUncertainty(variance)
 
-        unit = getattr(self.image, 'unit',
-                       u.Unit(self.unit) if self.unit is not None else u.Unit())
+        unit = getattr(image, 'unit',
+                       u.Unit(unit) if 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)
+        spectral_axis = getattr(image, 'spectral_axis',
+                                np.arange(img.shape[disp_axis]) * u.pix)
 
-        self.image = Spectrum1D(img * unit, spectral_axis=spectral_axis,
-                                uncertainty=variance, mask=mask)
+        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,
@@ -423,7 +458,7 @@ def __call__(self, image=None, trace_object=None,
         unit = unit if unit is not None else self.unit
 
         # parse image and replace optional arguments with updated values
-        self._parse_image(variance, mask, unit)
+        self.image = self._parse_image(image, variance, mask, unit, disp_axis)
         variance = self.image.uncertainty.array
         unit = self.image.unit
 

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,65 @@ 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 unitless flux and spectral axis in pixels
     err = np.ones_like(image)
-    ext = extract(variance=err)
+    ext = extract(image=image.data, variance=err, mask=None, unit=None)
     assert ext.unit == u.Unit()
+    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
+    err = image.uncertainty.array
+    err[0][0] = -1
     mask = np.zeros_like(image)
     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)