implement simple boxcar Background subtraction

* requires exposing some internal methods from Boxcar extract to re-use here
* adds and uses +/- offset support to traces
* updates the notebook with a basic use-case

specreduce/background.py

@@ -0,0 +1,136 @@
+# Licensed under a 3-clause BSD style license - see LICENSE.rst
+
+from dataclasses import dataclass
+
+import numpy as np
+from astropy.nddata import CCDData
+
+from specreduce.core import SpecreduceOperation
+from specreduce.extract import _ap_weight_image
+from specreduce.tracing import Trace, FlatTrace
+
+__all__ = ['Background']
+
+
+@dataclass
+class Background(SpecreduceOperation):
+    """
+    Determine the background from an image for subtraction
+
+    Parameters
+    ----------
+    image : nddata-compatible image
+        image with 2-D spectral image data
+    trace_object : Trace
+        trace object
+    width : float
+        width of extraction aperture in pixels
+    disp_axis : int
+        dispersion axis
+    crossdisp_axis : int
+        cross-dispersion axis
+
+    Returns
+    -------
+    spec : `~specutils.Spectrum1D`
+        The extracted 1d spectrum expressed in DN and pixel units
+    """
+    # required so numpy won't call __rsub__ on individual elements
+    # https://stackoverflow.com/a/58409215
+    __array_ufunc__ = None
+
+    image: CCDData
+    trace_object: Trace
+    separation: float = 5
+    width: float = 5
+    disp_axis: int = 1
+    crossdisp_axis: int = 0
+
+    def __post_init__(self):
+        """
+        Extract the 1D spectrum using the boxcar method.
+
+        Parameters
+        ----------
+        image : nddata-compatible image
+            image with 2-D spectral image data
+        trace_object : Trace or int
+            trace object or an integer to use a FloatTrace
+        separation: float
+            separation between trace and extraction apertures on each
+            side of the trace
+        width : float
+            width of each background aperture in pixels
+        disp_axis : int
+            dispersion axis
+        crossdisp_axis : int
+            cross-dispersion axis
+        """
+        if isinstance(self.trace_object, (int, float)):
+            self.trace_object = FlatTrace(self.image, self.trace_object)
+
+        # TODO: this check can be removed if/when implemented as a check in FlatTrace
+        if isinstance(self.trace_object, FlatTrace):
+            if self.trace_object.trace_pos < 1:
+                raise ValueError('trace_object.trace_pos must be >= 1')
+
+        bkg_wimage = _ap_weight_image(
+            self.trace_object-self.separation,
+            self.width,
+            self.disp_axis,
+            self.crossdisp_axis,
+            self.image.shape)
+
+        bkg_wimage += _ap_weight_image(
+            self.trace_object+self.separation,
+            self.width,
+            self.disp_axis,
+            self.crossdisp_axis,
+            self.image.shape)
+
+        self.bkg_wimage = bkg_wimage
+        self.bkg_array = np.average(self.image, weights=self.bkg_wimage, axis=0)
+
+    def bkg_image(self, image=None):
+        """
+        Expose the background tiled to the dimension of ``image``.
+
+        Parameters
+        ----------
+        image : nddata-compatible image or None
+            image with 2-D spectral image data.  If None, will use ``image`` passed
+            to extract the background.
+
+        Returns
+        -------
+        array with same shape as ``image``.
+        """
+        if image is None:
+            image = self.image
+
+        return np.tile(self.bkg_array, (image.shape[0], 1))
+
+    def sub_image(self, image=None):
+        """
+        Subtract the computed background from ``image``.
+
+        Parameters
+        ----------
+        image : nddata-compatible image or None
+            image with 2-D spectral image data.  If None, will use ``image`` passed
+            to extract the background.
+
+        Returns
+        -------
+        array with same shape as ``image``
+        """
+        if image is None:
+            image = self.image
+
+        return image - self.bkg_image(image)
+
+    def __rsub__(self, image):
+        """
+        Subtract the background from an image.
+        """
+        return self.sub_image(image)

specreduce/extract.py

@@ -13,6 +13,66 @@
 __all__ = ['BoxcarExtract']
 
 
