Consolidate image parsing in _ImageParser mixin

Author: ojustino

CHANGES.rst

@@ -7,6 +7,10 @@ New Features
 API Changes
 ^^^^^^^^^^^
 
+- All operations now accept Spectrum1D and Quantity-type images. All accepted
+image types are now processed internally as Spectrum1D objects. [#146]
+-
+
 Bug Fixes
 ^^^^^^^^^
 

specreduce/background.py

@@ -4,18 +4,19 @@
 from dataclasses import dataclass, field
 
 import numpy as np
-from astropy.nddata import NDData, VarianceUncertainty
+from astropy.nddata import NDData
 from astropy import units as u
 from specutils import Spectrum1D
 
+from specreduce.core import _ImageParser
 from specreduce.extract import _ap_weight_image, _to_spectrum1d_pixels
 from specreduce.tracing import Trace, FlatTrace
 
 __all__ = ['Background']
 
 
 @dataclass
-class Background:
+class Background(_ImageParser):
     """
     Determine the background from an image for subtraction.
 
@@ -55,41 +56,6 @@ 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.
@@ -122,7 +88,7 @@ def _to_trace(trace):
                     raise ValueError('trace_object.trace_pos must be >= 1')
             return trace
 
-        self._parse_image()
+        self.image = self._parse_image(self.image)
 
         if self.width < 0:
             raise ValueError("width must be positive")

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'))  # or u.Unit()?
+
+        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

@@ -9,7 +9,7 @@
 from astropy.modeling import Model, models, fitting
 from astropy.nddata import NDData, VarianceUncertainty
 
-from specreduce.core import SpecreduceOperation
+from specreduce.core import _ImageParser, SpecreduceOperation
 from specreduce.tracing import Trace, FlatTrace
 from specutils import Spectrum1D
 
@@ -164,41 +164,6 @@ 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):
         """
@@ -231,12 +196,7 @@ def __call__(self, image=None, trace_object=None, width=None,
         crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
 
         # handle image processing based on its type
-        if isinstance(image, Spectrum1D):
-            img = image.data
-            unit = image.unit
-        else:
-            img = image
-            unit = getattr(image, 'unit', u.DN)
+        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):
@@ -252,11 +212,11 @@ def __call__(self, image=None, trace_object=None, width=None,
             width,
             disp_axis,
             crossdisp_axis,
-            img.shape)
+            self.image.shape)
 
         # extract
-        ext1d = np.sum(img * wimg, axis=crossdisp_axis) * unit
-        return _to_spectrum1d_pixels(ext1d)
+        ext1d = np.sum(self.image.data * wimg, axis=crossdisp_axis)
+        return _to_spectrum1d_pixels(ext1d * self.image.unit)
 
 
 @dataclass

specreduce/tracing.py

@@ -12,6 +12,8 @@
 from specutils import Spectrum1D
 import numpy as np
 
+from specreduce.core import _ImageParser
+
 __all__ = ['Trace', 'FlatTrace', 'ArrayTrace', 'KosmosTrace']
 
 
@@ -39,41 +41,6 @@ def __post_init__(self):
     def __getitem__(self, i):
         return self.trace[i]
 
-    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)
-
     @property
     def shape(self):
         return self.trace.shape
@@ -116,7 +83,7 @@ def __sub__(self, delta):
 
 
 @dataclass
-class FlatTrace(Trace):
+class FlatTrace(Trace, _ImageParser):
     """
     Trace that is constant along the axis being traced
 
@@ -132,7 +99,7 @@ class FlatTrace(Trace):
     trace_pos: float
 
     def __post_init__(self):
-        super()._parse_image()
+        self.image = self._parse_image(self.image)
 
         self.set_position(self.trace_pos)
 
@@ -151,7 +118,7 @@ def set_position(self, trace_pos):
 
 
 @dataclass
-class ArrayTrace(Trace):
+class ArrayTrace(Trace, _ImageParser):
     """
     Define a trace given an array of trace positions
 
@@ -163,7 +130,7 @@ class ArrayTrace(Trace):
     trace: np.ndarray
 
     def __post_init__(self):
-        super()._parse_image()
+        self.image = self._parse_image(self.image)
 
         nx = self.image.shape[1]
         nt = len(self.trace)
@@ -180,7 +147,7 @@ def __post_init__(self):
 
 
 @dataclass
-class KosmosTrace(Trace):
+class KosmosTrace(Trace, _ImageParser):
     """
     Trace the spectrum aperture in an image.
 
@@ -241,7 +208,7 @@ class KosmosTrace(Trace):
     _disp_axis = 1
 
     def __post_init__(self):
-        super()._parse_image()
+        self.image = self._parse_image(self.image)
 
         # mask any previously uncaught invalid values
         or_mask = np.logical_or(self.image.mask,