+def _get_boxcar_weights(center, hwidth, npix):
+    """
+    Compute weights given an aperture center, half width,
+    and number of pixels
+    """
+    weights = np.zeros((npix))
+
+    # pixels with full weight
+    fullpixels = [max(0, int(center - hwidth + 1)),
+                  min(int(center + hwidth), npix)]
+    weights[fullpixels[0]:fullpixels[1]] = 1.0
+
+    # pixels at the edges of the boxcar with partial weight
+    if fullpixels[0] > 0:
+        w = hwidth - (center - fullpixels[0] + 0.5)
+        if w >= 0:
+            weights[fullpixels[0] - 1] = w
+        else:
+            weights[fullpixels[0]] = 1. + w
+    if fullpixels[1] < npix:
+        weights[fullpixels[1]] = hwidth - (fullpixels[1] - center - 0.5)
+
+    return weights
+
+
+def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
+
+    """
+    Create a weight image that defines the desired extraction aperture.
+
+    Parameters
+    ----------
+    trace : Trace
+        trace object
+    width : float
+        width of extraction aperture in pixels
+    disp_axis : int
+        dispersion axis
+    crossdisp_axis : int
+        cross-dispersion axis
+    image_shape : tuple with 2 elements
+        size (shape) of image
+
+    Returns
+    -------
+    wimage : 2D image
+        weight image defining the aperture
+    """
+    wimage = np.zeros(image_shape)
+    hwidth = 0.5 * width
+    image_sizes = image_shape[crossdisp_axis]
+
+    # loop in dispersion direction and compute weights.
+    for i in range(image_shape[disp_axis]):
+        # TODO trace must handle transposed data (disp_axis == 0)
+        wimage[:, i] = _get_boxcar_weights(trace[i], hwidth, image_sizes)
+
+    return wimage
+
+
 @dataclass
 class BoxcarExtract(SpecreduceOperation):
     """
@@ -64,64 +124,6 @@ def __call__(self, image, trace_object, width=5,
             The extracted 1d spectrum with flux expressed in the same
             units as the input image, or u.DN, and pixel units
         """
-        def _get_boxcar_weights(center, hwidth, npix):
-            """
-            Compute weights given an aperture center, half width,
-            and number of pixels
-            """
-            weights = np.zeros((npix))
-
-            # pixels with full weight
-            fullpixels = [max(0, int(center - hwidth + 1)),
-                          min(int(center + hwidth), npix)]
-            weights[fullpixels[0]:fullpixels[1]] = 1.0
-
-            # pixels at the edges of the boxcar with partial weight
-            if fullpixels[0] > 0:
-                w = hwidth - (center - fullpixels[0] + 0.5)
-                if w >= 0:
-                    weights[fullpixels[0] - 1] = w
-                else:
-                    weights[fullpixels[0]] = 1. + w
-            if fullpixels[1] < npix:
-                weights[fullpixels[1]] = hwidth - (fullpixels[1] - center - 0.5)
-
-            return weights
-
-        def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
-
-            """
-            Create a weight image that defines the desired extraction aperture.
-
-            Parameters
-            ----------
-            trace : Trace
-                trace object
-            width : float
-                width of extraction aperture in pixels
-            disp_axis : int
-                dispersion axis
-            crossdisp_axis : int
-                cross-dispersion axis
-            image_shape : tuple with 2 elements
-                size (shape) of image
-
-            Returns
-            -------
-            wimage : 2D image
-                weight image defining the aperture
-            """
-            wimage = np.zeros(image_shape)
-            hwidth = 0.5 * width
-            image_sizes = image_shape[crossdisp_axis]
-
-            # loop in dispersion direction and compute weights.
-            for i in range(image_shape[disp_axis]):
-                # TODO trace must handle transposed data (disp_axis == 0)
-                wimage[:, i] = _get_boxcar_weights(trace[i], hwidth, image_sizes)
-
-            return wimage
-
         # 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:

specreduce/tracing.py

@@ -1,5 +1,6 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
+from copy import deepcopy
 from dataclasses import dataclass
 
 import numpy as np
@@ -45,6 +46,8 @@ def shift(self, delta):
         delta : float
             Shift to be applied to the trace
         """
+        if not isinstance(delta, (int, float)):
+            raise TypeError(f"{self.__class__.__name__} only supports shifting by floats/integers")
         self.trace += delta
         self._bound_trace()
 
@@ -55,6 +58,20 @@ def _bound_trace(self):
         ny = self.image.shape[0]
         self.trace = np.ma.masked_outside(self.trace, 0, ny-1)
 
+    def __add__(self, delta):
+        """
+        Return a copy of the trace shifted by delta pixels perpendicular to the axis being traced
+        """
+        copy = deepcopy(self)
+        copy.shift(delta)
+        return copy
+
+    def __sub__(self, delta):
+        """
+        Return a copy of the trace shifted by delta pixels perpendicular to the axis being traced
+        """
+        return self.__add__(-delta)
+
 
 @dataclass
 class FlatTrace(Trace):