From 12eb7051d5ebf46677ad8d328116ccb613a2c48d Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Tue, 20 Feb 2024 18:38:59 -0500
Subject: [PATCH 01/25] fix input for background.traces, raise error in
 FlatTrace for negative trace

.

.

.

..

,

.

.

code style

.
---
 CHANGES.rst                         |   7 +-
 specreduce/background.py            |  90 ++++++-
 specreduce/core.py                  | 120 ++++++++-
 specreduce/extract.py               |  44 +++-
 specreduce/tests/test_background.py | 193 +++++++++++++-
 specreduce/tests/test_extract.py    |  71 ++++-
 specreduce/tests/test_tracing.py    | 394 +++++++++++++++++++++++-----
 specreduce/tracing.py               | 107 +++++++-
 8 files changed, 902 insertions(+), 124 deletions(-)

diff --git a/CHANGES.rst b/CHANGES.rst
index e7974c8..66c11f3 100644
--- a/CHANGES.rst
+++ b/CHANGES.rst
@@ -70,7 +70,12 @@ Bug Fixes
   peaks. Previously for Gaussian, the entire fit failed. [#205, #206]
 
 - Fixed input of `traces` in `Background`. Added a condition to 'FlatTrace' that
-  trace position must be a positive number. [#211]
+
+- Fix in FitTrace to set fully-masked column bin peaks to NaN. Previously, for
+  peak_method='max' these were set to 0.0, and for peak_method='centroid' they
+  were set to the number of rows in the image, biasing the final fit to all bin
+  peaks. [#206]
+
 
 Other changes
 ^^^^^^^^^^^^^
diff --git a/specreduce/background.py b/specreduce/background.py
index 5b8c6b6..29c81ca 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -32,9 +32,9 @@ class Background(_ImageParser):
     ----------
     image : `~astropy.nddata.NDData`-like or array-like
         image with 2-D spectral image data
-    traces : trace, int, float (single or list)
+    traces : List, `tracing.Trace`, int, float
         Individual or list of trace object(s) (or integers/floats to define
-        FlatTraces) to extract the background. If None, a FlatTrace at the
+        FlatTraces) to extract the background. If None, a `FlatTrace` at the
         center of the image (according to `disp_axis`) will be used.
     width : float
         width of extraction aperture in pixels
@@ -44,8 +44,22 @@ class Background(_ImageParser):
         pixels.
     disp_axis : int
         dispersion axis
+        [default: 1]
     crossdisp_axis : int
         cross-dispersion axis
+        [default: 0]
+    mask_treatment : string, optional
+        The method for handling masked or non-finite data. Choice of `filter`,
+        `omit`, or `zero-fill`. If `filter` is chosen, masked and non-finite
+        data will not contribute to the background statistic that is calculated
+        in each column along `disp_axis`. If `omit` is chosen, columns along
+        disp_axis with any masked/non-finite data values will be fully masked
+        (i.e, 2D mask is collapsed to 1D and applied). If `zero-fill` is chosen,
+        masked/non-finite data will be replaced with 0.0 in the input image,
+        and the mask will then be dropped. For all three options, the input mask
+        (optional on input NDData object) will be combined with a mask generated
+        from any non-finite values in the image data.
+        [default: ``filter``]
     """
     # required so numpy won't call __rsub__ on individual elements
     # https://stackoverflow.com/a/58409215
@@ -57,6 +71,8 @@ class Background(_ImageParser):
     statistic: str = 'average'
     disp_axis: int = 1
     crossdisp_axis: int = 0
+    mask_treatment: str = 'filter'
+    _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
 
     # TO-DO: update bkg_array with Spectrum1D alternative (is bkg_image enough?)
     bkg_array = deprecated_attribute('bkg_array', '1.3')
@@ -82,9 +98,29 @@ def __post_init__(self):
             dispersion axis
         crossdisp_axis : int
             cross-dispersion axis
+        mask_treatment : string
+            The method for handling masked or non-finite data. Choice of `filter`,
+            `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+            will be filtered during the fit to each bin/column (along disp. axis) to
+            find the peak. If `omit` is chosen, columns along disp_axis with any
+            masked/non-finite data values will be fully masked (i.e, 2D mask is
+            collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+            data will be replaced with 0.0 in the input image, and the mask will then
+            be dropped. For all three options, the input mask (optional on input
+            NDData object) will be combined with a mask generated from any non-finite
+            values in the image data.
         """
+
+        # Parse image, including masked/nonfinite data handling based on
+        # choice of `mask_treatment`. Any uncaught nonfinte data values will be
+        # masked as well. Returns a Spectrum1D.
         self.image = self._parse_image(self.image)
 
+        # always work with masked array, even if there is no masked
+        # or nonfinite data, in case padding is needed. if not, mask will be
+        # dropped at the end and a regular array will be returned.
+        img = np.ma.masked_array(self.image.data, self.image.mask)
+
         if self.width < 0:
             raise ValueError("width must be positive")
         if self.width == 0:
@@ -95,9 +131,13 @@ def __post_init__(self):
 
         bkg_wimage = np.zeros_like(self.image.data, dtype=np.float64)
         for trace in self.traces:
+            # note: ArrayTrace can have masked values, but if it does a MaskedArray
+            # will be returned so this should be reflected in the window size here
+            # (i.e, np.nanmax is not required.)
             windows_max = trace.trace.data.max() + self.width/2
             windows_min = trace.trace.data.min() - self.width/2
-            if windows_max >= self.image.shape[self.crossdisp_axis]:
+
+            if windows_max > self.image.shape[self.crossdisp_axis]:
                 warnings.warn("background window extends beyond image boundaries " +
                               f"({windows_max} >= {self.image.shape[self.crossdisp_axis]})")
             if windows_min < 0:
@@ -115,6 +155,9 @@ def __post_init__(self):
             raise ValueError("background regions overlapped")
         if np.any(np.sum(bkg_wimage, axis=self.crossdisp_axis) == 0):
             raise ValueError("background window does not remain in bounds across entire dispersion axis")  # noqa
+        # check if image contained within background window is fully-nonfinite and raise an error
+        if np.all(img.mask[bkg_wimage > 0]):
+            raise ValueError("Image is fully masked within background window determined by `width`.")  # noqa
 
         if self.statistic == 'median':
             # make it clear in the expose image that partial pixels are fully-weighted
@@ -122,20 +165,16 @@ def __post_init__(self):
 
         self.bkg_wimage = bkg_wimage
 
-        # mask user-highlighted and invalid values (if any) before taking stats
-        or_mask = (np.logical_or(~np.isfinite(self.image.data), self.image.mask)
-                   if self.image.mask is not None
-                   else ~np.isfinite(self.image.data))
-
         if self.statistic == 'average':
-            image_ma = np.ma.masked_array(self.image.data, mask=or_mask)
-            self._bkg_array = np.ma.average(image_ma,
+            self._bkg_array = np.ma.average(img,
                                             weights=self.bkg_wimage,
-                                            axis=self.crossdisp_axis).data
+                                            axis=self.crossdisp_axis)
+
         elif self.statistic == 'median':
-            med_mask = np.logical_or(self.bkg_wimage == 0, or_mask)
-            image_ma = np.ma.masked_array(self.image.data, mask=med_mask)
-            self._bkg_array = np.ma.median(image_ma, axis=self.crossdisp_axis).data
+            # combine where background weight image is 0 with image masked (which already
+            # accounts for non-finite data that wasn't already masked)
+            img.mask = np.logical_or(self.bkg_wimage == 0, self.image.mask)
+            self._bkg_array = np.ma.median(img, axis=self.crossdisp_axis)
         else:
             raise ValueError("statistic must be 'average' or 'median'")
 
@@ -204,7 +243,19 @@ def two_sided(cls, image, trace_object, separation, **kwargs):
             dispersion axis
         crossdisp_axis : int
             cross-dispersion axis
+        mask_treatment : string
+            The method for handling masked or non-finite data. Choice of `filter`,
+            `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+            will be filtered during the fit to each bin/column (along disp. axis) to
+            find the peak. If `omit` is chosen, columns along disp_axis with any
+            masked/non-finite data values will be fully masked (i.e, 2D mask is
+            collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+            data will be replaced with 0.0 in the input image, and the mask will then
+            be dropped. For all three options, the input mask (optional on input
+            NDData object) will be combined with a mask generated from any non-finite
+            values in the image data.
         """
+
         image = _ImageParser._get_data_from_image(image) if image is not None else cls.image
         kwargs['traces'] = [trace_object-separation, trace_object+separation]
         return cls(image=image, **kwargs)
@@ -241,6 +292,17 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
             dispersion axis
         crossdisp_axis : int
             cross-dispersion axis
+        mask_treatment : string
+            The method for handling masked or non-finite data. Choice of `filter`,
+            `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+            will be filtered during the fit to each bin/column (along disp. axis) to
+            find the peak. If `omit` is chosen, columns along disp_axis with any
+            masked/non-finite data values will be fully masked (i.e, 2D mask is
+            collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+            data will be replaced with 0.0 in the input image, and the mask will then
+            be dropped. For all three options, the input mask (optional on input
+            NDData object) will be combined with a mask generated from any non-finite
+            values in the image data.
         """
         image = _ImageParser._get_data_from_image(image) if image is not None else cls.image
         kwargs['traces'] = [trace_object+separation]
diff --git a/specreduce/core.py b/specreduce/core.py
index 1737d0f..4fc6ded 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -1,5 +1,6 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
+from copy import deepcopy
 import inspect
 from dataclasses import dataclass
 
@@ -67,14 +68,20 @@ def _get_data_from_image(image, disp_axis=1):
         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.isfinite(img)
+        mask = getattr(image, 'mask', None)
+
+        # next, handle masked and nonfinite data in image.
+        # A mask will be created from any nonfinite image data, and combined
+        # with any additional 'mask' passed in. If image is being parsed within
+        # a specreduce operation that has 'mask_treatment' options, this will be
+        # handled as well. Note that input data may be modified if a fill value
+        # is chosen to handle masked data. The returned image will always have
+        # `image.mask` even if there are no nonfinte or masked values.
+        img, mask = self._mask_and_nonfinite_data_handling(image=img, mask=mask)
 
+        # mask (handled above) 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, 'uncertainty', None) is not None:
             uncertainty = image.uncertainty
         else:
@@ -85,11 +92,106 @@ def _get_data_from_image(image, disp_axis=1):
         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)
+        img = Spectrum1D(img * unit, spectral_axis=spectral_axis,
+                         uncertainty=uncertainty, mask=mask)
+
+        return img
+
+    @staticmethod
+    def _get_data_from_image(image):
+        """Extract data array from various input types for `image`.
+           Retruns `np.ndarray` of image data."""
 
+        if isinstance(image, u.quantity.Quantity):
+            img = image.value
+        if isinstance(image, np.ndarray):
+            img = image
+        else:  # NDData, including CCDData and Spectrum1D
+            img = image.data
         return img
 
+    def _mask_and_nonfinite_data_handling(self, image, mask):
+        """
+        This function handles the treatment of masked and nonfinite data,
+        including input validation.
+
+        All operations in Specreduce can take in a mask for the data as
+        part of the input NDData. Additionally, any non-finite values in the
+        data that aren't in the user-supplied mask will be combined bitwise
+        with the input mask.
+
+        There are three options currently implemented for the treatment
+        of masked and nonfinite data - filter, omit, and zero-fill.
+        Depending on the step, all or a subset of these three options are valid.
+
+        """
+
+        # valid options depend on Specreduce step, and are set accordingly there
+        # for steps that this isn't implemeted for yet, default to 'filter',
+        # which will return unmodified input image and mask
+        mask_treatment = getattr(self, 'mask_treatment', 'filter')
+
+        # make sure chosen option is valid. if _valid_mask_treatment_methods
+        # is not an attribue, proceed with 'filter' to return back inupt data
+        # and mask that is combined with nonfinite data.
+        if mask_treatment is not None:  # None in operations where masks aren't relevant (FlatTrace)
+            valid_mask_treatment_methods = getattr(self, '_valid_mask_treatment_methods', ['filter'])  # noqa
+            if mask_treatment not in valid_mask_treatment_methods:
+                raise ValueError(f"`mask_treatment` must be one of {valid_mask_treatment_methods}")
+
+        # make sure there is always a 'mask', even when all data is unmasked and finite.
+        if mask is not None:
+            mask = self.image.mask
+            # always mask any previously uncaught nonfinite values in image array
+            # combining these with the (optional) user-provided mask on `image.mask`
+            mask = np.logical_or(mask, ~np.isfinite(image))
+        else:
+            mask = ~np.isfinite(image)
+
+        # if mask option is the default 'filter' option, or None,
+        # nothing needs to be done. input mask (combined with nonfinite data)
+        # remains with data as-is.
+
+        if mask_treatment == 'zero-fill':
+            # make a copy of the input image since we will be modifying it
+            image = deepcopy(image)
+
+            # if mask_treatment is 'zero_fill', set masked values to zero in
+            # image data and drop image.mask. note that this is done after
+            # _combine_mask_with_nonfinite_from_data, so non-finite values in
+            # data (which are now in the mask) will also be set to zero.
+            # set masked values to zero
+            image[mask] = 0.
+
+            # masked array with no masked values, so accessing image.mask works
+            # but we don't need the actual mask anymore since data has been set to 0
+            mask = np.zeros(image.shape)
+
+        elif mask_treatment == 'omit':
+            # collapse 2d mask (after accounting for addl non-finite values in
+            # data) to a 1d mask, along dispersion axis, to fully mask columns
+            # that have any masked values.
+
+            # must have a crossdisp_axis specified to use 'omit' optoin
+            if hasattr(self, 'crossdisp_axis'):
+                crossdisp_axis = self.crossdisp_axis
+            if hasattr(self, '_crossdisp_axis'):
+                crossdisp_axis = self._crossdisp_axis
+
+            # create a 1d mask along crossdisp axis - if any column has a single nan,
+            # the entire column should be masked
+            reduced_mask = np.logical_or.reduce(mask,
+                                                axis=crossdisp_axis)
+
+            # back to a 2D mask
+            shape = (image.shape[0], 1) if crossdisp_axis == 0 else (1, image.shape[1])
+            mask = np.tile(reduced_mask, shape)
+
+        # check for case where entire image is masked.
+        if mask.all():
+            raise ValueError('Image is fully masked. Check for invalid values.')
+
+        return image, mask
 
 @dataclass
 class SpecreduceOperation(_ImageParser):
diff --git a/specreduce/extract.py b/specreduce/extract.py
index 4b41ca2..96128af 100644
--- a/specreduce/extract.py
+++ b/specreduce/extract.py
@@ -113,6 +113,10 @@ def _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image_shape):
     for i in range(image_shape[disp_axis]):
         # TODO trace must handle transposed data (disp_axis == 0)
         # pass trace.trace.data[i] to avoid any mask if part of the regions is out-of-bounds
+
+        # ArrayTrace can have nonfinite or masked data in trace, and this will fail,
+        # so figure out how to handle that...
+
         wimage[:, i] = _get_boxcar_weights(trace.trace.data[i], hwidth, image_sizes)
 
     return wimage
@@ -154,6 +158,8 @@ class BoxcarExtract(SpecreduceOperation):
     disp_axis: int = 1
     crossdisp_axis: int = 0
     # TODO: should disp_axis and crossdisp_axis be defined in the Trace object?
+    mask_treatment: str = 'filter'
+    _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
 
     @property
     def spectrum(self):
@@ -176,6 +182,20 @@ def __call__(self, image=None, trace_object=None, width=None,
             dispersion axis [default: 1]
         crossdisp_axis : int, optional
             cross-dispersion axis [default: 0]
+        mask_treatment : string, optional
+            The method for handling masked or non-finite data. Choice of `filter`,
+            `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+            will be filtered during the fit to each bin/column (along disp. axis) to
+            find the peak. If `omit` is chosen, columns along disp_axis with any
+            masked/non-finite data values will be fully masked (i.e, 2D mask is
+            collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+            data will be replaced with 0.0 in the input image, and the mask will then
+            be dropped. For all three options, the input mask (optional on input
+            NDData object) will be combined with a mask generated from any non-finite
+            values in the image data. Also note that because binning is an option in
+            FitTrace, that masked data will contribute zero to the sum when binning
+            adjacent columns.
+            [default: ``filter``]
 
 
         Returns
@@ -190,24 +210,24 @@ def __call__(self, image=None, trace_object=None, width=None,
         disp_axis = disp_axis if disp_axis is not None else self.disp_axis
         crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
 
-        # handle image processing based on its type
+        # Parse image, including masked/nonfinite data handling based on
+        # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
+        # omit. non-finite data will be masked, always. Returns a Spectrum1D.
         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):
-            if trace_object.trace_pos < 1:
-                raise ValueError('trace_object.trace_pos must be >= 1')
+        # # _parse_image returns a Spectrum1D. convert this to a masked array
+        # # for ease of calculations here (even if there is no masked data).
+        # img = np.ma.masked_array(self.image.data, self.image.mask)
 
-        if width < 0:
+        if width <= 0:
             raise ValueError("width must be positive")
 
         # weight image to use for extraction
-        wimg = _ap_weight_image(
-            trace_object,
-            width,
-            disp_axis,
-            crossdisp_axis,
-            self.image.shape)
+        wimg = _ap_weight_image(trace_object,
+                                width,
+                                disp_axis,
+                                crossdisp_axis,
+                                self.image.shape)
 
         # extract, assigning no weight to non-finite pixels outside the window
         # (non-finite pixels inside the window will still make it into the sum)
diff --git a/specreduce/tests/test_background.py b/specreduce/tests/test_background.py
index 09364c3..388da5c 100644
--- a/specreduce/tests/test_background.py
+++ b/specreduce/tests/test_background.py
@@ -1,3 +1,5 @@
+from astropy.nddata import NDData
+import astropy.units as u
 import numpy as np
 import pytest
 from specutils import Spectrum1D
@@ -22,13 +24,13 @@ def test_background(mk_test_img_raw, mk_test_spec_no_spectral_axis,
 
     # all the following should be equivalent, whether image's spectral axis
     # is in pixels or physical units:
-    bg1 = Background(image, [trace-bkg_sep, trace+bkg_sep], width=bkg_width)
+    bg1 = Background(image, [trace - bkg_sep, trace + bkg_sep], width=bkg_width)
     bg2 = Background.two_sided(image, trace, bkg_sep, width=bkg_width)
     bg3 = Background.two_sided(image, trace_pos, bkg_sep, width=bkg_width)
     assert np.allclose(bg1.bkg_image().flux, bg2.bkg_image().flux)
     assert np.allclose(bg1.bkg_image().flux, bg3.bkg_image().flux)
 
-    bg4 = Background(image_um, [trace-bkg_sep, trace+bkg_sep], width=bkg_width)
+    bg4 = Background(image_um, [trace - bkg_sep, trace + bkg_sep], width=bkg_width)
     bg5 = Background.two_sided(image_um, trace, bkg_sep, width=bkg_width)
     bg6 = Background.two_sided(image_um, trace_pos, bkg_sep, width=bkg_width)
     assert np.allclose(bg1.bkg_image().flux, bg4.bkg_image().flux)
@@ -72,7 +74,7 @@ def test_background(mk_test_img_raw, mk_test_spec_no_spectral_axis,
     stats = ['average', 'median']
 
     for st in stats:
-        bg = Background(img, trace-bkg_sep, width=bkg_width, statistic=st)
+        bg = Background(img, trace - bkg_sep, width=bkg_width, statistic=st)
         assert np.isnan(bg.image.flux).sum() == 2
         assert np.isnan(bg._bkg_array).sum() == 0
         assert np.isnan(bg.bkg_spectrum().flux).sum() == 0
@@ -99,7 +101,7 @@ def test_warnings_errors(mk_test_spec_no_spectral_axis):
     with pytest.warns(match="background window extends beyond image boundaries"):
         Background.two_sided(image, 7, 5, width=6)
 
-    trace = ArrayTrace(image, trace=np.arange(10)+20)  # from 20 to 29
+    trace = ArrayTrace(image, trace=np.arange(10) + 20)  # from 20 to 29
     with pytest.warns(match="background window extends beyond image boundaries"):
         with pytest.raises(ValueError,
                            match="background window does not remain in bounds across entire dispersion axis"):  # noqa
@@ -112,6 +114,11 @@ def test_warnings_errors(mk_test_spec_no_spectral_axis):
 
 
 def test_trace_inputs(mk_test_img_raw):
+    """
+    Tests for the input argument 'traces' to `Background`. This should accept
+    a list of or a single Trace object, or a list of or a single (positive)
+    number to define a FlatTrace.
+    """
 
     image = mk_test_img_raw
 
@@ -135,3 +142,181 @@ def test_trace_inputs(mk_test_img_raw):
                 'or None to use a FlatTrace in the middle of the image.'
     with pytest.raises(ValueError, match=match_str):
         Background(image, 'non_valid_trace_pos')
+
+
+class TestMasksBackground():
+
+    """
+    Various test functions to test how masked and non-finite data is handled
+    in `Background. There are three currently implemented options for masking
+    in Background: filter, omit, and zero-fill.
+    """
+
+    def mk_img(self, nrows=4, ncols=5, nan_slices=None):
+        """
+        Make a simple gradient image to test masking in Background.
+        Optionally add NaNs to data with `nan_slices`. Returned array is in
+        u.DN.
+        """
+
+        img = np.tile((np.arange(1., ncols + 1)), (nrows, 1))
+
+        if nan_slices:  # add nans in data
+            for s in nan_slices:
+                img[s] = np.nan
+
+        return img * u.DN
+
+    @pytest.mark.parametrize("mask", ["filter", "omit", "zero-fill"])
+    def test_fully_masked_column(self, mask):
+        """
+        Test background with some fully-masked columns (not fully masked image).
+        In this case, the background value for that fully-masked column should
+        be 0.0, with no error or warning raised. This is the case for
+        mask_treatment=filter, omit, or zero-fill.
+        """
+
+        img = self.mk_img(nrows=10, ncols=10)
+        img[:, 0:1] = np.nan
+
+        bkg = Background(img, traces=FlatTrace(img, 6), mask_treatment=mask)
+        assert np.all(bkg.bkg_image().data[:, 0:1] == 0.0)
+
+    @pytest.mark.parametrize("mask", ["filter", "omit"])
+    def test_fully_masked_image(self, mask):
+        """
+        Test that the appropriate error is raised by `Background` when image
+        is fully masked/NaN.
+        """
+
+        with pytest.raises(ValueError, match='Image is fully masked.'):
+            # fully NaN image
+            img = self.mk_img() * np.nan
+            Background(img, traces=FlatTrace(self.mk_img(), 2), mask_treatment=mask)
+
+        with pytest.raises(ValueError, match='Image is fully masked.'):
+            # fully masked image (should be equivalent)
+            img = NDData(np.ones((4, 5)), mask=np.ones((4, 5)))
+            Background(img, traces=FlatTrace(self.mk_img(), 2), mask_treatment=mask)
+
+        # Now test that an image that isn't fully masked, but is fully masked
+        # within the window determined by `width`, produces the correct result.
+        # only applicable for mask_treatment=filter, because this is the only
+        # option that allows a slice of masked values that don't span all rows.
+        msg = 'Image is fully masked within background window determined by `width`.'
+        with pytest.raises(ValueError, match=msg):
+            img = self.mk_img(nrows=12, ncols=12, nan_slices=[np.s_[3:10, :]])
+            Background(img, traces=FlatTrace(img, 6), width=7)
+
+    @pytest.mark.filterwarnings("ignore:background window extends beyond image boundaries")
+    @pytest.mark.parametrize("method,expected",
+                             [("filter", np.array([1., 2., 3., 4., 5., 6., 7.,
+                                                  8., 9., 10., 11., 12.])),
+                              ("omit", np.array([0., 2., 3., 0., 5., 6.,
+                                                 7., 0., 9., 10., 11., 12.])),
+                              ("zero-fill", np.array([0.58333333, 2., 3.,
+                                                      2.33333333, 5., 6., 7.,
+                                                      7.33333333, 9., 10., 11.,
+                                                      12.]))])
+    def test_mask_treatment_bkg_img_spectrum(self, method, expected):
+        """
+        This test function tests `Background.bkg_image` and
+        `Background.bkg_spectrum` when there is masked data. It also tests
+        background subtracting the image, and returning the spectrum of the
+        background subtracted image. This test is parameterized over all
+        currently implemented mask handling methods (filter, omit, and
+        zero-fill) to test that all three work as intended. The window size is
+        set to use the entire image array, so warning about background window
+        is ignored."""
+
+        img_size = 12  # square 12 x 12 image
+
+        # make image, set some value to nan, which will be masked in the function
+        image1 = self.mk_img(nrows=img_size, ncols=img_size,
+                             nan_slices=[np.s_[5:10, 0], np.s_[7:12, 3],
+                                         np.s_[2, 7]])
+
+        # also make an image that doesn't have nonf data values, but has
+        # masked values at the same locations, to make sure they give the same
+        # results
+        mask = ~np.isfinite(image1)
+        dat = self.mk_img(nrows=img_size, ncols=img_size)
+        image2 = NDData(dat, mask=mask)
+
+        for image in [image1, image2]:
+
+            # construct a flat trace in center of image
+            trace = FlatTrace(image, img_size / 2)
+
+            # create 'Background' object with `mask_treatment` set
+            # 'width' should be > size of image to use all pix (but warning will
+            # be raised, which we ignore.)
+            background = Background(image, mask_treatment=method,
+                                    traces=trace, width=img_size + 1)
+
+            # test background image matches 'expected'
+            bk_img = background.bkg_image()
+            # change this and following assertions to assert_quantity_allclose once
+            # issue #213 is fixed
+            np.testing.assert_allclose(bk_img.flux.value,
+                                       np.tile(expected, (img_size, 1)))
+
+            # test background spectrum matches 'expected' times the number of rows
+            # in cross disp axis, since this is a sum and all values in a col are
+            # the same.
+            bk_spec = background.bkg_spectrum()
+            np.testing.assert_allclose(bk_spec.flux.value, expected * img_size)
+
+    def test_sub_bkg_image(self):
+        """
+        Test that masked and nonfinite data is handled correctly when subtracting
+        background from image, for all currently implemented masking
+        options ('filter', 'omit', and 'zero-fill').
+        """
+
+        # make image, set some value to nan, which will be masked in the function
+        image = self.mk_img(nrows=12, ncols=12,
+                            nan_slices=[np.s_[5:10, 0], np.s_[7:12, 3],
+                                        np.s_[2, 7]])
+
+        # Calculate a background value using mask_treatment = 'filter'.
+        # For 'filter', the flag applies to how masked values are handled during
+        # calculation of background for each column, but nonfinite data will
+        # remain in input data array
+        background_filter = Background(image, mask_treatment='filter',
+                                       traces=FlatTrace(image, 6),
+                                       width=2)
+        subtracted_img_filter = background_filter.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_filter.data) == np.isfinite(image.data))
+
+        # Calculate a background value using mask_treatment = 'omit'. The input
+        # 2d mask is reduced to a 1d mask to mask out full columns in the
+        # presence of any nans - this means that (as tested above in
+        # `test_mask_treatment_bkg_img_spectrum`) those columns will have 0.0
+        # background. In this case, image.mask is expanded to mask full
+        # columns - the image itself will not have full columns set to np.nan,
+        # so there are still valid background subtracted data values in this
+        # case, but the corresponding mask for that entire column will be masked.
+
+        background_omit = Background(image, mask_treatment='omit',
+                                     traces=FlatTrace(image, 6),
+                                     width=2)
+        subtracted_img_omit = background_omit.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_omit.data) == np.isfinite(image.data))
+
+        # Calculate a background value using mask_treatment = 'zero-fill'. Data
+        # values at masked locations are set to 0 in the image array, and the
+        # background value calculated for that column will be subtracted
+        # resulting in a negative value. The resulting background subtracted
+        # image should be fully finite and the mask should be zero everywhere
+        # (all unmasked)
+
+        background_zero_fill = Background(image, mask_treatment='zero-fill',
+                                          traces=FlatTrace(image, 6),
+                                          width=2)
+        subtracted_img_zero_fill = background_zero_fill.sub_image()
+
+        assert np.all(np.isfinite(subtracted_img_zero_fill.data))
+        assert np.all(subtracted_img_zero_fill.mask == 0)
diff --git a/specreduce/tests/test_extract.py b/specreduce/tests/test_extract.py
index 4465a1b..7200d2d 100644
--- a/specreduce/tests/test_extract.py
+++ b/specreduce/tests/test_extract.py
@@ -2,20 +2,23 @@
 import pytest
 from astropy import units as u
 from astropy.modeling import models
-from astropy.nddata import VarianceUncertainty, UnknownUncertainty
+from astropy.nddata import NDData, VarianceUncertainty, UnknownUncertainty
 from astropy.tests.helper import assert_quantity_allclose
+from specutils import Spectrum1D
 
+from specreduce.background import Background
 from specreduce.extract import (
     BoxcarExtract, HorneExtract, OptimalExtract, _align_along_trace
 )
-from specreduce.tracing import FlatTrace, ArrayTrace
+from specreduce.tracing import FitTrace, FlatTrace, ArrayTrace
 
 
 def add_gaussian_source(image, amps=2, stddevs=2, means=None):
 
     """ Modify `image.data` to add a horizontal spectrum across the image.
         Each column can have a different amplitude, stddev or mean position
-        if these are arrays (otherwise, constant across image)."""
+        if these are arrays (otherwise, constant across image).
+    """
 
     nrows, ncols = image.shape
 
@@ -325,3 +328,65 @@ def test_horne_interpolated_nbins_fails(mk_test_img):
                           spatial_profile={'name': 'interpolated_profile',
                                            'n_bins_interpolated_profile': 100})
         ex.spectrum
+
+
+class TestMasksExtract():
+
+    def mk_flat_gauss_img(self, nrows=200, ncols=160, nan_slices=None, add_noise=True):
+
+        """
+        Makes a flat gaussian image for testing, with optional added gaussian
+        nosie and optional data values set to NaN. Variance is included, which
+        is required by HorneExtract. Returns a Spectrum1D with flux, spectral
+        axis, and uncertainty.
+        """
+
+        sigma_pix = 4
+        col_model = models.Gaussian1D(amplitude=1, mean=nrows/2,
+                                      stddev=sigma_pix)
+        spec2dvar = np.ones((nrows, ncols))
+        noise = 0
+        if add_noise:
+            np.random.seed(7)
+            sigma_noise = 1
+            noise = np.random.normal(scale=sigma_noise, size=(nrows, ncols))
+
+        index_arr = np.tile(np.arange(nrows), (ncols, 1))
+        img = col_model(index_arr.T) + noise
+
+        if nan_slices:  # add nans in data
+            for s in nan_slices:
+                img[s] = np.nan
+
+        wave = np.arange(0, img.shape[1], 1)
+        objectspec = Spectrum1D(spectral_axis=wave*u.m, flux=img*u.Jy,
+                                uncertainty=VarianceUncertainty(spec2dvar*u.Jy*u.Jy))
+
+        return objectspec
+
+    def test_boxcar_fully_masked(self):
+        """
+        Test that the appropriate error is raised by `BoxcarExtract` when image
+        is fully masked/NaN.
+        """
+        return
+
+        img = self.mk_flat_gauss_img()
+        trace = FitTrace(img)
+
+        with pytest.raises(ValueError, match='Image is fully masked.'):
+            # fully NaN image
+            img = np.zeros((4, 5)) * np.nan
+            Background(img, traces=trace, width=2)
+
+        with pytest.raises(ValueError, match='Image is fully masked.'):
+            # fully masked image (should be equivalent)
+            img = NDData(np.ones((4, 5)), mask=np.ones((4, 5)))
+            Background(img, traces=trace, width=2)
+
+        # Now test that an image that isn't fully masked, but is fully masked
+        # within the window determined by `width`, produces the correct result
+        msg = 'Image is fully masked within background window determined by `width`.'
+        with pytest.raises(ValueError, match=msg):
+            img = self.mk_img(nrows=12, ncols=12, nan_slices=[np.s_[3:10, :]])
+            Background(img, traces=FlatTrace(img, 6), width=7)
diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py
index 867ef2c..03cc294 100644
--- a/specreduce/tests/test_tracing.py
+++ b/specreduce/tests/test_tracing.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from astropy.modeling import models
+from astropy.modeling import fitting, models
 from astropy.nddata import NDData
 import astropy.units as u
 from specreduce.utils.synth_data import make_2d_trace_image
@@ -9,6 +9,34 @@
 IM = make_2d_trace_image()
 
 
+def mk_img(nrows=200, ncols=160, nan_slices=None, add_noise=True):
+
+    """
+    Makes a gaussian image for testing, with optional added gaussian
+    nosie and optional data values set to NaN.
+    """
+
+    # NOTE: Will move this to a fixture at some point.
+
+    sigma_pix = 4
+    col_model = models.Gaussian1D(amplitude=1, mean=nrows/2,
+                                  stddev=sigma_pix)
+    noise = 0
+    if add_noise:
+        np.random.seed(7)
+        sigma_noise = 1
+        noise = np.random.normal(scale=sigma_noise, size=(nrows, ncols))
+
+    index_arr = np.tile(np.arange(nrows), (ncols, 1))
+    img = col_model(index_arr.T) + noise
+
+    if nan_slices:  # add nans in data
+        for s in nan_slices:
+            img[s] = np.nan
+
+    return img * u.DN
+
+
 # test basic trace class
 def test_basic_trace():
     t_pos = IM.shape[0] / 2
@@ -42,6 +70,8 @@ def test_negative_flat_trace_err():
     # negative trace_pos
     with pytest.raises(ValueError, match='must be positive.'):
         FlatTrace(IM, trace_pos=-1)
+    with pytest.raises(ValueError, match='must be positive.'):
+        FlatTrace(IM, trace_pos=0)
 
 
 # test array traces
@@ -70,6 +100,12 @@ def test_array_trace():
     assert np.ma.is_masked(t_short[-1])
     assert t_short.shape[0] == IM.shape[1]
 
+    # make sure nonfinite data in input `trace` is masked
+    arr[0:5] = np.nan
+    t = ArrayTrace(IM, arr)
+    assert np.all(t.trace.mask[0:5])
+    assert np.all(t.trace.mask[5:] == 0)
+
 
 # test fitted traces
 @pytest.mark.filterwarnings("ignore:Model is linear in parameters")
@@ -143,35 +179,103 @@ def test_fit_trace():
         FitTrace(img, bins=ncols + 1)
 
 
+@pytest.mark.filterwarnings("ignore:The fit may be unsuccessful")
+@pytest.mark.filterwarnings("ignore:Model is linear in parameters")
 class TestMasksTracing():
-
-    def mk_img(self, nrows=200, ncols=160):
-
-        np.random.seed(7)
-
-        sigma_pix = 4
-        sigma_noise = 1
-
-        col_model = models.Gaussian1D(amplitude=1, mean=nrows/2, stddev=sigma_pix)
-        noise = np.random.normal(scale=sigma_noise, size=(nrows, ncols))
-
-        index_arr = np.tile(np.arange(nrows), (ncols, 1))
-        img = col_model(index_arr.T) + noise
-
-        return img * u.DN
-
-    def test_window_fit_trace(self):
-
-        """This test function will test that masked values are treated correctly in
-        FitTrace, and produce the correct results and warning messages based on
-        `peak_method`."""
-        img = self.mk_img()
-
-        # create same-shaped variations of image with invalid values
+    """
+    There are three currently implemented options for masking in FitTrace: filter,
+    omit, and zero-fill. Trace, FlatTrace, and ArrayTrace do not have
+    `mask_treatment` options as input because masked/nonfinite values in the data
+    are not relevant for those trace types as they are not affected by masked
+    input data. The tests in this class test masking options for FitTrace, as
+    well as some basic tests (errors, etc) for the other trace types.
+    """
+
+    def test_flat_and_basic_trace_mask(self):
+        """
+        Mask handling is not relevant for basic and flat trace - nans or masked
+        values in the input image will not impact the trace value. The attribute
+        should be initialized though, and be one of the valid options ([None]
+        in this case), for consistancy with all other Specreduce operations.
+        Note that unlike FitTrace, a fully-masked image should NOT result in an
+        error raised because the trace does not depend on the data.
+        """
+
+        img = mk_img(nrows=5, ncols=5)
+
+        basic_trace = Trace(img)
+        assert basic_trace.mask_treatment is None
+
+        flat_trace = FlatTrace(img, trace_pos=2)
+        assert flat_trace.mask_treatment is None
+
+        arr = [1, 2, np.nan, 3, 4]
+        array_trace = ArrayTrace(img, arr)
+        assert array_trace.mask_treatment is None
+
+    def test_array_trace_masking(self):
+        """
+        The `trace` input to ArrayTrace can be a masked array, or an array
+        containing nonfinite data which will be converted to a masked array.
+        Additionally, if any padding needs to be added, the returned trace will
+        be a masked array. Otherwise, it should be a regular array.
+
+        Even though an ArrayTrace may have nans or masked values
+        in the input 1D array for the trace, `mask_treatment_method` refers
+        to how masked values in the input data should be treated. Nans / masked
+        values passed in the array trace should be considered intentional, so
+        also test that `mask_treatment` is initialized to None.
+        """
+        img = mk_img(nrows=10, ncols=10)
+
+        # non-finite data in input trace should be masked out
+        trace_arr = np.array((1, 2, np.nan, 4, 5))
+        array_trace = ArrayTrace(img, trace_arr)
+        assert array_trace.trace.mask[2]
+        assert isinstance(array_trace.trace, np.ma.MaskedArray)
+
+        # and combined with any input masked data
+        trace_arr = np.ma.MaskedArray([1, 2, np.nan, 4, 5], mask=[1, 0, 0, 0, 0])
+        array_trace = ArrayTrace(img, trace_arr)
+        assert array_trace.trace.mask[0]
+        assert array_trace.trace.mask[2]
+        assert isinstance(array_trace.trace, np.ma.MaskedArray)
+
+        # check that mask_treatment is None as there are no valid choices
+        assert array_trace.mask_treatment is None
+
+        # check that if array is fully finite and not masked, that the returned
+        # trace is a notrmal array, not a masked array
+        trace = ArrayTrace(img, np.ones(100))
+        assert isinstance(trace.trace, np.ndarray)
+        assert not isinstance(trace.trace, np.ma.MaskedArray)
+
+        # ensure correct warning is raised when entire trace is masked.
+        trace_arr = np.ma.MaskedArray([1, 2, np.nan, 4, 5], mask=[1, 1, 0, 1, 1])
+        with pytest.raises(UserWarning, match=r'Entire trace array is masked.'):
+            array_trace = ArrayTrace(img, trace_arr)
+
+    def test_fit_trace_fully_masked_image(self):
+        """
+        Test that the correct warning is raised when a fully maksed image is
+        encountered. Also test that when a non-fully masked image is provided,
+        but `window` is set and the image is fully masked within that window,
+        that the correct error is raised.
+        """
+
+        # make simple gaussian image.
+        img = mk_img()
+
+        # create same-shaped variations of image with nans in data array
+        # which will be masked within FitTrace.
         nrows = 200
         ncols = 160
         img_all_nans = np.tile(np.nan, (nrows, ncols))
 
+        # error on trace of all-nan image
+        with pytest.raises(ValueError, match=r'Image is fully masked. Check for invalid values.'):
+            FitTrace(img_all_nans)
+
         window = 10
         guess = int(nrows/2)
         img_win_nans = img.copy()
@@ -181,49 +285,16 @@ def test_window_fit_trace(self):
         with pytest.raises(ValueError, match='pixels in window region are masked'):
             FitTrace(img_win_nans, guess=guess, window=window)
 
-        # error on trace of all-nan image
-        with pytest.raises(ValueError, match=r'image is fully masked'):
-            FitTrace(img_all_nans)
-
-    @pytest.mark.filterwarnings("ignore:The fit may be unsuccessful")
-    @pytest.mark.filterwarnings("ignore:Model is linear in parameters")
-    @pytest.mark.filterwarnings("ignore:All pixels in bins")
-    def test_fit_trace_all_nan_cols(self):
-
-        # make sure that the actual trace that is fit is correct when
-        # all-masked bin peaks are set to NaN
-        img = self.mk_img(nrows=10, ncols=11)
-
-        img[:, 7] = np.nan
-        img[:, 4] = np.nan
-        img[:, 0] = np.nan
-
-        # peak_method = 'max'
-        truth = [1.6346154, 2.2371795, 2.8397436, 3.4423077, 4.0448718,
-                 4.6474359, 5.25, 5.8525641, 6.4551282, 7.0576923,
-                 7.6602564]
-        max_trace = FitTrace(img, peak_method='max')
-        np.testing.assert_allclose(truth, max_trace.trace)
-
-        # peak_method = 'gaussian'
-        truth = [1.947455, 2.383634, 2.8198131, 3.2559921, 3.6921712,
-                 4.1283502, 4.5645293, 5.0007083, 5.4368874, 5.8730665,
-                 6.3092455]
-        max_trace = FitTrace(img, peak_method='gaussian')
-        np.testing.assert_allclose(truth, max_trace.trace)
-
-        # peak_method = 'centroid'
-        truth = [2.5318835, 2.782069, 3.0322546, 3.2824402, 3.5326257,
-                 3.7828113, 4.0329969, 4.2831824, 4.533368, 4.7835536,
-                 5.0337391]
-        max_trace = FitTrace(img, peak_method='centroid')
-        np.testing.assert_allclose(truth, max_trace.trace)
-
-    @pytest.mark.filterwarnings("ignore:The fit may be unsuccessful")
-    @pytest.mark.filterwarnings("ignore:Model is linear in parameters")
-    def test_warn_msg_fit_trace_all_nan_cols(self):
-
-        img = self.mk_img()
+    def test_fit_trace_fully_masked_columns_warn_msg(self):
+        """
+        Test that the correct warning message is raised when fully masked columns
+        (in a not-fully-masked image) are encountered in FitTrace. These columns
+        will be set to NaN and filtered from the final all-bin fit (as tested in
+        test_fit_trace_fully_masked_cols), but a warning message is raised. This
+        should happen for mask_treatment='filter' and 'omit' (for 'zero-fill',
+        all NaN columns become all-zero columns).
+        """
+        img = mk_img()
 
         # test that warning (dependent on choice of `peak_method`) is raised when a
         # few bins are masked, and that theyre listed individually
@@ -252,3 +323,190 @@ def test_warn_msg_fit_trace_all_nan_cols(self):
                           '0, 1, 2, 3, 4, 5, 6, 7, 8, 9, ..., 20 are '
                           'fully masked. Setting bin peaks to NaN.'):
             FitTrace(nddat)
+
+    @pytest.mark.filterwarnings("ignore:All pixels in bins")
+    @pytest.mark.parametrize("mask_treatment", ['filter', 'omit'])
+    def test_fit_trace_fully_masked_cols(self, mask_treatment):
+        """
+        Create a test image with some fully-nan/masked columns, and test that
+        when the final fit to all bin peaks is done for the trace, that these
+        fully-masked columns are set to NaN and filtered during the final all-bin
+        fit. This should happen for mask_treatment = 'filter' and 'omit'
+        (for 'zero-fill', all NaN columns become all-zero columns). Ignore the
+        warning that is produced when this case is encountered (that is tested
+        in `test_fit_trace_fully_masked_cols_warn_msg`.)
+        """
+        img = mk_img(nrows=10, ncols=11)
+
+        # set some columns fully to nan, which will be masked out
+        img[:, 7] = np.nan
+        img[:, 4] = np.nan
+        img[:, 0] = np.nan
+
+        # also create an image that doesn't have nans in the data, but
+        # is masked in the same locations, to make sure that is equivilant.
+
+        # test peak_method = 'max'
+        truth = [1.6346154, 2.2371795, 2.8397436, 3.4423077, 4.0448718,
+                 4.6474359, 5.25, 5.8525641, 6.4551282, 7.0576923,
+                 7.6602564]
+        max_trace = FitTrace(img, peak_method='max',
+                             mask_treatment=mask_treatment)
+        np.testing.assert_allclose(truth, max_trace.trace)
+
+        # peak_method = 'gaussian'
+        truth = [1.947455, 2.383634, 2.8198131, 3.2559921, 3.6921712,
+                 4.1283502, 4.5645293, 5.0007083, 5.4368874, 5.8730665,
+                 6.3092455]
+        max_trace = FitTrace(img, peak_method='gaussian',
+                             mask_treatment=mask_treatment)
+        np.testing.assert_allclose(truth, max_trace.trace)
+
+        # peak_method = 'centroid'
+        truth = [2.5318835, 2.782069, 3.0322546, 3.2824402, 3.5326257,
+                 3.7828113, 4.0329969, 4.2831824, 4.533368, 4.7835536,
+                 5.0337391]
+        max_trace = FitTrace(img, peak_method='centroid', mask_treatment=mask_treatment)
+        np.testing.assert_allclose(truth, max_trace.trace)
+
+    @pytest.mark.filterwarnings("ignore:All pixels in bins")
+    @pytest.mark.parametrize("peak_method,expected",
+                             [("max", [5., 3., 5., 5., 7., 5.,
+                                       np.nan, 5., 5., 5., 2., 5.]),
+                              ("gaussian", [5., 2.10936004, 5., 5., 7.80744334,
+                                            5., np.nan, 5., 5., 5., 1.28216332, 5.]),
+                              ("centroid", [4.27108332, 2.24060342, 4.27108332,
+                                            4.27108332, 6.66827608, 4.27108332,
+                                            np.nan, 4.27108332, 4.27108332,
+                                            4.27108332, 1.19673467, 4.27108332])])
+    def test_mask_treatment_filter(self, peak_method, expected):
+        """
+        Test for mask_treatment=filter for FitTrace.
+        With this masking option, masked and nonfinite data should be filtered
+        when determining bin/column peak. Fully masked bins should be omitted
+        from the final all-bin-peak fit for the Trace. Parametrized over different
+        `peak_method` options.
+        """
+
+        # Make an image with some nonfinite values.
+        image1 = mk_img(nan_slices=[np.s_[4:8, 1:2], np.s_[2:7, 4:5],
+                                    np.s_[:, 6:7], np.s_[3:9, 10:11]],
+                        nrows=10, ncols=12, add_noise=False)
+
+        # Also make an image that doesn't have nonf data values, but has masked
+        # values at the same locations, to make sure they give the same results.
+        mask = ~np.isfinite(image1)
+        dat = mk_img(nrows=10, ncols=12, add_noise=False)
+        image2 = NDData(dat, mask=mask)
+
+        for imgg in [image1, image2]:
+            # run FitTrace, with the testing-only flag _save_bin_peaks_testing set
+            # to True to return the bin peak values before fitting the trace
+            trace = FitTrace(imgg, peak_method=peak_method,
+                             _save_bin_peaks_testing=True)
+            x_bins, y_bins = trace._bin_peaks_testing
+            np.testing.assert_allclose(y_bins, expected)
+
+            # check that final fit to all bins, accouting for fully-masked bins,
+            # matches the trace
+            fitter = fitting.LevMarLSQFitter()
+            mask = np.isfinite(y_bins)
+            all_bin_fit = fitter(trace.trace_model, x_bins[mask], y_bins[mask])
+            all_bin_fit = all_bin_fit((np.arange(12)))
+
+            np.testing.assert_allclose(trace.trace, all_bin_fit)
+
+    @pytest.mark.filterwarnings("ignore:All pixels in bins")
+    @pytest.mark.parametrize("peak_method,expected",
+                             [("max", [5., 3., 5., 5., 7., 5.,
+                                       0., 5., 5., 5., 2., 5.]),
+                              ("gaussian", [5., 1.09382384, 5., 5., 7.81282206,
+                                            5., 0., 5., 5., 5., 1.28216332, 5.]),
+                              ("centroid", [4.27108332, 2.24060342, 4.27108332,
+                                            4.27108332, 6.66827608, 4.27108332,
+                                            9., 4.27108332, 4.27108332,
+                                            4.27108332, 1.19673467, 4.27108332])])
+    def test_mask_treatment_zero_fill(self, peak_method, expected):
+        """
+        Test for mask_treatment=`zero_fill` for FitTrace.
+        Masked and nonfinite data are replaced with zero in the data array,
+        and the input mask is then dropped. Parametrized over different
+        `peak_method` options.
+        """
+
+        # Make an image with some nonfinite values.
+        image1 = mk_img(nan_slices=[np.s_[4:8, 1:2], np.s_[2:7, 4:5],
+                                    np.s_[:, 6:7], np.s_[3:9, 10:11]],
+                        nrows=10, ncols=12, add_noise=False)
+
+        # Also make an image that doesn't have nonf data values, but has masked
+        # values at the same locations, to make sure they give the same results.
+        mask = ~np.isfinite(image1)
+        dat = mk_img(nrows=10, ncols=12, add_noise=False)
+        image2 = NDData(dat, mask=mask)
+
+        for imgg in [image1, image2]:
+            # run FitTrace, with the testing-only flag _save_bin_peaks_testing set
+            # to True to return the bin peak values before fitting the trace
+            trace = FitTrace(imgg, peak_method=peak_method,
+                             mask_treatment='zero-fill',
+                             _save_bin_peaks_testing=True)
+            x_bins, y_bins = trace._bin_peaks_testing
+            np.testing.assert_allclose(y_bins, expected)
+
+            # check that final fit to all bins, accouting for fully-masked bins,
+            # matches the trace
+            fitter = fitting.LevMarLSQFitter()
+            mask = np.isfinite(y_bins)
+            all_bin_fit = fitter(trace.trace_model, x_bins[mask], y_bins[mask])
+            all_bin_fit = all_bin_fit((np.arange(12)))
+
+            np.testing.assert_allclose(trace.trace, all_bin_fit)
+
+    @pytest.mark.filterwarnings("ignore:All pixels in bins")
+    @pytest.mark.parametrize("peak_method,expected",
+                             [("max", [5., np.nan, 5., 5., np.nan, 5.,
+                                       np.nan, 5., 5., 5., np.nan,  5.]),
+                              ("gaussian", [5., np.nan, 5., 5., np.nan, 5.,
+                                            np.nan, 5., 5., 5., np.nan, 5.]),
+                              ("centroid", [4.27108332, np.nan, 4.27108332,
+                                            4.27108332, np.nan, 4.27108332,
+                                            np.nan, 4.27108332, 4.27108332,
+                                            4.27108332, np.nan, 4.27108332])])
+    def test_mask_treatment_omit(self, peak_method, expected):
+        """
+        Test for mask_treatment=`omit` for FitTrace. Columns (assuming
+        disp_axis==1) with any masked data values will be fully masked and
+        therefore not contribute to the bin peaks. Parametrized over different
+        `peak_method` options.
+        """
+
+        # Make an image with some nonfinite values.
+        image1 = mk_img(nan_slices=[np.s_[4:8, 1:2], np.s_[2:7, 4:5],
+                                    np.s_[:, 6:7], np.s_[3:9, 10:11]],
+                        nrows=10, ncols=12, add_noise=False)
+
+        # Also make an image that doesn't have nonfinite data values, but has masked
+        # values at the same locations, to make sure those cases are equivalent
+        mask = ~np.isfinite(image1)
+        dat = mk_img(nrows=10, ncols=12, add_noise=False)
+        image2 = NDData(dat, mask=mask)
+
+        for imgg in [image1, image2]:
+
+            # run FitTrace, with the testing-only flag _save_bin_peaks_testing set
+            # to True to return the bin peak values before fitting the trace
+            trace = FitTrace(imgg, peak_method=peak_method,
+                             mask_treatment='omit',
+                             _save_bin_peaks_testing=True)
+            x_bins, y_bins = trace._bin_peaks_testing
+            np.testing.assert_allclose(y_bins, expected)
+
+            # check that final fit to all bins, accouting for fully-masked bins,
+            # matches the trace
+            fitter = fitting.LevMarLSQFitter()
+            mask = np.isfinite(y_bins)
+            all_bin_fit = fitter(trace.trace_model, x_bins[mask], y_bins[mask])
+            all_bin_fit = all_bin_fit((np.arange(12)))
+
+            np.testing.assert_allclose(trace.trace, all_bin_fit)
diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index 7964480..f124ddb 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -36,6 +36,13 @@ def __post_init__(self):
         self.trace_pos = self.image.shape[0] / 2
         self.trace = np.ones_like(self.image[0]) * self.trace_pos
 
+        # masking options not relevant for basic Trace
+        self._mask_treatment = None
+        self._valid_mask_treatment_methods = [None]
+
+        # eventually move this to common SpecreduceOperation base class
+        self.validate_masking_options()
+
     def __getitem__(self, i):
         return self.trace[i]
 
@@ -43,6 +50,11 @@ def __getitem__(self, i):
     def shape(self):
         return self.trace.shape
 
+    def validate_masking_options(self):
+        if self.mask_treatment not in self.valid_mask_treatment_methods:
+            raise ValueError(
+                f'`mask_treatment` {self.mask_treatment} not one of {self.valid_mask_treatment_methods}')  # noqa
+
     def shift(self, delta):
         """
         Shift the trace by delta pixels perpendicular to the axis being traced
@@ -61,7 +73,7 @@ def _bound_trace(self):
         Mask trace positions that are outside the upper/lower bounds of the image.
         """
         ny = self.image.shape[0]
-        self.trace = np.ma.masked_outside(self.trace, 0, ny-1)
+        self.trace = np.ma.masked_outside(self.trace, 0, ny - 1)
 
     def __add__(self, delta):
         """
@@ -79,6 +91,14 @@ def __sub__(self, delta):
         """
         return self.__add__(-delta)
 
+    @property
+    def mask_treatment(self):
+        return self._mask_treatment
+
+    @property
+    def valid_mask_treatment_methods(self):
+        return self._valid_mask_treatment_methods
+
 
 @dataclass
 class FlatTrace(Trace, _ImageParser):
@@ -101,6 +121,10 @@ def __post_init__(self):
 
         self.set_position(self.trace_pos)
 
+        # masking options not relevant for basic Trace
+        self._mask_treatment = None
+        self._valid_mask_treatment_methods = [None]
+
     def set_position(self, trace_pos):
         """
         Set the trace position within the image
@@ -124,12 +148,33 @@ class ArrayTrace(Trace, _ImageParser):
 
     Parameters
     ----------
-    trace : `numpy.ndarray`
-        Array containing trace positions
+    trace : `numpy.ndarray` or `numpy.ma.MaskedArray`
+        Array containing trace positions.
     """
     trace: np.ndarray
 
     def __post_init__(self):
+
+        # masking options not relevant for ArrayTrace. any non-finite or masked
+        # data in `image` will not affect array trace
+        self._mask_treatment = None
+        self._valid_mask_treatment_methods = [None]
+
+        # masked array will have a .data, regular array will not.
+        trace_data = getattr(self.trace, 'data', self.trace)
+
+        # but we do need to mask uncaught non-finite values in input trace array
+        # which should also be combined with any existing mask in the input `trace`
+        if hasattr(self.trace, 'mask'):
+            total_mask = np.logical_or(self.trace.mask, ~np.isfinite(trace_data))
+        else:
+            total_mask = ~np.isfinite(trace_data)
+
+        # always work with masked array, even if there is no masked
+        # or nonfinite data, in case padding is needed. if not, mask will be
+        # dropped at the end and a regular array will be returned.
+        self.trace = np.ma.MaskedArray(trace_data, total_mask)
+
         self.image = self._parse_image(self.image)
 
         nx = self.image.shape[1]
@@ -143,8 +188,17 @@ def __post_init__(self):
                 # padding will be the last value of the trace, but will be masked out.
                 padding = np.ma.MaskedArray(np.ones(nx - nt) * self.trace[-1], mask=True)
                 self.trace = np.ma.hstack([self.trace, padding])
+
         self._bound_trace()
 
+        # warn if entire trace is masked
+        if np.all(self.trace.mask):
+            warnings.warn("Entire trace array is masked.")
+
+        # and return plain array if nothing is masked
+        if not np.any(self.trace.mask):
+            self.trace = self.trace.data
+
 
 @dataclass
 class FitTrace(Trace, _ImageParser):
@@ -201,6 +255,21 @@ class FitTrace(Trace, _ImageParser):
         ``centroid``: Takes the centroid of the window within in bin.
         ``max``: Saves the position with the maximum flux in each bin.
         [default: ``max``]
+    mask_treatment : string, optional
+        The method for handling masked or non-finite data. Choice of `filter`,
+        `omit`, or `zero-fill`. If `filter` is chosen, masked/non-finite data
+        will be filtered during the fit to each bin/column (along disp. axis) to
+        find the peak. If `omit` is chosen, columns along disp_axis with any
+        masked/non-finite data values will be fully masked (i.e, 2D mask is
+        collapsed to 1D and applied). If `zero-fill` is chosen, masked/non-finite
+        data will be replaced with 0.0 in the input image, and the mask will then
+        be dropped. For all three options, the input mask (optional on input
+        NDData object) will be combined with a mask generated from any non-finite
+        values in the image data. Also note that because binning is an option in
+        FitTrace, that masked data will contribute zero to the sum when binning
+        adjacent columns.
+        [default: ``filter``]
+
     """
     bins: int = None
     guess: float = None
@@ -209,24 +278,29 @@ class FitTrace(Trace, _ImageParser):
     peak_method: str = 'max'
     _crossdisp_axis = 0
     _disp_axis = 1
+    mask_treatment: str = 'filter'
+    _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
+    # for testing purposes only, save bin peaks if requested
+    _save_bin_peaks_testing: bool = False
 
     def __post_init__(self):
 
-        # parse image
+        # Parse image, including masked/nonfinite data handling based on
+        # choice of `mask_treatment`. returns a Spectrum1D
         self.image = self._parse_image(self.image)
 
-        # mask any previously uncaught invalid values
-        or_mask = np.logical_or(self.image.mask, ~np.isfinite(self.image.data))
-        img = np.ma.masked_array(self.image.data, or_mask)
+        # _parse_image returns a Spectrum1D. convert this to a masked array
+        # for ease of calculations here (even if there is no masked data).
+        # Note: uncertainties are dropped, this should also be addressed at
+        # some point probably across the package.
+        img = np.ma.masked_array(self.image.data, self.image.mask)
+        self._mask_temp = self.image.mask
 
-        # validate arguments
+        # validate input arguments
         valid_peak_methods = ('gaussian', 'centroid', 'max')
         if self.peak_method not in valid_peak_methods:
             raise ValueError(f"peak_method must be one of {valid_peak_methods}")
 
-        if img.mask.all():
-            raise ValueError('image is fully masked. Check for invalid values')
-
         if self._crossdisp_axis != 0:
             raise ValueError('cross-dispersion axis must equal 0')
 
@@ -317,7 +391,7 @@ def _fit_trace(self, img):
 
             # binned columns, summed along disp. axis.
             # or just a single, unbinned column if no bins
-            z_i = img[ilum2, x_bins[i]:x_bins[i+1]].sum(axis=self._disp_axis)
+            z_i = img[ilum2, x_bins[i]:x_bins[i + 1]].sum(axis=self._disp_axis)
 
             # if this bin is fully masked, set bin peak to NaN so it can be
             # filtered in the final fit to all bin peaks for the trace
@@ -359,7 +433,7 @@ def _fit_trace(self, img):
                 z_i_cumsum = np.cumsum(z_i)
                 # find the interpolated index where the cumulative array reaches
                 # half the total cumulative values
-                y_bins[i] = np.interp(z_i_cumsum[-1]/2., z_i_cumsum, ilum2)
+                y_bins[i] = np.interp(z_i_cumsum[-1] / 2., z_i_cumsum, ilum2)
 
                 # NOTE this reflects current behavior, should eventually be changed
                 # to set to nan by default (or zero fill / interpoate option once
@@ -388,6 +462,12 @@ def _fit_trace(self, img):
         x_bins = (x_bins[:-1] + x_bins[1:]) / 2
 
         # interpolate the fitted trace over the entire wavelength axis
+
+        # for testing purposes only, save bin peaks if requested
+        if self._save_bin_peaks_testing:
+            self._bin_peaks_testing = (x_bins, y_bins)
+
+        # filter non-finite bin peaks before filtering to all bin peaks
         y_finite = np.where(np.isfinite(y_bins))[0]
         if y_finite.size > 0:
             x_bins = x_bins[y_finite]
@@ -397,6 +477,7 @@ def _fit_trace(self, img):
             fitter = (fitting.SplineSmoothingFitter()
                       if isinstance(self.trace_model, models.Spline1D)
                       else fitting.LevMarLSQFitter())
+            self._y_bins = y_bins
             self.trace_model_fit = fitter(self.trace_model, x_bins, y_bins)
 
             trace_x = np.arange(img.shape[self._disp_axis])

From d4d528dddbb91897cbb4d21f70f081a5cfb038fe Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Wed, 18 Sep 2024 15:58:27 -0400
Subject: [PATCH 02/25] rtd fixes

---
 specreduce/background.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/specreduce/background.py b/specreduce/background.py
index 29c81ca..6e1167a 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -32,9 +32,9 @@ class Background(_ImageParser):
     ----------
     image : `~astropy.nddata.NDData`-like or array-like
         image with 2-D spectral image data
-    traces : List, `tracing.Trace`, int, float
+    traces : List, `specreduce.tracing.Trace`, int, float
         Individual or list of trace object(s) (or integers/floats to define
-        FlatTraces) to extract the background. If None, a `FlatTrace` at the
+        FlatTraces) to extract the background. If None, a ``FlatTrace`` at the
         center of the image (according to `disp_axis`) will be used.
     width : float
         width of extraction aperture in pixels
@@ -50,11 +50,11 @@ class Background(_ImageParser):
         [default: 0]
     mask_treatment : string, optional
         The method for handling masked or non-finite data. Choice of `filter`,
-        `omit`, or `zero-fill`. If `filter` is chosen, masked and non-finite
+        ``omit``, or ``zero-fill``. If ``filter`` is chosen, masked and non-finite
         data will not contribute to the background statistic that is calculated
         in each column along `disp_axis`. If `omit` is chosen, columns along
         disp_axis with any masked/non-finite data values will be fully masked
-        (i.e, 2D mask is collapsed to 1D and applied). If `zero-fill` is chosen,
+        (i.e, 2D mask is collapsed to 1D and applied). If ``zero-fill`` is chosen,
         masked/non-finite data will be replaced with 0.0 in the input image,
         and the mask will then be dropped. For all three options, the input mask
         (optional on input NDData object) will be combined with a mask generated

From 9d20d1262fc6fbc1db9d70c9c8b918f6585f7a71 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Sun, 3 Nov 2024 19:33:16 +0000
Subject: [PATCH 03/25] Changed
 `_ImageParser._mask_and_nonfinite_data_handling` into a static method and
 modified the method's logic accordingly.

---
 specreduce/core.py | 98 +++++++++++++++++++++++++---------------------
 1 file changed, 53 insertions(+), 45 deletions(-)

diff --git a/specreduce/core.py b/specreduce/core.py
index 4fc6ded..273dce4 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -28,6 +28,8 @@ class _ImageParser:
         - `~numpy.ndarray`
     """
 
+    implemented_mask_treatment_methods = 'filter', 'zero-fill', 'omit'
+
     def _parse_image(self, image, disp_axis=1):
         """
         Convert all accepted image types to a consistently formatted
@@ -57,9 +59,28 @@ def _parse_image(self, image, disp_axis=1):
         return img
 
     @staticmethod
-    def _get_data_from_image(image, disp_axis=1):
-        """Extract data array from various input types for `image`.
-           Retruns `np.ndarray` of image data."""
+    def _get_data_from_image(image,
+                             disp_axis: int = 1,
+                             mask_treatment: str = 'filter') -> Spectrum1D:
+        """
+        Extract data array from various input types for `image`.
+
+        Parameters
+        ----------
+        image : array-like or Quantity
+            Input image from which data is extracted. This can be a 2D numpy
+            array, Quantity, or an NDData object.
+        disp_axis : int, optional
+            The dispersion axis of the image.
+        mask_treatment : str, optional
+            Treatment method for the mask.
+
+        Returns
+        -------
+        Spectrum1D
+        """
+        # This works only with 2D images.
+        crossdisp_axis = (disp_axis + 1) % 2
 
         if isinstance(image, u.quantity.Quantity):
             img = image.value
@@ -77,12 +98,15 @@ def _get_data_from_image(image, disp_axis=1):
         # handled as well. Note that input data may be modified if a fill value
         # is chosen to handle masked data. The returned image will always have
         # `image.mask` even if there are no nonfinte or masked values.
-        img, mask = self._mask_and_nonfinite_data_handling(image=img, mask=mask)
+        img, mask = _ImageParser._mask_and_nonfinite_data_handling(image=img,
+                                                                   mask=mask,
+                                                                   mask_treatment=mask_treatment,
+                                                                   crossdisp_axis=crossdisp_axis)
 
         # mask (handled above) 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, 'uncertainty', None) is not None:
+        if hasattr(image, 'uncertainty'):
             uncertainty = image.uncertainty
         else:
             uncertainty = VarianceUncertainty(np.ones(img.shape))
@@ -94,26 +118,14 @@ def _get_data_from_image(image, disp_axis=1):
 
         img = Spectrum1D(img * unit, spectral_axis=spectral_axis,
                          uncertainty=uncertainty, mask=mask)
-
         return img
 
     @staticmethod
-    def _get_data_from_image(image):
-        """Extract data array from various input types for `image`.
-           Retruns `np.ndarray` of image data."""
-
-        if isinstance(image, u.quantity.Quantity):
-            img = image.value
-        if isinstance(image, np.ndarray):
-            img = image
-        else:  # NDData, including CCDData and Spectrum1D
-            img = image.data
-        return img
-
-    def _mask_and_nonfinite_data_handling(self, image, mask):
+    def _mask_and_nonfinite_data_handling(image, mask,
+                                          mask_treatment: str = 'filter',
+                                          crossdisp_axis: int = 0) -> tuple[np.ndarray, np.ndarray]:
         """
-        This function handles the treatment of masked and nonfinite data,
-        including input validation.
+        Handle the treatment of masked and nonfinite data.
 
         All operations in Specreduce can take in a mask for the data as
         part of the input NDData. Additionally, any non-finite values in the
@@ -124,31 +136,33 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
         of masked and nonfinite data - filter, omit, and zero-fill.
         Depending on the step, all or a subset of these three options are valid.
 
+        Parameters
+        ----------
+        image : array-like
+            The input image data array that may contain nonfinite values.
+        mask : array-like or None
+            An optional mask array. Nonfinite values in the image will be added to this mask.
+        mask_treatment : str
+            Specifies how to handle masked data:
+            - 'filter' (default): Returns the unmodified input image and combined mask.
+            - 'zero-fill': Sets masked values in the image to zero.
+            - 'omit': Masks entire columns or rows if any value is masked.
+        crossdisp_axis : int
+            Axis along which to collapse the 2D mask into a 1D mask for treatment 'omit'.
         """
-
-        # valid options depend on Specreduce step, and are set accordingly there
-        # for steps that this isn't implemeted for yet, default to 'filter',
-        # which will return unmodified input image and mask
-        mask_treatment = getattr(self, 'mask_treatment', 'filter')
-
-        # make sure chosen option is valid. if _valid_mask_treatment_methods
-        # is not an attribue, proceed with 'filter' to return back inupt data
-        # and mask that is combined with nonfinite data.
-        if mask_treatment is not None:  # None in operations where masks aren't relevant (FlatTrace)
-            valid_mask_treatment_methods = getattr(self, '_valid_mask_treatment_methods', ['filter'])  # noqa
-            if mask_treatment not in valid_mask_treatment_methods:
-                raise ValueError(f"`mask_treatment` must be one of {valid_mask_treatment_methods}")
+        if mask_treatment not in _ImageParser.implemented_mask_treatment_methods:
+            raise ValueError("`mask_treatment` must be one of "
+                             f"{_ImageParser.implemented_mask_treatment_methods}")
 
         # make sure there is always a 'mask', even when all data is unmasked and finite.
         if mask is not None:
-            mask = self.image.mask
             # always mask any previously uncaught nonfinite values in image array
             # combining these with the (optional) user-provided mask on `image.mask`
             mask = np.logical_or(mask, ~np.isfinite(image))
         else:
             mask = ~np.isfinite(image)
 
-        # if mask option is the default 'filter' option, or None,
+        # if mask option is the default 'filter' option,
         # nothing needs to be done. input mask (combined with nonfinite data)
         # remains with data as-is.
 
@@ -165,23 +179,16 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
 
             # masked array with no masked values, so accessing image.mask works
             # but we don't need the actual mask anymore since data has been set to 0
-            mask = np.zeros(image.shape)
+            mask = np.zeros(image.shape, dtype=bool)
 
         elif mask_treatment == 'omit':
             # collapse 2d mask (after accounting for addl non-finite values in
             # data) to a 1d mask, along dispersion axis, to fully mask columns
             # that have any masked values.
 
-            # must have a crossdisp_axis specified to use 'omit' optoin
-            if hasattr(self, 'crossdisp_axis'):
-                crossdisp_axis = self.crossdisp_axis
-            if hasattr(self, '_crossdisp_axis'):
-                crossdisp_axis = self._crossdisp_axis
-
             # create a 1d mask along crossdisp axis - if any column has a single nan,
             # the entire column should be masked
-            reduced_mask = np.logical_or.reduce(mask,
-                                                axis=crossdisp_axis)
+            reduced_mask = np.logical_or.reduce(mask, axis=crossdisp_axis)
 
             # back to a 2D mask
             shape = (image.shape[0], 1) if crossdisp_axis == 0 else (1, image.shape[1])
@@ -193,6 +200,7 @@ def _mask_and_nonfinite_data_handling(self, image, mask):
 
         return image, mask
 
+
 @dataclass
 class SpecreduceOperation(_ImageParser):
     """

From ca9d9aa7051aad2a384d4e68a479dbdb510884ce Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Sun, 3 Nov 2024 20:02:19 +0000
Subject: [PATCH 04/25] Fixed a minor typo in mask_treatment docstring.

---
 specreduce/background.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/specreduce/background.py b/specreduce/background.py
index 6e1167a..ed3fc60 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -49,7 +49,7 @@ class Background(_ImageParser):
         cross-dispersion axis
         [default: 0]
     mask_treatment : string, optional
-        The method for handling masked or non-finite data. Choice of `filter`,
+        The method for handling masked or non-finite data. Choice of ``filter``,
         ``omit``, or ``zero-fill``. If ``filter`` is chosen, masked and non-finite
         data will not contribute to the background statistic that is calculated
         in each column along `disp_axis`. If `omit` is chosen, columns along

From 4fed07d87927dc784b64e9e31733c9139da3a2a2 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Sun, 3 Nov 2024 20:03:42 +0000
Subject: [PATCH 05/25] Added a mask treatment method as an optional argument
 to `_ImageParser._parse_image`.

---
 specreduce/core.py | 28 ++++++++++++++++++----------
 1 file changed, 18 insertions(+), 10 deletions(-)

diff --git a/specreduce/core.py b/specreduce/core.py
index 273dce4..0591b30 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -30,22 +30,28 @@ class _ImageParser:
 
     implemented_mask_treatment_methods = 'filter', 'zero-fill', 'omit'
 
-    def _parse_image(self, image, disp_axis=1):
+    def _parse_image(self, image,
+                     disp_axis: int = 1,
+                     mask_treatment: str = 'filter') -> Spectrum1D:
         """
-        Convert all accepted image types to a consistently formatted
-        Spectrum1D object.
+        Convert all accepted image types to a consistently formatted Spectrum1D object.
 
         Parameters
         ----------
-        image : `~astropy.nddata.NDData`-like or array-like, required
+        image : `~astropy.nddata.NDData`-like or array-like
             The image to be parsed. If None, defaults to class' own
             image attribute.
-        disp_axis : int, optional
+        disp_axis
             The index of the image's dispersion axis. Should not be
             changed until operations can handle variable image
             orientations. [default: 1]
-        """
+        mask_treatment
+            Treatment method for the mask.
 
+        Returns
+        -------
+        Spectrum1D
+        """
         # 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...
@@ -54,9 +60,8 @@ def _parse_image(self, image, disp_axis=1):
             # useful for Background's instance methods
             return self.image
 
-        img = self._get_data_from_image(image, disp_axis=disp_axis)
-
-        return img
+        return self._get_data_from_image(image, disp_axis=disp_axis,
+                                        mask_treatment=mask_treatment)
 
     @staticmethod
     def _get_data_from_image(image,
@@ -73,7 +78,10 @@ def _get_data_from_image(image,
         disp_axis : int, optional
             The dispersion axis of the image.
         mask_treatment : str, optional
-            Treatment method for the mask.
+            Treatment method for the mask:
+            - 'filter' (default): Return the unmodified input image and combined mask.
+            - 'zero-fill': Set masked values in the image to zero.
+            - 'omit': Mask all pixels along the cross dispersion axis if any value is masked.
 
         Returns
         -------

From 6ccb4571cfdab43c9ed24559aad9951ba7eeac84 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Sun, 3 Nov 2024 20:04:17 +0000
Subject: [PATCH 06/25] Updated image parsing method in BoxcarExtract.

---
 specreduce/extract.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/specreduce/extract.py b/specreduce/extract.py
index 96128af..6456401 100644
--- a/specreduce/extract.py
+++ b/specreduce/extract.py
@@ -213,7 +213,7 @@ def __call__(self, image=None, trace_object=None, width=None,
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
         # omit. non-finite data will be masked, always. Returns a Spectrum1D.
-        self.image = self._parse_image(image)
+        self.image = self._parse_image(image, disp_axis=disp_axis, mask_treatment=self.mask_treatment)
 
         # # _parse_image returns a Spectrum1D. convert this to a masked array
         # # for ease of calculations here (even if there is no masked data).

From 340d866ff2f06eebcefcec6f5f6063d012dc7838 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Sun, 3 Nov 2024 20:08:39 +0000
Subject: [PATCH 07/25] Fixed too long lines.

---
 specreduce/core.py    | 2 +-
 specreduce/extract.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/specreduce/core.py b/specreduce/core.py
index 0591b30..00b1f1b 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -61,7 +61,7 @@ def _parse_image(self, image,
             return self.image
 
         return self._get_data_from_image(image, disp_axis=disp_axis,
-                                        mask_treatment=mask_treatment)
+                                         mask_treatment=mask_treatment)
 
     @staticmethod
     def _get_data_from_image(image,
diff --git a/specreduce/extract.py b/specreduce/extract.py
index 6456401..00d6896 100644
--- a/specreduce/extract.py
+++ b/specreduce/extract.py
@@ -213,7 +213,8 @@ def __call__(self, image=None, trace_object=None, width=None,
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
         # omit. non-finite data will be masked, always. Returns a Spectrum1D.
-        self.image = self._parse_image(image, disp_axis=disp_axis, mask_treatment=self.mask_treatment)
+        self.image = self._parse_image(image, disp_axis=disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # # _parse_image returns a Spectrum1D. convert this to a masked array
         # # for ease of calculations here (even if there is no masked data).

From 6f2dbd38b271040083d8c1a157356d5a6a2a43c9 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Mon, 4 Nov 2024 19:22:00 +0000
Subject: [PATCH 08/25] Updated _InmageParser._parse_image method call to
 include disp_axis and mask_treatment.

---
 specreduce/background.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/specreduce/background.py b/specreduce/background.py
index ed3fc60..ac1f08c 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -114,7 +114,8 @@ def __post_init__(self):
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. Any uncaught nonfinte data values will be
         # masked as well. Returns a Spectrum1D.
-        self.image = self._parse_image(self.image)
+        self.image = self._parse_image(self.image, disp_axis=self.disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # always work with masked array, even if there is no masked
         # or nonfinite data, in case padding is needed. if not, mask will be

From c115357d773949db1cc5a6cb332936ac5ee47900 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Mon, 4 Nov 2024 19:23:15 +0000
Subject: [PATCH 09/25] Added a small comment to '_ImageParser' about
 'implemented_mask_treatment_methods'.

---
 specreduce/core.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/specreduce/core.py b/specreduce/core.py
index 00b1f1b..76894b1 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -28,6 +28,8 @@ class _ImageParser:
         - `~numpy.ndarray`
     """
 
+    # The '_valid_mask_treatment_methods' in the Background, Trace, and Extract
+    # classes is a subset of implemented methods.
     implemented_mask_treatment_methods = 'filter', 'zero-fill', 'omit'
 
     def _parse_image(self, image,

From a7553ee9ad25881f61beeb0b23f10773f931b083 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Mon, 4 Nov 2024 19:23:29 +0000
Subject: [PATCH 10/25] Updated _InmageParser._parse_image method call to
 include disp_axis and mask_treatment.

---
 specreduce/tracing.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index f124ddb..9017cfd 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -287,7 +287,8 @@ def __post_init__(self):
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. returns a Spectrum1D
-        self.image = self._parse_image(self.image)
+        self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
+                                       mask_treatment=self.mask_treatment)
 
         # _parse_image returns a Spectrum1D. convert this to a masked array
         # for ease of calculations here (even if there is no masked data).

From 1cb42aa42e2e9b883204d7053ac90d905c4e45f0 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Mon, 4 Nov 2024 19:25:01 +0000
Subject: [PATCH 11/25] Changed 'TestMaskTracing.test_array_trace_masking' to
 use 'pytest.warns' instead of 'pytest.raises' to test for an expected
 warning.

---
 specreduce/tests/test_tracing.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py
index 03cc294..b3e33d4 100644
--- a/specreduce/tests/test_tracing.py
+++ b/specreduce/tests/test_tracing.py
@@ -252,7 +252,7 @@ def test_array_trace_masking(self):
 
         # ensure correct warning is raised when entire trace is masked.
         trace_arr = np.ma.MaskedArray([1, 2, np.nan, 4, 5], mask=[1, 1, 0, 1, 1])
-        with pytest.raises(UserWarning, match=r'Entire trace array is masked.'):
+        with pytest.warns(UserWarning, match=r'Entire trace array is masked.'):
             array_trace = ArrayTrace(img, trace_arr)
 
     def test_fit_trace_fully_masked_image(self):

From 039faf4bf6609c1f6e9c967341a7241a03d95868 Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Fri, 27 Sep 2024 11:47:31 -0400
Subject: [PATCH 12/25] MNT: Use hash for Action workflow versions and update
 if needed

---
 .github/workflows/cron-tests.yml      | 2 +-
 .github/workflows/publish-to-pypi.yml | 2 +-
 .github/workflows/tox-tests.yml       | 2 +-
 3 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/cron-tests.yml b/.github/workflows/cron-tests.yml
index 2bb4da1..c66f63a 100644
--- a/.github/workflows/cron-tests.yml
+++ b/.github/workflows/cron-tests.yml
@@ -22,7 +22,7 @@ concurrency:
 jobs:
   tests:
     if: (github.repository == 'astropy/specreduce' && (github.event_name == 'schedule' || github.event_name == 'push' || contains(github.event.pull_request.labels.*.name, 'Extra CI')))
-    uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@v1
+    uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@924441154cf3053034c6513d5e06c69d262fb9a6  # v1.13.0
     with:
       submodules: false
       coverage: ''
diff --git a/.github/workflows/publish-to-pypi.yml b/.github/workflows/publish-to-pypi.yml
index 60b4296..249c452 100644
--- a/.github/workflows/publish-to-pypi.yml
+++ b/.github/workflows/publish-to-pypi.yml
@@ -8,7 +8,7 @@ on:
 
 jobs:
   publish:
-    uses: OpenAstronomy/github-actions-workflows/.github/workflows/publish_pure_python.yml@v1
+    uses: OpenAstronomy/github-actions-workflows/.github/workflows/publish_pure_python.yml@924441154cf3053034c6513d5e06c69d262fb9a6  # v1.13.0
     # NOTE: Uncomment "if" if you do not want this to run for every PR.
     # if: ((github.event_name == 'push' && startsWith(github.ref, 'refs/tags')) || contains(github.event.pull_request.labels.*.name, 'Build wheels'))
     with:
diff --git a/.github/workflows/tox-tests.yml b/.github/workflows/tox-tests.yml
index 24e7f09..ef33fdd 100644
--- a/.github/workflows/tox-tests.yml
+++ b/.github/workflows/tox-tests.yml
@@ -21,7 +21,7 @@ concurrency:
 
 jobs:
   tests:
-    uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@v1
+    uses: OpenAstronomy/github-actions-workflows/.github/workflows/tox.yml@924441154cf3053034c6513d5e06c69d262fb9a6  # v1.13.0
     secrets:
       CODECOV_TOKEN: ${{ secrets.CODECOV }}
     with:

From 20b846678c26da707e601d9e2ebdea16c231578d Mon Sep 17 00:00:00 2001
From: "P. L. Lim" <2090236+pllim@users.noreply.github.com>
Date: Fri, 27 Sep 2024 17:35:05 -0400
Subject: [PATCH 13/25] Create dependabot.yml

---
 .github/dependabot.yml | 15 +++++++++++++++
 1 file changed, 15 insertions(+)
 create mode 100644 .github/dependabot.yml

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..1a218f5
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,15 @@
+# To get started with Dependabot version updates, you'll need to specify which
+# package ecosystems to update and where the package manifests are located.
+# Please see the documentation for all configuration options:
+# https://docs.github.com/github/administering-a-repository/configuration-options-for-dependency-updates
+
+version: 2
+updates:
+  - package-ecosystem: "github-actions" # See documentation for possible values
+    directory: ".github/workflows" # Location of package manifests
+    schedule:
+      interval: "monthly"
+    groups:
+      actions:
+        patterns:
+          - "*"

From 6e61444a50674a15caf4602b81dbac32fe705f41 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Tue, 5 Nov 2024 20:41:13 +0000
Subject: [PATCH 14/25] Created a 'specreduce.image.Image' class that can
 replace the 'specreduce.core._ImageParser' mixin class.

---
 specreduce/image.py | 120 ++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 120 insertions(+)
 create mode 100644 specreduce/image.py

diff --git a/specreduce/image.py b/specreduce/image.py
new file mode 100644
index 0000000..2a5f742
--- /dev/null
+++ b/specreduce/image.py
@@ -0,0 +1,120 @@
+from copy import deepcopy
+
+import astropy.units as u
+import numpy as np
+
+from astropy.nddata import NDData, VarianceUncertainty
+
+
+class Image:
+    implemented_masking_methods = 'filter', 'zero-fill', 'omit'
+
+    def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None):
+        self._image: NDData = Image._parse_image(image)
+        self._disp_axis: int = disp_axis
+        self._crossdisp_axis: int = crossdisp_axis
+
+        if self._crossdisp_axis is None and self._image.data.ndim == 2:
+            self._crossdisp_axis = (self._disp_axis + 1) % 2
+        else:
+            raise ValueError('The cross dispersion axis must be given for for image cubes with ndim > 2.')
+
+    def __repr__(self) -> str:
+        return f"<Image unit={self.unit} shape={self.shape}>"
+
+    @staticmethod
+    def _parse_image(image) -> NDData:
+        """
+        Parse any of the supported image data types into NDData.
+
+        Parameters
+        ----------
+        image : Quantity, numpy.ndarray, NDData
+            The input image which can be of type Quantity, numpy.ndarray, or NDData.
+
+        Returns
+        -------
+        NDData
+            An NDData object containing the parsed data, with units and uncertainty.
+
+        Raises
+        ------
+        ValueError
+            If the image type is unrecognized.
+            If the image is fully masked or contains only invalid values.
+        """
+        if isinstance(image, u.quantity.Quantity):
+            data = image.value
+        elif isinstance(image, np.ndarray):
+            data = image
+        elif isinstance(image, NDData):
+            data = image.data
+        else:
+            raise ValueError('Unrecognized image type.')
+
+        if getattr(image, 'mask', None) is not None:
+            mask = image.mask | (~np.isfinite(data))
+        else:
+            mask =  ~np.isfinite(data)
+        if mask.all():
+            raise ValueError('Image is fully masked. Check for invalid values.')
+
+        unit = getattr(image, 'unit', u.Unit('DN'))
+        uncertainty = getattr(image, 'uncertainty', VarianceUncertainty(np.ones(data.shape)))
+        return NDData(data * unit, uncertainty=uncertainty, mask=mask)
+
+    @property
+    def unit(self) -> u.Unit:
+        return self._image.unit
+
+    @property
+    def shape(self) -> tuple:
+        return self._image.data.shape
+
+    @property
+    def data(self) -> np.ndarray:
+        return self._image.data
+
+    @property
+    def mask(self) -> np.ndarray:
+        return self._image.mask
+
+    def masked(self, mask_treatment: str = 'filter') -> NDData:
+        """
+        Get the NDData image with the given mask treatment applied.
+
+        Parameters
+        ----------
+        mask_treatment : str, optional
+            The method to be used for handling the mask in the image data.
+            Allowed values are 'filter', 'zero-fill', and 'omit'.
+            - 'filter': Leaves the masked data as is.
+            - 'zero-fill': Sets the values of the masked elements to zero and removes the mask.
+            - 'omit': Collapses the mask across the specified cross-dispersion axis.
+
+        Returns
+        -------
+        NDData
+            An image object with the specified mask treatment applied.
+
+        Raises
+        ------
+        ValueError
+            If the `mask_treatment` is not one of the implemented masking methods.
+        """
+        if mask_treatment not in self.implemented_masking_methods:
+            raise ValueError("`mask_treatment` must be one of "
+                             f"{self.implemented_masking_methods}")
+
+        image = deepcopy(self._image)
+        if mask_treatment == 'filter':
+            return image
+        elif mask_treatment == 'zero-fill':
+            image = deepcopy(self._image)
+            image.data[image.mask] = 0.
+            image.mask[:] = False
+            return image
+        elif mask_treatment == 'omit':
+            image = deepcopy(self._image)
+            image.mask[:] = np.any(image.mask, axis=self._crossdisp_axis, keepdims=True)
+            return image

From 7ca476da413bcff8ea7ee8f4f39276d7afb02dc1 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Tue, 5 Nov 2024 20:46:13 +0000
Subject: [PATCH 15/25] Added some sanity checks to the Image class.

---
 specreduce/image.py | 14 ++++++++++++++
 1 file changed, 14 insertions(+)

diff --git a/specreduce/image.py b/specreduce/image.py
index 2a5f742..0f6fc36 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -19,6 +19,16 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
         else:
             raise ValueError('The cross dispersion axis must be given for for image cubes with ndim > 2.')
 
+        if self._disp_axis == self._crossdisp_axis:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
+
+        if self._disp_axis < 0 or self._crossdisp_axis < 0:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be negative.')
+
+        if self._disp_axis >= self.ndim or self._crossdisp_axis >= self.ndim:
+            raise ValueError('The dispersion and cross-dispersion axes '
+                             'must be smaller than the number of image dimensions.')
+
     def __repr__(self) -> str:
         return f"<Image unit={self.unit} shape={self.shape}>"
 
@@ -71,6 +81,10 @@ def unit(self) -> u.Unit:
     def shape(self) -> tuple:
         return self._image.data.shape
 
+    @property
+    def ndim(self) -> int:
+        return self._image.data.ndim
+
     @property
     def data(self) -> np.ndarray:
         return self._image.data

From d2f4645eb22d3d78bbf223e55a89552487102efe Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Wed, 6 Nov 2024 19:47:37 +0000
Subject: [PATCH 16/25] Renamed Image class to SRImage and enhanced axis
 handling.

Renamed the `Image` class to `SRImage` for better clarity and added comprehensive validation for the dispersion and cross-dispersion axes. Introduced properties to provide reordered views of data, mask, and uncertainty arrays.
---
 specreduce/image.py | 117 ++++++++++++++++++++++++++++++++++++--------
 1 file changed, 96 insertions(+), 21 deletions(-)

diff --git a/specreduce/image.py b/specreduce/image.py
index 0f6fc36..f1b22d5 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -2,15 +2,23 @@
 
 import astropy.units as u
 import numpy as np
-
 from astropy.nddata import NDData, VarianceUncertainty
 
 
-class Image:
+class SRImage:
     implemented_masking_methods = 'filter', 'zero-fill', 'omit'
 
     def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None):
-        self._image: NDData = Image._parse_image(image)
+        self._image: NDData = SRImage._parse_image(image)
+
+        if disp_axis == crossdisp_axis:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
+        if disp_axis < 0 or crossdisp_axis < 0:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be negative.')
+        if disp_axis >= self.ndim or crossdisp_axis >= self.ndim:
+            raise ValueError('The dispersion and cross-dispersion axes '
+                             'must be smaller than the number of image dimensions.')
+
         self._disp_axis: int = disp_axis
         self._crossdisp_axis: int = crossdisp_axis
 
@@ -19,15 +27,23 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
         else:
             raise ValueError('The cross dispersion axis must be given for for image cubes with ndim > 2.')
 
-        if self._disp_axis == self._crossdisp_axis:
-            raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
+        # A view to the image data with cross-dispersion axis as the first dimension
+        # and dispersion axis as the second dimension.
+        self._st_data = np.moveaxis(self._image.data,
+                                    [self._crossdisp_axis, self._disp_axis],
+                                    [0, 1])
 
-        if self._disp_axis < 0 or self._crossdisp_axis < 0:
-            raise ValueError('The dispersion and cross-dispersion axes cannot be negative.')
+        # A view to the image mask with cross-dispersion axis as the first dimension
+        # and dispersion axis as the second dimension.
+        self._st_mask = np.moveaxis(self._image.mask,
+                                    [self._crossdisp_axis, self._disp_axis],
+                                    [0, 1])
 
-        if self._disp_axis >= self.ndim or self._crossdisp_axis >= self.ndim:
-            raise ValueError('The dispersion and cross-dispersion axes '
-                             'must be smaller than the number of image dimensions.')
+        # A view to the image mask with cross-dispersion axis as the first dimension
+        # and dispersion axis as the second dimension.
+        self._st_uncertainty = np.moveaxis(self._image.uncertainty,
+                                    [self._crossdisp_axis, self._disp_axis],
+                                    [0, 1])
 
     def __repr__(self) -> str:
         return f"<Image unit={self.unit} shape={self.shape}>"
@@ -73,6 +89,51 @@ def _parse_image(image) -> NDData:
         uncertainty = getattr(image, 'uncertainty', VarianceUncertainty(np.ones(data.shape)))
         return NDData(data * unit, uncertainty=uncertainty, mask=mask)
 
+    @property
+    def nddata(self) -> NDData:
+        return self._image
+
+    @property
+    def data(self) -> np.ndarray:
+        return self._image.data
+
+    @property
+    def ordered_data(self) -> np.ndarray:
+        """Image data arranged to a shape (crossdisp_axis, disp_axis).
+
+        A view to the data with cross-dispersion axis as the first dimension
+        and the dispersion axis as the second dimension.
+        """
+        return self._st_data
+
+    @property
+    def mask(self) -> np.ndarray:
+        return self._image.mask
+
+    @property
+    def ordered_mask(self) -> np.ndarray:
+        """
+        Image mask arranged to a shape (crossdisp_axis, disp_axis).
+
+        A view to the image mask with cross-dispersion axis as the first
+        dimension and the dispersion axis as the second dimension.
+        """
+        return self._st_mask
+
+    @property
+    def uncertainty(self):
+        """
+        Image uncertainty arranged to a shape (crossdisp_axis, disp_axis, ...).
+
+        A view to the image uncertainty with cross-dispersion axis as the first
+        dimension and the dispersion axis as the second dimension.
+        """
+        return self._image.uncertainty
+
+    @property
+    def ordered_uncertainty(self):
+        return self._st_uncertainty
+
     @property
     def unit(self) -> u.Unit:
         return self._image.unit
@@ -81,19 +142,35 @@ def unit(self) -> u.Unit:
     def shape(self) -> tuple:
         return self._image.data.shape
 
+    @property
+    def ordered_shape(self) -> tuple:
+        return self._st_data.shape
+
     @property
     def ndim(self) -> int:
         return self._image.data.ndim
 
     @property
-    def data(self) -> np.ndarray:
-        return self._image.data
+    def wcs(self):
+        return self._image.wcs
 
     @property
-    def mask(self) -> np.ndarray:
-        return self._image.mask
+    def meta(self):
+        return self._image.meta
+
+    @property
+    def disp_axis(self) -> int:
+        return self._disp_axis
+
+    @property
+    def crossdisp_axis(self) -> int:
+        return self._crossdisp_axis
+
+    def copy(self) -> 'SRImage':
+        """Copy the image object."""
+        return deepcopy(self)
 
-    def masked(self, mask_treatment: str = 'filter') -> NDData:
+    def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
         """
         Get the NDData image with the given mask treatment applied.
 
@@ -120,15 +197,13 @@ def masked(self, mask_treatment: str = 'filter') -> NDData:
             raise ValueError("`mask_treatment` must be one of "
                              f"{self.implemented_masking_methods}")
 
-        image = deepcopy(self._image)
+        image = self.copy()
         if mask_treatment == 'filter':
             return image
         elif mask_treatment == 'zero-fill':
-            image = deepcopy(self._image)
-            image.data[image.mask] = 0.
-            image.mask[:] = False
+            image._image.data[image.mask] = 0.
+            image._image.mask[:] = False
             return image
         elif mask_treatment == 'omit':
-            image = deepcopy(self._image)
-            image.mask[:] = np.any(image.mask, axis=self._crossdisp_axis, keepdims=True)
+            image._image.mask[:] = np.any(image.mask, axis=self._crossdisp_axis, keepdims=True)
             return image

From a314aa25758c35425ec29e1fcdd492e10cd1b745 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Wed, 6 Nov 2024 20:55:01 +0000
Subject: [PATCH 17/25] Renamed some SRImage properties.

---
 specreduce/image.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/specreduce/image.py b/specreduce/image.py
index f1b22d5..bfc9913 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -98,7 +98,7 @@ def data(self) -> np.ndarray:
         return self._image.data
 
     @property
-    def ordered_data(self) -> np.ndarray:
+    def cdata(self) -> np.ndarray:
         """Image data arranged to a shape (crossdisp_axis, disp_axis).
 
         A view to the data with cross-dispersion axis as the first dimension
@@ -111,7 +111,7 @@ def mask(self) -> np.ndarray:
         return self._image.mask
 
     @property
-    def ordered_mask(self) -> np.ndarray:
+    def cmask(self) -> np.ndarray:
         """
         Image mask arranged to a shape (crossdisp_axis, disp_axis).
 
@@ -131,7 +131,7 @@ def uncertainty(self):
         return self._image.uncertainty
 
     @property
-    def ordered_uncertainty(self):
+    def cuncertainty(self):
         return self._st_uncertainty
 
     @property
@@ -167,8 +167,8 @@ def crossdisp_axis(self) -> int:
         return self._crossdisp_axis
 
     def copy(self) -> 'SRImage':
-        """Copy the image object."""
-        return deepcopy(self)
+        """Copy the SRImage object."""
+        return SRImage(self._image, disp_axis=self._disp_axis, crossdisp_axis=self._crossdisp_axis)
 
     def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
         """

From ebf0cd64a8013374c8b1a8565c5d8c32006c06a1 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Thu, 7 Nov 2024 20:51:37 +0000
Subject: [PATCH 18/25] Refactored SRImage initialization and axis handling.

Redesigned the SRImage class to centralise image initialisation and sanity checks. Standardised axis handling across the class.
---
 specreduce/image.py | 175 ++++++++++++++++----------------------------
 1 file changed, 62 insertions(+), 113 deletions(-)

diff --git a/specreduce/image.py b/specreduce/image.py
index bfc9913..c404e39 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -1,74 +1,18 @@
-from copy import deepcopy
-
 import astropy.units as u
 import numpy as np
 from astropy.nddata import NDData, VarianceUncertainty
 
+CROSSDISP_AXIS: int = 0
+DISP_AXIS: int = 1
 
 class SRImage:
     implemented_masking_methods = 'filter', 'zero-fill', 'omit'
 
     def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None):
-        self._image: NDData = SRImage._parse_image(image)
-
-        if disp_axis == crossdisp_axis:
-            raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
-        if disp_axis < 0 or crossdisp_axis < 0:
-            raise ValueError('The dispersion and cross-dispersion axes cannot be negative.')
-        if disp_axis >= self.ndim or crossdisp_axis >= self.ndim:
-            raise ValueError('The dispersion and cross-dispersion axes '
-                             'must be smaller than the number of image dimensions.')
-
-        self._disp_axis: int = disp_axis
-        self._crossdisp_axis: int = crossdisp_axis
-
-        if self._crossdisp_axis is None and self._image.data.ndim == 2:
-            self._crossdisp_axis = (self._disp_axis + 1) % 2
-        else:
-            raise ValueError('The cross dispersion axis must be given for for image cubes with ndim > 2.')
-
-        # A view to the image data with cross-dispersion axis as the first dimension
-        # and dispersion axis as the second dimension.
-        self._st_data = np.moveaxis(self._image.data,
-                                    [self._crossdisp_axis, self._disp_axis],
-                                    [0, 1])
-
-        # A view to the image mask with cross-dispersion axis as the first dimension
-        # and dispersion axis as the second dimension.
-        self._st_mask = np.moveaxis(self._image.mask,
-                                    [self._crossdisp_axis, self._disp_axis],
-                                    [0, 1])
-
-        # A view to the image mask with cross-dispersion axis as the first dimension
-        # and dispersion axis as the second dimension.
-        self._st_uncertainty = np.moveaxis(self._image.uncertainty,
-                                    [self._crossdisp_axis, self._disp_axis],
-                                    [0, 1])
-
-    def __repr__(self) -> str:
-        return f"<Image unit={self.unit} shape={self.shape}>"
-
-    @staticmethod
-    def _parse_image(image) -> NDData:
-        """
-        Parse any of the supported image data types into NDData.
-
-        Parameters
-        ----------
-        image : Quantity, numpy.ndarray, NDData
-            The input image which can be of type Quantity, numpy.ndarray, or NDData.
-
-        Returns
-        -------
-        NDData
-            An NDData object containing the parsed data, with units and uncertainty.
+        self._nddata: NDData | None = None
 
-        Raises
-        ------
-        ValueError
-            If the image type is unrecognized.
-            If the image is fully masked or contains only invalid values.
-        """
+        # Extract the ndarray from the image object
+        # -----------------------------------------
         if isinstance(image, u.quantity.Quantity):
             data = image.value
         elif isinstance(image, np.ndarray):
@@ -76,8 +20,26 @@ def _parse_image(image) -> NDData:
         elif isinstance(image, NDData):
             data = image.data
         else:
-            raise ValueError('Unrecognized image type.')
+                raise ValueError('Unrecognized image type.')
 
+        # Carry out dispersion and cross-dispersion axis sanity checks
+        # ------------------------------------------------------------
+        if crossdisp_axis is None:
+            if data.ndim == 2:
+                crossdisp_axis = (disp_axis + 1) % 2
+            else:
+                raise ValueError('The cross-dispersion axis must be given for for image cubes with ndim > 2.')
+
+        if disp_axis == crossdisp_axis:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
+        if disp_axis < 0 or crossdisp_axis < 0:
+            raise ValueError('The dispersion and cross-dispersion axes cannot be negative.')
+        if disp_axis >= data.ndim or crossdisp_axis >= data.ndim:
+            raise ValueError('The dispersion and cross-dispersion axes '
+                             'must be smaller than the number of image dimensions.')
+
+        # Create the mask
+        # ---------------
         if getattr(image, 'mask', None) is not None:
             mask = image.mask | (~np.isfinite(data))
         else:
@@ -85,90 +47,77 @@ def _parse_image(image) -> NDData:
         if mask.all():
             raise ValueError('Image is fully masked. Check for invalid values.')
 
-        unit = getattr(image, 'unit', u.Unit('DN'))
-        uncertainty = getattr(image, 'uncertainty', VarianceUncertainty(np.ones(data.shape)))
-        return NDData(data * unit, uncertainty=uncertainty, mask=mask)
+        # Extract the unit and uncertainty
+        # --------------------------------
+        unit = getattr(image, 'unit', u.DN)
+        uncertainty = getattr(image, 'uncertainty', None) or VarianceUncertainty(np.ones_like(data))
 
-    @property
-    def nddata(self) -> NDData:
-        return self._image
+        # Standardise the axes
+        # --------------------
+        # This step forces the cross-dispersion axis as the first axis and the dispersion
+        # axis as the second axis. This could be done using transpose as well, but this
+        # approach works also with image cubes (although we're not supporting them yet).
+        data = np.moveaxis(data, [crossdisp_axis, disp_axis], [0, 1])
+        mask = np.moveaxis(mask, [crossdisp_axis, disp_axis], [0, 1])
+        uncertainty._array = np.moveaxis(uncertainty._array, [crossdisp_axis, disp_axis], [0, 1])
 
-    @property
-    def data(self) -> np.ndarray:
-        return self._image.data
+        self._nddata = NDData(data * unit, uncertainty=uncertainty, mask=mask)
 
-    @property
-    def cdata(self) -> np.ndarray:
-        """Image data arranged to a shape (crossdisp_axis, disp_axis).
 
-        A view to the data with cross-dispersion axis as the first dimension
-        and the dispersion axis as the second dimension.
-        """
-        return self._st_data
+    def __repr__(self) -> str:
+        return f"<Image unit={self.unit} shape={self.shape}>"
 
-    @property
-    def mask(self) -> np.ndarray:
-        return self._image.mask
 
     @property
-    def cmask(self) -> np.ndarray:
-        """
-        Image mask arranged to a shape (crossdisp_axis, disp_axis).
-
-        A view to the image mask with cross-dispersion axis as the first
-        dimension and the dispersion axis as the second dimension.
-        """
-        return self._st_mask
+    def nddata(self) -> NDData:
+        return self._nddata
 
     @property
-    def uncertainty(self):
-        """
-        Image uncertainty arranged to a shape (crossdisp_axis, disp_axis, ...).
+    def data(self) -> np.ndarray:
+        return self._nddata.data
 
-        A view to the image uncertainty with cross-dispersion axis as the first
-        dimension and the dispersion axis as the second dimension.
-        """
-        return self._image.uncertainty
+    @property
+    def mask(self) -> np.ndarray:
+        return self._nddata.mask
 
     @property
-    def cuncertainty(self):
-        return self._st_uncertainty
+    def uncertainty(self):
+        return self._nddata.uncertainty
 
     @property
     def unit(self) -> u.Unit:
-        return self._image.unit
+        return self._nddata.unit
 
     @property
     def shape(self) -> tuple:
-        return self._image.data.shape
-
-    @property
-    def ordered_shape(self) -> tuple:
-        return self._st_data.shape
+        return self.data.shape
 
     @property
     def ndim(self) -> int:
-        return self._image.data.ndim
+        return self.data.ndim
 
     @property
     def wcs(self):
-        return self._image.wcs
+        return self._nddata.wcs
 
     @property
     def meta(self):
-        return self._image.meta
+        return self._nddata.meta
 
     @property
     def disp_axis(self) -> int:
-        return self._disp_axis
+        return DISP_AXIS
 
     @property
     def crossdisp_axis(self) -> int:
-        return self._crossdisp_axis
+        return CROSSDISP_AXIS
+
+    def to_masked_array(self) -> np.ma.masked_array:
+        return np.ma.masked_array(self.data, mask=self.mask)
 
     def copy(self) -> 'SRImage':
         """Copy the SRImage object."""
-        return SRImage(self._image, disp_axis=self._disp_axis, crossdisp_axis=self._crossdisp_axis)
+        return SRImage(self._nddata, disp_axis=DISP_AXIS, crossdisp_axis=CROSSDISP_AXIS)
 
     def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
         """
@@ -201,9 +150,9 @@ def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
         if mask_treatment == 'filter':
             return image
         elif mask_treatment == 'zero-fill':
-            image._image.data[image.mask] = 0.
-            image._image.mask[:] = False
+            image._nddata.data[image.mask] = 0.
+            image._nddata.mask[:] = False
             return image
         elif mask_treatment == 'omit':
-            image._image.mask[:] = np.any(image.mask, axis=self._crossdisp_axis, keepdims=True)
+            image._nddata.mask[:] = np.any(image.mask, axis=CROSSDISP_AXIS, keepdims=True)
             return image

From 9a7d869faa7c0117b12ded66215ba1b7e8d2c760 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Thu, 7 Nov 2024 20:51:50 +0000
Subject: [PATCH 19/25] Added the first SRImage tests.

---
 specreduce/tests/test_srimage.py | 36 ++++++++++++++++++++++++++++++++
 1 file changed, 36 insertions(+)
 create mode 100644 specreduce/tests/test_srimage.py

diff --git a/specreduce/tests/test_srimage.py b/specreduce/tests/test_srimage.py
new file mode 100644
index 0000000..606301d
--- /dev/null
+++ b/specreduce/tests/test_srimage.py
@@ -0,0 +1,36 @@
+import astropy.units as u
+import numpy as np
+import pytest
+
+from astropy.nddata import NDData, VarianceUncertainty
+from specreduce.image import SRImage
+
+img = np.tile((np.arange(5, 15)), (7, 1)).astype('d')
+
+def test_init_ndarray():
+    image = SRImage(img)
+    image = SRImage(img, disp_axis=1)
+    image = SRImage(img, disp_axis=1, crossdisp_axis=0)
+    image = SRImage(img, disp_axis=0, crossdisp_axis=1)
+
+    with pytest.raises(ValueError):
+        image = SRImage(img, disp_axis=1, crossdisp_axis=1)
+
+    with pytest.raises(ValueError):
+        image = SRImage(img, disp_axis=2)
+
+    with pytest.raises(ValueError):
+        image = SRImage(img, disp_axis=-1)
+
+
+def test_init_quantity():
+    image = SRImage(img * u.DN, disp_axis=1)
+
+
+def test_init_nddata():
+    image = SRImage(NDData(img * u.DN), disp_axis=1)
+
+
+def test_init_bad():
+    with pytest.raises(ValueError, match='Unrecognized image type.'):
+        image = SRImage([[0.0, 1.0],[2.0, 3.0]])
\ No newline at end of file

From e5bc07b3f29b8b3b90d218c198bcef6084b0fef2 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hpparvi@gmail.com>
Date: Fri, 8 Nov 2024 19:56:51 +0000
Subject: [PATCH 20/25] Refactored `Background` class to use `SRImage` for
 image handling

Replaced `_ImageParser` with `SRImage` in `Background` class to standardize image handling. Added utility functions in `SRImage` for masking and subtraction, and updated related methods and tests to accommodate these changes.
---
 specreduce/background.py            | 105 ++++++++++------------------
 specreduce/image.py                 |  40 +++++++----
 specreduce/tests/test_background.py |   4 +-
 specreduce/tests/test_srimage.py    |  33 ++++++---
 4 files changed, 90 insertions(+), 92 deletions(-)

diff --git a/specreduce/background.py b/specreduce/background.py
index ac1f08c..07f0a4e 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -1,7 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import warnings
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, InitVar
 
 import numpy as np
 from astropy import units as u
@@ -9,15 +9,15 @@
 from astropy.utils.decorators import deprecated_attribute
 from specutils import Spectrum1D
 
-from specreduce.core import _ImageParser
 from specreduce.extract import _ap_weight_image
 from specreduce.tracing import Trace, FlatTrace
+from specreduce.image import SRImage, as_image
 
 __all__ = ['Background']
 
 
 @dataclass
-class Background(_ImageParser):
+class Background:
     """
     Determine the background from an image for subtraction.
 
@@ -65,19 +65,19 @@ class Background(_ImageParser):
     # https://stackoverflow.com/a/58409215
     __array_ufunc__ = None
 
-    image: NDData
+    image: SRImage
     traces: list = field(default_factory=list)
-    width: float = 5
+    width: float = 5.
+    disp_axis: InitVar[int] = 1
+    crossdisp_axis: InitVar[int | None] = None
     statistic: str = 'average'
-    disp_axis: int = 1
-    crossdisp_axis: int = 0
     mask_treatment: str = 'filter'
     _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
 
     # TO-DO: update bkg_array with Spectrum1D alternative (is bkg_image enough?)
     bkg_array = deprecated_attribute('bkg_array', '1.3')
 
-    def __post_init__(self):
+    def __post_init__(self, disp_axis, crossdisp_axis):
         """
         Determine the background from an image for subtraction.
 
@@ -110,27 +110,24 @@ def __post_init__(self):
             NDData object) will be combined with a mask generated from any non-finite
             values in the image data.
         """
-
-        # Parse image, including masked/nonfinite data handling based on
-        # choice of `mask_treatment`. Any uncaught nonfinte data values will be
-        # masked as well. Returns a Spectrum1D.
-        self.image = self._parse_image(self.image, disp_axis=self.disp_axis,
-                                       mask_treatment=self.mask_treatment)
+        self.image = as_image(self.image,
+                              disp_axis=disp_axis,
+                              crossdisp_axis=crossdisp_axis).apply_mask(self.mask_treatment)
 
         # always work with masked array, even if there is no masked
         # or nonfinite data, in case padding is needed. if not, mask will be
         # dropped at the end and a regular array will be returned.
-        img = np.ma.masked_array(self.image.data, self.image.mask)
+        img = self.image.to_masked_array()
 
         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.image.disp_axis])
             return
 
         self._set_traces()
 
-        bkg_wimage = np.zeros_like(self.image.data, dtype=np.float64)
+        bkg_wimage = np.zeros_like(self.image.flux, dtype=np.float64)
         for trace in self.traces:
             # note: ArrayTrace can have masked values, but if it does a MaskedArray
             # will be returned so this should be reflected in the window size here
@@ -138,9 +135,9 @@ def __post_init__(self):
             windows_max = trace.trace.data.max() + self.width/2
             windows_min = trace.trace.data.min() - self.width/2
 
-            if windows_max > self.image.shape[self.crossdisp_axis]:
+            if windows_max > self.image.shape[self.image.crossdisp_axis]:
                 warnings.warn("background window extends beyond image boundaries " +
-                              f"({windows_max} >= {self.image.shape[self.crossdisp_axis]})")
+                              f"({windows_max} >= {self.image.shape[self.image.crossdisp_axis]})")
             if windows_min < 0:
                 warnings.warn("background window extends beyond image boundaries " +
                               f"({windows_min} < 0)")
@@ -148,13 +145,13 @@ def __post_init__(self):
             # pass trace.trace.data to ignore any mask on the trace
             bkg_wimage += _ap_weight_image(trace,
                                            self.width,
-                                           self.disp_axis,
-                                           self.crossdisp_axis,
+                                           self.image.disp_axis,
+                                           self.image.crossdisp_axis,
                                            self.image.shape)
 
         if np.any(bkg_wimage > 1):
             raise ValueError("background regions overlapped")
-        if np.any(np.sum(bkg_wimage, axis=self.crossdisp_axis) == 0):
+        if np.any(np.sum(bkg_wimage, axis=self.image.crossdisp_axis) == 0):
             raise ValueError("background window does not remain in bounds across entire dispersion axis")  # noqa
         # check if image contained within background window is fully-nonfinite and raise an error
         if np.all(img.mask[bkg_wimage > 0]):
@@ -169,13 +166,12 @@ def __post_init__(self):
         if self.statistic == 'average':
             self._bkg_array = np.ma.average(img,
                                             weights=self.bkg_wimage,
-                                            axis=self.crossdisp_axis)
-
+                                            axis=self.image.crossdisp_axis)
         elif self.statistic == 'median':
             # combine where background weight image is 0 with image masked (which already
             # accounts for non-finite data that wasn't already masked)
             img.mask = np.logical_or(self.bkg_wimage == 0, self.image.mask)
-            self._bkg_array = np.ma.median(img, axis=self.crossdisp_axis)
+            self._bkg_array = np.ma.median(img, axis=self.image.crossdisp_axis)
         else:
             raise ValueError("statistic must be 'average' or 'median'")
 
@@ -189,8 +185,8 @@ def _set_traces(self):
 
         if self.traces == []:
             # assume a flat trace at the image center if nothing is passed in.
-            trace_pos = self.image.shape[self.disp_axis] / 2.
-            self.traces = [FlatTrace(self.image, trace_pos)]
+            trace_pos = self.image.shape[self.image.disp_axis] / 2.
+            self.traces = [FlatTrace(self.image.nddata, trace_pos)]
 
         if isinstance(self.traces, Trace):
             # if just one trace, turn it into iterable.
@@ -202,7 +198,7 @@ def _set_traces(self):
             self.traces = [self.traces]
 
         if np.all([isinstance(x, (float, int)) for x in self.traces]):
-            self.traces = [FlatTrace(self.image, trace_pos) for trace_pos in self.traces]
+            self.traces = [FlatTrace(self.image.nddata, trace_pos) for trace_pos in self.traces]
             return
 
         else:
@@ -256,8 +252,6 @@ def two_sided(cls, image, trace_object, separation, **kwargs):
             NDData object) will be combined with a mask generated from any non-finite
             values in the image data.
         """
-
-        image = _ImageParser._get_data_from_image(image) if image is not None else cls.image
         kwargs['traces'] = [trace_object-separation, trace_object+separation]
         return cls(image=image, **kwargs)
 
@@ -305,11 +299,10 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
             NDData object) will be combined with a mask generated from any non-finite
             values in the image data.
         """
-        image = _ImageParser._get_data_from_image(image) if image is not None else cls.image
         kwargs['traces'] = [trace_object+separation]
         return cls(image=image, **kwargs)
 
-    def bkg_image(self, image=None):
+    def bkg_image(self, image=None, disp_axis: int = 1, crossdisp_axis: int | None = None) -> SRImage:
         """
         Expose the background tiled to the dimension of ``image``.
 
@@ -323,14 +316,12 @@ def bkg_image(self, image=None):
 
         Returns
         -------
-        `~specutils.Spectrum1D` object with same shape as ``image``.
+        `~specutils.SRImage` object with same shape as ``image``.
         """
-        image = self._parse_image(image)
-        return Spectrum1D(np.tile(self._bkg_array,
-                                  (image.shape[0], 1)) * image.unit,
-                          spectral_axis=image.spectral_axis)
+        image = self.image if image is None else as_image(image, disp_axis, crossdisp_axis)
+        return SRImage(np.tile(self._bkg_array,(image.shape[image.crossdisp_axis], 1)) * image.unit)
 
-    def bkg_spectrum(self, image=None):
+    def bkg_spectrum(self, image=None) -> Spectrum1D:
         """
         Expose the 1D spectrum of the background.
 
@@ -349,17 +340,10 @@ 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)
-
-        try:
-            return bkg_image.collapse(np.nansum, 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.nansum(bkg_image.flux, axis=self.crossdisp_axis)
-            return Spectrum1D(ext1d, bkg_image.spectral_axis)
+        img = self.bkg_image(image)
+        return Spectrum1D(np.nansum(img.flux, axis=img.crossdisp_axis) * img.unit)
 
-    def sub_image(self, image=None):
+    def sub_image(self, image=None, disp_axis: int = 1, crossdisp_axis: int | None = None):
         """
         Subtract the computed background from ``image``.
 
@@ -371,20 +355,12 @@ def sub_image(self, image=None):
 
         Returns
         -------
-        `~specutils.Spectrum1D` object with same shape as ``image``.
+        `~specreduce.SRImage` object with same shape as ``image``.
         """
-        image = self._parse_image(image)
+        image = self.image if image is None else as_image(image, disp_axis, crossdisp_axis)
+        return image.subtract(self.bkg_image(image).nddata)
 
-        # 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):
+    def sub_spectrum(self, image=None) -> Spectrum1D:
         """
         Expose the 1D spectrum of the background-subtracted image.
 
@@ -402,14 +378,7 @@ def sub_spectrum(self, image=None):
             the spectral axis expressed in pixel units.
         """
         sub_image = self.sub_image(image=image)
-
-        try:
-            return sub_image.collapse(np.nansum, 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.nansum(sub_image.flux, axis=self.crossdisp_axis)
-            return Spectrum1D(ext1d, spectral_axis=sub_image.spectral_axis)
+        return Spectrum1D(np.nansum(sub_image.flux, axis=sub_image.crossdisp_axis) * sub_image.unit)
 
     def __rsub__(self, image):
         """
diff --git a/specreduce/image.py b/specreduce/image.py
index c404e39..743b5fc 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -1,15 +1,22 @@
 import astropy.units as u
 import numpy as np
-from astropy.nddata import NDData, VarianceUncertainty
+from astropy.nddata import NDData, NDDataRef, VarianceUncertainty
 
 CROSSDISP_AXIS: int = 0
 DISP_AXIS: int = 1
 
+def as_image(image, disp_axis: int = 1, crossdisp_axis: int | None = None):
+    if isinstance(image, SRImage):
+        return image
+    else:
+        return SRImage(image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
+
+
 class SRImage:
     implemented_masking_methods = 'filter', 'zero-fill', 'omit'
 
     def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None):
-        self._nddata: NDData | None = None
+        self._nddata: NDDataRef | None = None
 
         # Extract the ndarray from the image object
         # -----------------------------------------
@@ -20,7 +27,7 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
         elif isinstance(image, NDData):
             data = image.data
         else:
-                raise ValueError('Unrecognized image type.')
+            raise ValueError('Unrecognized image type.')
 
         # Carry out dispersion and cross-dispersion axis sanity checks
         # ------------------------------------------------------------
@@ -41,7 +48,7 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
         # Create the mask
         # ---------------
         if getattr(image, 'mask', None) is not None:
-            mask = image.mask | (~np.isfinite(data))
+            mask = image.mask.astype(bool) | (~np.isfinite(data))
         else:
             mask =  ~np.isfinite(data)
         if mask.all():
@@ -61,7 +68,7 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
         mask = np.moveaxis(mask, [crossdisp_axis, disp_axis], [0, 1])
         uncertainty._array = np.moveaxis(uncertainty._array, [crossdisp_axis, disp_axis], [0, 1])
 
-        self._nddata = NDData(data * unit, uncertainty=uncertainty, mask=mask)
+        self._nddata = NDDataRef(data * unit, uncertainty=uncertainty, mask=mask)
 
 
     def __repr__(self) -> str:
@@ -69,13 +76,17 @@ def __repr__(self) -> str:
 
 
     @property
-    def nddata(self) -> NDData:
+    def nddata(self) -> NDDataRef:
         return self._nddata
 
     @property
     def data(self) -> np.ndarray:
         return self._nddata.data
 
+    @property
+    def flux(self) -> np.ndarray:
+        return self._nddata.data
+
     @property
     def mask(self) -> np.ndarray:
         return self._nddata.mask
@@ -90,11 +101,11 @@ def unit(self) -> u.Unit:
 
     @property
     def shape(self) -> tuple:
-        return self.data.shape
+        return self.flux.shape
 
     @property
     def ndim(self) -> int:
-        return self.data.ndim
+        return self.flux.ndim
 
     @property
     def wcs(self):
@@ -112,16 +123,21 @@ def disp_axis(self) -> int:
     def crossdisp_axis(self) -> int:
         return CROSSDISP_AXIS
 
+    def subtract(self, image, propagate_uncertainties: bool = True) -> 'SRImage':
+        image = SRImage(image)
+        return SRImage(self._nddata.subtract(image.nddata,
+                                             propagate_uncertainties=propagate_uncertainties))
+
     def to_masked_array(self) -> np.ma.masked_array:
-        return np.ma.masked_array(self.data, mask=self.mask)
+        return np.ma.masked_array(self.flux, mask=self.mask)
 
     def copy(self) -> 'SRImage':
         """Copy the SRImage object."""
         return SRImage(self._nddata, disp_axis=DISP_AXIS, crossdisp_axis=CROSSDISP_AXIS)
 
-    def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
+    def apply_mask(self, mask_treatment: str = 'filter') -> 'SRImage':
         """
-        Get the NDData image with the given mask treatment applied.
+        Get the image with the given mask treatment applied.
 
         Parameters
         ----------
@@ -134,7 +150,7 @@ def masked(self, mask_treatment: str = 'filter') -> 'SRImage':
 
         Returns
         -------
-        NDData
+        SRImage
             An image object with the specified mask treatment applied.
 
         Raises
diff --git a/specreduce/tests/test_background.py b/specreduce/tests/test_background.py
index 388da5c..db3fe29 100644
--- a/specreduce/tests/test_background.py
+++ b/specreduce/tests/test_background.py
@@ -180,7 +180,7 @@ def test_fully_masked_column(self, mask):
         img[:, 0:1] = np.nan
 
         bkg = Background(img, traces=FlatTrace(img, 6), mask_treatment=mask)
-        assert np.all(bkg.bkg_image().data[:, 0:1] == 0.0)
+        assert np.all(bkg.bkg_image().flux[:, 0:1] == 0.0)
 
     @pytest.mark.parametrize("mask", ["filter", "omit"])
     def test_fully_masked_image(self, mask):
@@ -258,7 +258,7 @@ def test_mask_treatment_bkg_img_spectrum(self, method, expected):
             bk_img = background.bkg_image()
             # change this and following assertions to assert_quantity_allclose once
             # issue #213 is fixed
-            np.testing.assert_allclose(bk_img.flux.value,
+            np.testing.assert_allclose(bk_img.flux,
                                        np.tile(expected, (img_size, 1)))
 
             # test background spectrum matches 'expected' times the number of rows
diff --git a/specreduce/tests/test_srimage.py b/specreduce/tests/test_srimage.py
index 606301d..90a7a2d 100644
--- a/specreduce/tests/test_srimage.py
+++ b/specreduce/tests/test_srimage.py
@@ -5,13 +5,34 @@
 from astropy.nddata import NDData, VarianceUncertainty
 from specreduce.image import SRImage
 
+def create_image(imtype: str, nrow: int = 7, ncol: int = 11):
+    img = np.tile((np.arange(ncol)), (nrow, 1)).astype('d')
+    if imtype == 'ndarray':
+        return img
+    elif imtype == 'quantity':
+        return img * u.DN
+    elif imtype == 'nddata':
+        return NDData(img*u.DN)
+    else:
+        raise NotImplementedError('unkown image type')
+
 img = np.tile((np.arange(5, 15)), (7, 1)).astype('d')
 
-def test_init_ndarray():
+
+@pytest.mark.parametrize("imtype", ["ndarray", "quantity", 'nddata'])
+def test_init(imtype):
+    img = create_image(imtype)
     image = SRImage(img)
-    image = SRImage(img, disp_axis=1)
+    assert image.shape == (7, 11)
+    assert image.unit == u.DN
+
     image = SRImage(img, disp_axis=1, crossdisp_axis=0)
+    assert image.shape == (7, 11)
+    assert image.unit == u.DN
+
     image = SRImage(img, disp_axis=0, crossdisp_axis=1)
+    assert image.shape == (11, 7)
+    assert image.unit == u.DN
 
     with pytest.raises(ValueError):
         image = SRImage(img, disp_axis=1, crossdisp_axis=1)
@@ -23,14 +44,6 @@ def test_init_ndarray():
         image = SRImage(img, disp_axis=-1)
 
 
-def test_init_quantity():
-    image = SRImage(img * u.DN, disp_axis=1)
-
-
-def test_init_nddata():
-    image = SRImage(NDData(img * u.DN), disp_axis=1)
-
-
 def test_init_bad():
     with pytest.raises(ValueError, match='Unrecognized image type.'):
         image = SRImage([[0.0, 1.0],[2.0, 3.0]])
\ No newline at end of file

From 67ac82133984573440389df07f63735072485efc Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hannu@iac.es>
Date: Sat, 9 Nov 2024 13:51:01 +0000
Subject: [PATCH 21/25] Refactored the Tracing classes to use the `SRImage`
 class instead of `_ImageParser`.

---
 specreduce/tracing.py | 93 ++++++++++++++++++++-----------------------
 1 file changed, 44 insertions(+), 49 deletions(-)

diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index 9017cfd..f2bd7a6 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -2,7 +2,7 @@
 
 import warnings
 from copy import deepcopy
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, InitVar
 
 import numpy as np
 from astropy.modeling import Model, fitting, models
@@ -10,7 +10,7 @@
 from astropy.stats import gaussian_sigma_to_fwhm
 from astropy.utils.decorators import deprecated
 
-from specreduce.core import _ImageParser
+from specreduce.image import SRImage, as_image
 
 __all__ = ['Trace', 'FlatTrace', 'ArrayTrace', 'FitTrace']
 
@@ -30,38 +30,38 @@ class Trace:
     shape : tuple
         Shape of the array describing the trace
     """
-    image: NDData
+    image: SRImage
 
     def __post_init__(self):
-        self.trace_pos = self.image.shape[0] / 2
-        self.trace = np.ones_like(self.image[0]) * self.trace_pos
+        self.image = im = as_image(self.image)
+        self.trace_pos: float = im.shape[im.crossdisp_axis] / 2
+        self.trace: np.ma.MaskedArray = np.ma.masked_array(np.full(im.shape[im.disp_axis], self.trace_pos))
+        self._bound_trace()
 
         # masking options not relevant for basic Trace
         self._mask_treatment = None
         self._valid_mask_treatment_methods = [None]
-
-        # eventually move this to common SpecreduceOperation base class
         self.validate_masking_options()
 
-    def __getitem__(self, i):
+    def __getitem__(self, i: int):
         return self.trace[i]
 
     @property
-    def shape(self):
+    def shape(self) -> tuple:
         return self.trace.shape
 
-    def validate_masking_options(self):
+    def validate_masking_options(self) -> None:
         if self.mask_treatment not in self.valid_mask_treatment_methods:
             raise ValueError(
                 f'`mask_treatment` {self.mask_treatment} not one of {self.valid_mask_treatment_methods}')  # noqa
 
-    def shift(self, delta):
+    def shift(self, delta: float) -> None:
         """
         Shift the trace by delta pixels perpendicular to the axis being traced
 
         Parameters
         ----------
-        delta : float
+        delta
             Shift to be applied to the trace
         """
         # act on self.trace.data to ignore the mask and then re-mask when calling _bound_trace
@@ -72,10 +72,10 @@ def _bound_trace(self):
         """
         Mask trace positions that are outside the upper/lower bounds of the image.
         """
-        ny = self.image.shape[0]
+        ny = self.image.shape[self.image.crossdisp_axis]
         self.trace = np.ma.masked_outside(self.trace, 0, ny - 1)
 
-    def __add__(self, delta):
+    def __add__(self, delta: float) -> 'Trace':
         """
         Return a copy of the trace shifted "forward" by delta pixels perpendicular to the axis
         being traced
@@ -84,7 +84,7 @@ def __add__(self, delta):
         copy.shift(delta)
         return copy
 
-    def __sub__(self, delta):
+    def __sub__(self, delta: float) -> 'Trace':
         """
         Return a copy of the trace shifted "backward" by delta pixels perpendicular to the axis
         being traced
@@ -92,16 +92,16 @@ def __sub__(self, delta):
         return self.__add__(-delta)
 
     @property
-    def mask_treatment(self):
+    def mask_treatment(self) -> str | None:
         return self._mask_treatment
 
     @property
-    def valid_mask_treatment_methods(self):
+    def valid_mask_treatment_methods(self) -> list[str | None]:
         return self._valid_mask_treatment_methods
 
 
 @dataclass
-class FlatTrace(Trace, _ImageParser):
+class FlatTrace(Trace):
     """
     Trace that is constant along the axis being traced.
 
@@ -117,32 +117,31 @@ class FlatTrace(Trace, _ImageParser):
     trace_pos: float
 
     def __post_init__(self):
-        self.image = self._parse_image(self.image)
-
+        self.image = as_image(self.image)
         self.set_position(self.trace_pos)
 
         # masking options not relevant for basic Trace
         self._mask_treatment = None
         self._valid_mask_treatment_methods = [None]
 
-    def set_position(self, trace_pos):
+    def set_position(self, trace_pos: float) -> None:
         """
         Set the trace position within the image
 
         Parameters
         ----------
-        trace_pos : float
+        trace_pos
             Position of the trace
         """
         if trace_pos < 1:
             raise ValueError('`trace_pos` must be positive.')
         self.trace_pos = trace_pos
-        self.trace = np.ones_like(self.image.data[0]) * self.trace_pos
+        self.trace = np.ma.masked_array(np.full(self.image.shape[self.image.disp_axis], self.trace_pos))
         self._bound_trace()
 
 
 @dataclass
-class ArrayTrace(Trace, _ImageParser):
+class ArrayTrace(Trace):
     """
     Define a trace given an array of trace positions.
 
@@ -151,7 +150,7 @@ class ArrayTrace(Trace, _ImageParser):
     trace : `numpy.ndarray` or `numpy.ma.MaskedArray`
         Array containing trace positions.
     """
-    trace: np.ndarray
+    trace: np.ma.MaskedArray
 
     def __post_init__(self):
 
@@ -175,9 +174,9 @@ def __post_init__(self):
         # dropped at the end and a regular array will be returned.
         self.trace = np.ma.MaskedArray(trace_data, total_mask)
 
-        self.image = self._parse_image(self.image)
+        self.image = as_image(self.image)
 
-        nx = self.image.shape[1]
+        nx = self.image.shape[self.image.disp_axis]
         nt = len(self.trace)
         if nt != nx:
             if nt > nx:
@@ -201,7 +200,7 @@ def __post_init__(self):
 
 
 @dataclass
-class FitTrace(Trace, _ImageParser):
+class FitTrace(Trace):
     """
     Trace the spectrum aperture in an image.
 
@@ -276,25 +275,27 @@ class FitTrace(Trace, _ImageParser):
     window: int = None
     trace_model: Model = field(default=models.Polynomial1D(degree=1))
     peak_method: str = 'max'
-    _crossdisp_axis = 0
-    _disp_axis = 1
+    disp_axis: InitVar[int] = 1
+    crossdisp_axis: InitVar[int | None] = None
     mask_treatment: str = 'filter'
     _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
     # for testing purposes only, save bin peaks if requested
     _save_bin_peaks_testing: bool = False
 
-    def __post_init__(self):
+    def __post_init__(self, disp_axis, crossdisp_axis):
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. returns a Spectrum1D
-        self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
-                                       mask_treatment=self.mask_treatment)
+        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis).apply_mask(self.mask_treatment)
+        img = self.image.to_masked_array()
+        # self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
+        #                                mask_treatment=self.mask_treatment)
 
         # _parse_image returns a Spectrum1D. convert this to a masked array
         # for ease of calculations here (even if there is no masked data).
         # Note: uncertainties are dropped, this should also be addressed at
         # some point probably across the package.
-        img = np.ma.masked_array(self.image.data, self.image.mask)
+        #img = np.ma.masked_array(self.image.data, self.image.mask)
         self._mask_temp = self.image.mask
 
         # validate input arguments
@@ -302,19 +303,13 @@ def __post_init__(self):
         if self.peak_method not in valid_peak_methods:
             raise ValueError(f"peak_method must be one of {valid_peak_methods}")
 
-        if self._crossdisp_axis != 0:
-            raise ValueError('cross-dispersion axis must equal 0')
-
-        if self._disp_axis != 1:
-            raise ValueError('dispersion axis must equal 1')
-
         valid_models = (models.Spline1D, models.Legendre1D,
                         models.Chebyshev1D, models.Polynomial1D)
         if not isinstance(self.trace_model, valid_models):
             raise ValueError("trace_model must be one of "
                              f"{', '.join([m.name for m in valid_models])}.")
 
-        cols = img.shape[self._disp_axis]
+        cols = img.shape[self.image.disp_axis]
         model_deg = self.trace_model.degree
         if self.bins is None:
             self.bins = cols
@@ -333,7 +328,7 @@ def __post_init__(self):
             self.bins = int(self.bins)
 
         if (self.window is not None
-            and (self.window > img.shape[self._disp_axis]
+            and (self.window > img.shape[self.image.disp_axis]
                  or self.window < 1)):
             raise ValueError(f"window must be >= 2 and less than {cols}, the "
                              "length of the image's spatial direction")
@@ -346,10 +341,10 @@ def __post_init__(self):
 
     def _fit_trace(self, img):
 
-        yy = np.arange(img.shape[self._crossdisp_axis])
+        yy = np.arange(img.shape[self.image.crossdisp_axis])
 
         # set max peak location by user choice or wavelength with max avg flux
-        ztot = img.sum(axis=self._disp_axis) / img.shape[self._disp_axis]
+        ztot = img.sum(axis=self.image.disp_axis) / img.shape[self.image.disp_axis]
         peak_y = self.guess if self.guess is not None else ztot.argmax()
         # NOTE: peak finder can be bad if multiple objects are on slit
 
@@ -383,7 +378,7 @@ def _fit_trace(self, img):
             raise ValueError('All pixels in window region are masked. Check '
                              'for invalid values or use a larger window value.')
 
-        x_bins = np.linspace(0, img.shape[self._disp_axis],
+        x_bins = np.linspace(0, img.shape[self.image.disp_axis],
                              self.bins + 1, dtype=int)
         y_bins = np.tile(np.nan, self.bins)
 
@@ -392,9 +387,9 @@ def _fit_trace(self, img):
 
             # binned columns, summed along disp. axis.
             # or just a single, unbinned column if no bins
-            z_i = img[ilum2, x_bins[i]:x_bins[i + 1]].sum(axis=self._disp_axis)
+            z_i = img[ilum2, x_bins[i]:x_bins[i + 1]].sum(axis=self.image.disp_axis)
 
-            # if this bin is fully masked, set bin peak to NaN so it can be
+            # if this bin is fully masked, set bin peak to NaN, so it can be
             # filtered in the final fit to all bin peaks for the trace
             if z_i.mask.all():
                 warn_bins.append(i)
@@ -481,11 +476,11 @@ def _fit_trace(self, img):
             self._y_bins = y_bins
             self.trace_model_fit = fitter(self.trace_model, x_bins, y_bins)
 
-            trace_x = np.arange(img.shape[self._disp_axis])
+            trace_x = np.arange(img.shape[self.image.disp_axis])
             trace_y = self.trace_model_fit(trace_x)
         else:
             warnings.warn("TRACE ERROR: No valid points found in trace")
-            trace_y = np.tile(np.nan, img.shape[self._disp_axis])
+            trace_y = np.tile(np.nan, img.shape[self.image.disp_axis])
 
         self.trace = np.ma.masked_invalid(trace_y)
 

From 63631289f5254f0cb6e24e0ff0f0ebb6ced989d9 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hannu@iac.es>
Date: Sun, 10 Nov 2024 11:21:30 +0100
Subject: [PATCH 22/25] Added `size_disp_axis` and `size_crossdisp_axis`
 properties to the `SRImage` class to simplify dimension handling. Refactored
 trace calculations in `tracing.py` to utilize these new properties.

---
 specreduce/image.py   | 13 +++++++++++--
 specreduce/tracing.py | 33 ++++++++++++++++-----------------
 2 files changed, 27 insertions(+), 19 deletions(-)

diff --git a/specreduce/image.py b/specreduce/image.py
index 743b5fc..4659029 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -123,13 +123,22 @@ def disp_axis(self) -> int:
     def crossdisp_axis(self) -> int:
         return CROSSDISP_AXIS
 
+    @property
+    def size_disp_axis(self):
+        return self.shape[self.disp_axis]
+
+    @property
+    def size_crossdisp_axis(self):
+        return self.shape[self.crossdisp_axis]
+
     def subtract(self, image, propagate_uncertainties: bool = True) -> 'SRImage':
         image = SRImage(image)
         return SRImage(self._nddata.subtract(image.nddata,
                                              propagate_uncertainties=propagate_uncertainties))
 
-    def to_masked_array(self) -> np.ma.masked_array:
-        return np.ma.masked_array(self.flux, mask=self.mask)
+    def to_masked_array(self, mask_treatment: str = 'filter') -> np.ma.MaskedArray:
+        image = self.apply_mask(mask_treatment)
+        return np.ma.masked_array(image.flux, mask=image.mask)
 
     def copy(self) -> 'SRImage':
         """Copy the SRImage object."""
diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index f2bd7a6..900b785 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -6,11 +6,10 @@
 
 import numpy as np
 from astropy.modeling import Model, fitting, models
-from astropy.nddata import NDData
 from astropy.stats import gaussian_sigma_to_fwhm
 from astropy.utils.decorators import deprecated
 
-from specreduce.image import SRImage, as_image
+from specreduce.image import SRImage, as_image, DISP_AXIS, CROSSDISP_AXIS
 
 __all__ = ['Trace', 'FlatTrace', 'ArrayTrace', 'FitTrace']
 
@@ -34,8 +33,8 @@ class Trace:
 
     def __post_init__(self):
         self.image = im = as_image(self.image)
-        self.trace_pos: float = im.shape[im.crossdisp_axis] / 2
-        self.trace: np.ma.MaskedArray = np.ma.masked_array(np.full(im.shape[im.disp_axis], self.trace_pos))
+        self.trace_pos: float = im.size_crossdisp_axis / 2
+        self.trace: np.ma.MaskedArray = np.ma.masked_array(np.full(im.size_disp_axis, self.trace_pos))
         self._bound_trace()
 
         # masking options not relevant for basic Trace
@@ -72,7 +71,7 @@ def _bound_trace(self):
         """
         Mask trace positions that are outside the upper/lower bounds of the image.
         """
-        ny = self.image.shape[self.image.crossdisp_axis]
+        ny = self.image.size_crossdisp_axis
         self.trace = np.ma.masked_outside(self.trace, 0, ny - 1)
 
     def __add__(self, delta: float) -> 'Trace':
@@ -136,7 +135,7 @@ def set_position(self, trace_pos: float) -> None:
         if trace_pos < 1:
             raise ValueError('`trace_pos` must be positive.')
         self.trace_pos = trace_pos
-        self.trace = np.ma.masked_array(np.full(self.image.shape[self.image.disp_axis], self.trace_pos))
+        self.trace = np.ma.masked_array(np.full(self.image.size_disp_axis, self.trace_pos))
         self._bound_trace()
 
 
@@ -176,7 +175,7 @@ def __post_init__(self):
 
         self.image = as_image(self.image)
 
-        nx = self.image.shape[self.image.disp_axis]
+        nx = self.image.size_disp_axis
         nt = len(self.trace)
         if nt != nx:
             if nt > nx:
@@ -286,8 +285,8 @@ def __post_init__(self, disp_axis, crossdisp_axis):
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. returns a Spectrum1D
-        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis).apply_mask(self.mask_treatment)
-        img = self.image.to_masked_array()
+        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
+        img = self.image.to_masked_array(mask_treatment=self.mask_treatment)
         # self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
         #                                mask_treatment=self.mask_treatment)
 
@@ -309,7 +308,7 @@ def __post_init__(self, disp_axis, crossdisp_axis):
             raise ValueError("trace_model must be one of "
                              f"{', '.join([m.name for m in valid_models])}.")
 
-        cols = img.shape[self.image.disp_axis]
+        cols = img.shape[DISP_AXIS]
         model_deg = self.trace_model.degree
         if self.bins is None:
             self.bins = cols
@@ -328,7 +327,7 @@ def __post_init__(self, disp_axis, crossdisp_axis):
             self.bins = int(self.bins)
 
         if (self.window is not None
-            and (self.window > img.shape[self.image.disp_axis]
+            and (self.window > img.shape[DISP_AXIS]
                  or self.window < 1)):
             raise ValueError(f"window must be >= 2 and less than {cols}, the "
                              "length of the image's spatial direction")
@@ -341,10 +340,10 @@ def __post_init__(self, disp_axis, crossdisp_axis):
 
     def _fit_trace(self, img):
 
-        yy = np.arange(img.shape[self.image.crossdisp_axis])
+        yy = np.arange(img.shape[CROSSDISP_AXIS])
 
         # set max peak location by user choice or wavelength with max avg flux
-        ztot = img.sum(axis=self.image.disp_axis) / img.shape[self.image.disp_axis]
+        ztot = img.sum(axis=DISP_AXIS) / img.shape[DISP_AXIS]
         peak_y = self.guess if self.guess is not None else ztot.argmax()
         # NOTE: peak finder can be bad if multiple objects are on slit
 
@@ -378,7 +377,7 @@ def _fit_trace(self, img):
             raise ValueError('All pixels in window region are masked. Check '
                              'for invalid values or use a larger window value.')
 
-        x_bins = np.linspace(0, img.shape[self.image.disp_axis],
+        x_bins = np.linspace(0, img.shape[DISP_AXIS],
                              self.bins + 1, dtype=int)
         y_bins = np.tile(np.nan, self.bins)
 
@@ -387,7 +386,7 @@ def _fit_trace(self, img):
 
             # binned columns, summed along disp. axis.
             # or just a single, unbinned column if no bins
-            z_i = img[ilum2, x_bins[i]:x_bins[i + 1]].sum(axis=self.image.disp_axis)
+            z_i = img[ilum2, x_bins[i]:x_bins[i + 1]].sum(axis=DISP_AXIS)
 
             # if this bin is fully masked, set bin peak to NaN, so it can be
             # filtered in the final fit to all bin peaks for the trace
@@ -476,11 +475,11 @@ def _fit_trace(self, img):
             self._y_bins = y_bins
             self.trace_model_fit = fitter(self.trace_model, x_bins, y_bins)
 
-            trace_x = np.arange(img.shape[self.image.disp_axis])
+            trace_x = np.arange(img.shape[DISP_AXIS])
             trace_y = self.trace_model_fit(trace_x)
         else:
             warnings.warn("TRACE ERROR: No valid points found in trace")
-            trace_y = np.tile(np.nan, img.shape[self.image.disp_axis])
+            trace_y = np.tile(np.nan, img.shape[DISP_AXIS])
 
         self.trace = np.ma.masked_invalid(trace_y)
 

From 3d6cf898dfb261e7d5fc179c7daff47db539bd29 Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hannu@iac.es>
Date: Mon, 11 Nov 2024 12:58:29 +0100
Subject: [PATCH 23/25] Finalising the image handling refactoring for the Trace
 and Extract classes.

---
 specreduce/extract.py | 57 +++++++++++++++++++--------------
 specreduce/image.py   | 74 ++++++++++++++++++++++++++++++++++++-------
 specreduce/tracing.py | 57 ++++++++++++++-------------------
 3 files changed, 119 insertions(+), 69 deletions(-)

diff --git a/specreduce/extract.py b/specreduce/extract.py
index 00d6896..1f35352 100644
--- a/specreduce/extract.py
+++ b/specreduce/extract.py
@@ -1,7 +1,7 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
 import warnings
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, InitVar
 
 import numpy as np
 from astropy import units as u
@@ -12,6 +12,7 @@
 from specutils import Spectrum1D
 
 from specreduce.core import SpecreduceOperation
+from specreduce.image import SRImage, as_image, DISP_AXIS, CROSSDISP_AXIS
 from specreduce.tracing import Trace, FlatTrace
 
 __all__ = ['BoxcarExtract', 'HorneExtract', 'OptimalExtract']
@@ -44,7 +45,7 @@ def _get_boxcar_weights(center, hwidth, npix):
         A 2D image with weights assigned to pixels that fall within the
         defined aperture.
     """
-    weights = np.zeros((npix))
+    weights = np.zeros(npix)
     if hwidth == 0:
         # the logic below would return all zeros anyways, so might as well save the time
         # (negative widths should be avoided by earlier logic!)
@@ -152,7 +153,7 @@ class BoxcarExtract(SpecreduceOperation):
     spec : `~specutils.Spectrum1D`
         The extracted 1d spectrum expressed in DN and pixel units
     """
-    image: NDData
+    image: SRImage
     trace_object: Trace
     width: float = 5
     disp_axis: int = 1
@@ -166,7 +167,8 @@ def spectrum(self):
         return self.__call__()
 
     def __call__(self, image=None, trace_object=None, width=None,
-                 disp_axis=None, crossdisp_axis=None):
+                 disp_axis=None, crossdisp_axis=None,
+                 mask_treatment: str = 'filter') -> Spectrum1D:
         """
         Extract the 1D spectrum using the boxcar method.
 
@@ -204,7 +206,6 @@ def __call__(self, image=None, trace_object=None, width=None,
             The extracted 1d spectrum with flux expressed in the same
             units as the input image, or u.DN, and pixel units
         """
-        image = image if image is not None else self.image
         trace_object = trace_object if trace_object is not None else self.trace_object
         width = width if width is not None else self.width
         disp_axis = disp_axis if disp_axis is not None else self.disp_axis
@@ -213,8 +214,11 @@ def __call__(self, image=None, trace_object=None, width=None,
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
         # omit. non-finite data will be masked, always. Returns a Spectrum1D.
-        self.image = self._parse_image(image, disp_axis=disp_axis,
-                                       mask_treatment=self.mask_treatment)
+        # self.image = self._parse_image(image, disp_axis=disp_axis,
+        #                                mask_treatment=self.mask_treatment)
+        # TODO: Mask treatment
+        self.image = as_image(image if image is not None else self.image,
+                              self.disp_axis, self.crossdisp_axis)
 
         # # _parse_image returns a Spectrum1D. convert this to a masked array
         # # for ease of calculations here (even if there is no masked data).
@@ -226,16 +230,15 @@ def __call__(self, image=None, trace_object=None, width=None,
         # weight image to use for extraction
         wimg = _ap_weight_image(trace_object,
                                 width,
-                                disp_axis,
-                                crossdisp_axis,
+                                self.image.disp_axis,
+                                self.image.crossdisp_axis,
                                 self.image.shape)
 
         # extract, assigning no weight to non-finite pixels outside the window
         # (non-finite pixels inside the window will still make it into the sum)
         image_windowed = np.where(wimg, self.image.data*wimg, 0)
-        ext1d = np.sum(image_windowed, axis=crossdisp_axis)
-        return Spectrum1D(ext1d * self.image.unit,
-                          spectral_axis=self.image.spectral_axis)
+        ext1d = np.sum(image_windowed, axis=self.image.crossdisp_axis)
+        return Spectrum1D(ext1d * self.image.unit)
 
 
 @dataclass
@@ -314,16 +317,18 @@ class HorneExtract(SpecreduceOperation):
         fluxes are interpreted in DN. [default: None]
 
     """
-    image: NDData
+    image: SRImage
     trace_object: Trace
     bkgrd_prof: Model = field(default=models.Polynomial1D(2))
-    spatial_profile: str = 'gaussian'  # can actually be str, dict
-    variance: np.ndarray = field(default=None)
-    mask: np.ndarray = field(default=None)
-    unit: np.ndarray = field(default=None)
-    disp_axis: int = 1
-    crossdisp_axis: int = 0
-    # TODO: should disp_axis and crossdisp_axis be defined in the Trace object?
+    spatial_profile: str | dict = 'gaussian'  # can actually be str, dict
+    variance: np.ndarray | None = field(default=None)
+    mask: np.ndarray | None = field(default=None)
+    unit: u.Unit | None = field(default=None)
+    disp_axis: InitVar[int] = 1
+    crossdisp_axis: InitVar[int | None] = None
+
+    def __post_init__(self, disp_axis, crossdisp_axis):
+        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
 
     @property
     def spectrum(self):
@@ -661,7 +666,14 @@ def __call__(self, image=None, trace_object=None,
             raise ValueError('``spatial_profile`` must either be string or dictionary.')
 
         # parse image and replace optional arguments with updated values
-        self.image = self._parse_image(image, variance, mask, unit, disp_axis)
+        self.image = SRImage(image,
+                             disp_axis=disp_axis,
+                             unit=unit,
+                             mask=mask,
+                             variance=variance,
+                             ensure_var_uncertainty=True,
+                             require_uncertainty=True)
+        #self.image = self._parse_image(image, variance, mask, unit, disp_axis)
         variance = self.image.uncertainty.array
         mask = self.image.mask
         unit = self.image.unit
@@ -805,8 +817,7 @@ def __call__(self, image=None, trace_object=None,
         extraction = result * norms
 
         # convert the extraction to a Spectrum1D object
-        return Spectrum1D(extraction * unit,
-                          spectral_axis=self.image.spectral_axis)
+        return Spectrum1D(extraction * unit)
 
 
 def _align_along_trace(img, trace_array, disp_axis=1, crossdisp_axis=0):
diff --git a/specreduce/image.py b/specreduce/image.py
index 4659029..863790f 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -1,3 +1,5 @@
+import warnings
+
 import astropy.units as u
 import numpy as np
 from astropy.nddata import NDData, NDDataRef, VarianceUncertainty
@@ -5,6 +7,7 @@
 CROSSDISP_AXIS: int = 0
 DISP_AXIS: int = 1
 
+
 def as_image(image, disp_axis: int = 1, crossdisp_axis: int | None = None):
     if isinstance(image, SRImage):
         return image
@@ -15,7 +18,11 @@ def as_image(image, disp_axis: int = 1, crossdisp_axis: int | None = None):
 class SRImage:
     implemented_masking_methods = 'filter', 'zero-fill', 'omit'
 
-    def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None):
+    def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None,
+                 unit: u.Unit | None = None,
+                 mask: np.ndarray | None = None, mask_nonfinite: bool = True,
+                 variance: np.ndarray | None = None, ensure_var_uncertainty: bool = False,
+                 require_uncertainty: bool = False) -> None:
         self._nddata: NDDataRef | None = None
 
         # Extract the ndarray from the image object
@@ -35,7 +42,8 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
             if data.ndim == 2:
                 crossdisp_axis = (disp_axis + 1) % 2
             else:
-                raise ValueError('The cross-dispersion axis must be given for for image cubes with ndim > 2.')
+                raise ValueError('The cross-dispersion axis must be given '
+                                 ' for for image cubes with ndim > 2.')
 
         if disp_axis == crossdisp_axis:
             raise ValueError('The dispersion and cross-dispersion axes cannot be the same.')
@@ -47,21 +55,65 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
 
         # Create the mask
         # ---------------
+        # TODO: Shouldn't the user-given mask override the default mask in the image?
+        #       The order of tests should probably be changed.
         if getattr(image, 'mask', None) is not None:
-            mask = image.mask.astype(bool) | (~np.isfinite(data))
+            mask = image.mask.astype(bool)
+        elif mask is not None:
+            pass
         else:
-            mask =  ~np.isfinite(data)
+            mask = np.zeros(data.shape, dtype=bool)
+
+        if mask_nonfinite:
+            mask |= ~np.isfinite(data)
+
         if mask.all():
             raise ValueError('Image is fully masked. Check for invalid values.')
-
-        # Extract the unit and uncertainty
-        # --------------------------------
-        unit = getattr(image, 'unit', u.DN)
-        uncertainty = getattr(image, 'uncertainty', None) or VarianceUncertainty(np.ones_like(data))
+        if data.shape != mask.shape:
+            raise ValueError('Image and mask shapes must match.')
+
+        # Extract the unit
+        # ----------------
+        if unit is None:
+            unit = getattr(image, 'unit', u.DN)
+
+        # Extract the uncertainty
+        # -----------------------
+        if variance is not None:
+            uncertainty = VarianceUncertainty(variance)
+        elif getattr(image, 'uncertainty', None) is not None:
+            uncertainty = image.uncertainty
+        elif not require_uncertainty:
+            uncertainty = VarianceUncertainty(np.ones_like(data))
+        else:
+            raise ValueError('Uncertainty information is required but missing.')
+
+        # Try to convert the uncertainty into VarianceUncertainty
+        if ensure_var_uncertainty:
+            if uncertainty.uncertainty_type != 'var':
+                warnings.warn("image NDData object's uncertainty is not"
+                              "given as VarianceUncertainty. Trying to "
+                              "convert the uncertainty to VarianceUncertainty.")
+                uncertainty = uncertainty.represent_as(VarianceUncertainty)
+
+        if uncertainty.array.shape != data.shape:
+            raise ValueError('Image and uncertainty shapes must match.')
+        if np.any(uncertainty.array < 0):
+            raise ValueError("Uncertainty must be fully positive")
+        if np.all(uncertainty.array == 0):
+            # technically would result in infinities, but since they're all
+            # zeros, we can override ones to simulate an unweighted case
+            uncertainty._array[:] = 1.0
+        if np.any(uncertainty.array == 0):
+            m = uncertainty.array == 0
+            # exclude such elements by editing the input mask
+            mask[m] = True
+            # replace the variances to avoid a divide by zero warning
+            uncertainty._array[m] = np.nan
 
         # Standardise the axes
         # --------------------
-        # This step forces the cross-dispersion axis as the first axis and the dispersion
+        # Force the cross-dispersion axis as the first axis and the dispersion
         # axis as the second axis. This could be done using transpose as well, but this
         # approach works also with image cubes (although we're not supporting them yet).
         data = np.moveaxis(data, [crossdisp_axis, disp_axis], [0, 1])
@@ -70,11 +122,9 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None)
 
         self._nddata = NDDataRef(data * unit, uncertainty=uncertainty, mask=mask)
 
-
     def __repr__(self) -> str:
         return f"<Image unit={self.unit} shape={self.shape}>"
 
-
     @property
     def nddata(self) -> NDDataRef:
         return self._nddata
diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index 900b785..c5f18d9 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -30,17 +30,19 @@ class Trace:
         Shape of the array describing the trace
     """
     image: SRImage
+    disp_axis: InitVar[int] = field(default=1, kw_only=True)
+    crossdisp_axis: InitVar[int | None] = field(default=None, kw_only=True)
+    trace: np.ma.MaskedArray | None = field(default=None, init=False, repr=False)
+    mask_treatment: str | None = field(default=None, init=False, repr=False, kw_only=True)
+    _valid_mask_treatment_methods: tuple[str | None] = field(default=(None,), init=False, repr=False)
 
-    def __post_init__(self):
-        self.image = im = as_image(self.image)
-        self.trace_pos: float = im.size_crossdisp_axis / 2
-        self.trace: np.ma.MaskedArray = np.ma.masked_array(np.full(im.size_disp_axis, self.trace_pos))
-        self._bound_trace()
-
-        # masking options not relevant for basic Trace
-        self._mask_treatment = None
-        self._valid_mask_treatment_methods = [None]
+    def __post_init__(self, disp_axis, crossdisp_axis):
+        self.image = im = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
         self.validate_masking_options()
+        if self.trace is None:
+            trace_pos = im.size_crossdisp_axis / 2
+            self.trace = np.ma.masked_array(np.full(im.size_disp_axis, trace_pos))
+            self._bound_trace()
 
     def __getitem__(self, i: int):
         return self.trace[i]
@@ -91,11 +93,7 @@ def __sub__(self, delta: float) -> 'Trace':
         return self.__add__(-delta)
 
     @property
-    def mask_treatment(self) -> str | None:
-        return self._mask_treatment
-
-    @property
-    def valid_mask_treatment_methods(self) -> list[str | None]:
+    def valid_mask_treatment_methods(self) -> tuple[str]:
         return self._valid_mask_treatment_methods
 
 
@@ -115,14 +113,10 @@ class FlatTrace(Trace):
     """
     trace_pos: float
 
-    def __post_init__(self):
-        self.image = as_image(self.image)
+    def __post_init__(self, disp_axis, crossdisp_axis):
+        super().__post_init__(disp_axis, crossdisp_axis)
         self.set_position(self.trace_pos)
 
-        # masking options not relevant for basic Trace
-        self._mask_treatment = None
-        self._valid_mask_treatment_methods = [None]
-
     def set_position(self, trace_pos: float) -> None:
         """
         Set the trace position within the image
@@ -151,12 +145,8 @@ class ArrayTrace(Trace):
     """
     trace: np.ma.MaskedArray
 
-    def __post_init__(self):
-
-        # masking options not relevant for ArrayTrace. any non-finite or masked
-        # data in `image` will not affect array trace
-        self._mask_treatment = None
-        self._valid_mask_treatment_methods = [None]
+    def __post_init__(self, disp_axis, crossdisp_axis):
+        super().__post_init__(disp_axis, crossdisp_axis)
 
         # masked array will have a .data, regular array will not.
         trace_data = getattr(self.trace, 'data', self.trace)
@@ -248,7 +238,7 @@ class FitTrace(Trace):
         One of ``gaussian``, ``centroid``, or ``max``.
         ``gaussian``: Fits a gaussian to the window within each bin and
         adopts the central value as the peak. May work best with fewer
-        bins on faint targets. (Based on the "kosmos" algorithm from
+        bins on faint targets. (Based on the "Kosmos" algorithm from
         James Davenport's same-named repository.)
         ``centroid``: Takes the centroid of the window within in bin.
         ``max``: Saves the position with the maximum flux in each bin.
@@ -274,18 +264,17 @@ class FitTrace(Trace):
     window: int = None
     trace_model: Model = field(default=models.Polynomial1D(degree=1))
     peak_method: str = 'max'
-    disp_axis: InitVar[int] = 1
-    crossdisp_axis: InitVar[int | None] = None
     mask_treatment: str = 'filter'
-    _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
+    _valid_mask_treatment_methods: tuple[str | None] = field(default=('filter', 'omit', 'zero-fill'), init=False, repr=False)
     # for testing purposes only, save bin peaks if requested
     _save_bin_peaks_testing: bool = False
 
     def __post_init__(self, disp_axis, crossdisp_axis):
+        super().__post_init__(disp_axis, crossdisp_axis)
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`. returns a Spectrum1D
-        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
+        # self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
         img = self.image.to_masked_array(mask_treatment=self.mask_treatment)
         # self.image = self._parse_image(self.image, disp_axis=self._disp_axis,
         #                                mask_treatment=self.mask_treatment)
@@ -294,7 +283,7 @@ def __post_init__(self, disp_axis, crossdisp_axis):
         # for ease of calculations here (even if there is no masked data).
         # Note: uncertainties are dropped, this should also be addressed at
         # some point probably across the package.
-        #img = np.ma.masked_array(self.image.data, self.image.mask)
+        # img = np.ma.masked_array(self.image.data, self.image.mask)
         self._mask_temp = self.image.mask
 
         # validate input arguments
@@ -431,7 +420,7 @@ def _fit_trace(self, img):
                 y_bins[i] = np.interp(z_i_cumsum[-1] / 2., z_i_cumsum, ilum2)
 
                 # NOTE this reflects current behavior, should eventually be changed
-                # to set to nan by default (or zero fill / interpoate option once
+                # to set to nan by default (or zero fill / interpolate option once
                 # available)
 
             elif self.peak_method == 'max':
@@ -439,7 +428,7 @@ def _fit_trace(self, img):
                 y_bins[i] = ilum2[z_i.argmax()]
 
                 # NOTE: a fully masked should eventually be changed to set to
-                # nan by default (or zero fill / interpoate option once available)
+                # nan by default (or zero fill / interpolate option once available)
 
         # warn about fully-masked bins
         if len(warn_bins) > 0:

From b768d0c45cc1a87848066f27f12b940457d2678c Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hannu@iac.es>
Date: Mon, 11 Nov 2024 19:29:06 +0100
Subject: [PATCH 24/25] Improved the 'SRImage' class uncertainty handling.

---
 specreduce/image.py | 82 ++++++++++++++++++++++++++++++---------------
 1 file changed, 55 insertions(+), 27 deletions(-)

diff --git a/specreduce/image.py b/specreduce/image.py
index 863790f..18d9287 100644
--- a/specreduce/image.py
+++ b/specreduce/image.py
@@ -2,17 +2,28 @@
 
 import astropy.units as u
 import numpy as np
-from astropy.nddata import NDData, NDDataRef, VarianceUncertainty
+from astropy.nddata import (NDData, NDDataRef, VarianceUncertainty, NDUncertainty,
+                            StdDevUncertainty, InverseVariance)
 
 CROSSDISP_AXIS: int = 0
 DISP_AXIS: int = 1
 
 
-def as_image(image, disp_axis: int = 1, crossdisp_axis: int | None = None):
+def as_image(image, disp_axis: int = 1, crossdisp_axis: int | None = None,
+             unit: u.Unit | None = None,
+             mask: np.ndarray | None = None, mask_nonfinite: bool = True,
+             uncertainty: [np.ndarray | NDUncertainty | None] = None,
+             uncertainty_type: str = 'var',
+             ensure_var_uncertainty: bool = False,
+             require_uncertainty: bool = False):
     if isinstance(image, SRImage):
         return image
     else:
-        return SRImage(image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
+        return SRImage(image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis,
+                       unit=unit, mask=mask, mask_nonfinite=mask_nonfinite,
+                       uncertainty=uncertainty, uncertainty_type=uncertainty_type,
+                       ensure_var_uncertainty=ensure_var_uncertainty,
+                       require_uncertainty=require_uncertainty)
 
 
 class SRImage:
@@ -21,13 +32,18 @@ class SRImage:
     def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None,
                  unit: u.Unit | None = None,
                  mask: np.ndarray | None = None, mask_nonfinite: bool = True,
-                 variance: np.ndarray | None = None, ensure_var_uncertainty: bool = False,
+                 uncertainty: [np.ndarray | NDUncertainty | None] = None,
+                 uncertainty_type: str = 'var',
+                 ensure_var_uncertainty: bool = False,
                  require_uncertainty: bool = False) -> None:
         self._nddata: NDDataRef | None = None
 
         # Extract the ndarray from the image object
         # -----------------------------------------
-        if isinstance(image, u.quantity.Quantity):
+        if isinstance(image, SRImage):
+            data = image.data
+            disp_axis = image.disp_axis
+        elif isinstance(image, u.quantity.Quantity):
             data = image.value
         elif isinstance(image, np.ndarray):
             data = image
@@ -60,7 +76,7 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None,
         if getattr(image, 'mask', None) is not None:
             mask = image.mask.astype(bool)
         elif mask is not None:
-            pass
+            mask = mask.astype(bool)
         else:
             mask = np.zeros(data.shape, dtype=bool)
 
@@ -79,46 +95,58 @@ def __init__(self, image, disp_axis: int = 1, crossdisp_axis: int | None = None,
 
         # Extract the uncertainty
         # -----------------------
-        if variance is not None:
-            uncertainty = VarianceUncertainty(variance)
+        unc_types = {'var': VarianceUncertainty,
+                     'std': StdDevUncertainty,
+                     'ivar': InverseVariance}
+        if uncertainty is not None:
+            if isinstance(uncertainty, np.ndarray):
+                uncertainty = unc_types[uncertainty_type](uncertainty)
+            elif isinstance(uncertainty, NDUncertainty):
+                pass
+            else:
+                raise ValueError('Unrecognized uncertainty type.')
         elif getattr(image, 'uncertainty', None) is not None:
             uncertainty = image.uncertainty
         elif not require_uncertainty:
-            uncertainty = VarianceUncertainty(np.ones_like(data))
+            uncertainty = None
         else:
             raise ValueError('Uncertainty information is required but missing.')
 
         # Try to convert the uncertainty into VarianceUncertainty
         if ensure_var_uncertainty:
             if uncertainty.uncertainty_type != 'var':
-                warnings.warn("image NDData object's uncertainty is not"
+                warnings.warn("image object's uncertainty is not"
                               "given as VarianceUncertainty. Trying to "
                               "convert the uncertainty to VarianceUncertainty.")
                 uncertainty = uncertainty.represent_as(VarianceUncertainty)
 
-        if uncertainty.array.shape != data.shape:
-            raise ValueError('Image and uncertainty shapes must match.')
-        if np.any(uncertainty.array < 0):
-            raise ValueError("Uncertainty must be fully positive")
-        if np.all(uncertainty.array == 0):
-            # technically would result in infinities, but since they're all
-            # zeros, we can override ones to simulate an unweighted case
-            uncertainty._array[:] = 1.0
-        if np.any(uncertainty.array == 0):
-            m = uncertainty.array == 0
-            # exclude such elements by editing the input mask
-            mask[m] = True
-            # replace the variances to avoid a divide by zero warning
-            uncertainty._array[m] = np.nan
+        if uncertainty is not None:
+            if uncertainty.array.shape != data.shape:
+                raise ValueError('Image and uncertainty shapes must match.')
+            if np.any(uncertainty.array < 0):
+                raise ValueError("Uncertainty must be fully positive")
+            if np.all(uncertainty.array == 0):
+                # technically would result in infinities, but since they're all
+                # zeros, we can override ones to simulate an unweighted case
+                uncertainty._array[:] = 1.0
+            if np.any(uncertainty.array == 0):
+                m = uncertainty.array == 0
+                # exclude such elements by editing the input mask
+                mask[m] = True
+                # replace the variances to avoid a divide by zero warning
+                uncertainty._array[m] = np.nan
 
         # Standardise the axes
         # --------------------
         # Force the cross-dispersion axis as the first axis and the dispersion
         # axis as the second axis. This could be done using transpose as well, but this
         # approach works also with image cubes (although we're not supporting them yet).
-        data = np.moveaxis(data, [crossdisp_axis, disp_axis], [0, 1])
-        mask = np.moveaxis(mask, [crossdisp_axis, disp_axis], [0, 1])
-        uncertainty._array = np.moveaxis(uncertainty._array, [crossdisp_axis, disp_axis], [0, 1])
+        if (crossdisp_axis, disp_axis) != (0, 1):
+            data = np.moveaxis(data, [crossdisp_axis, disp_axis], [0, 1])
+            mask = np.moveaxis(mask, [crossdisp_axis, disp_axis], [0, 1])
+            uncertainty._array = np.moveaxis(uncertainty._array,
+                                             [crossdisp_axis, disp_axis],
+                                             [0, 1])
 
         self._nddata = NDDataRef(data * unit, uncertainty=uncertainty, mask=mask)
 

From eace0ae16fcd5cb59e1cf44f2d755ecccbba4b0b Mon Sep 17 00:00:00 2001
From: Hannu Parviainen <hannu@iac.es>
Date: Mon, 11 Nov 2024 19:30:08 +0100
Subject: [PATCH 25/25] Refactored the image handling logic to further use the
 SRImage class, replacing the _ImageParser. Updated related tests and modules
 to accommodate SRImage-specific methods and streamlined code a bit.

---
 specreduce/background.py         |   9 +-
 specreduce/core.py               | 201 +-----------------------
 specreduce/extract.py            | 254 +++++++++----------------------
 specreduce/tests/test_extract.py |  11 +-
 specreduce/tests/test_srimage.py |   7 +-
 specreduce/tracing.py            |   8 +-
 specreduce/utils/utils.py        |  26 +---
 7 files changed, 96 insertions(+), 420 deletions(-)

diff --git a/specreduce/background.py b/specreduce/background.py
index 07f0a4e..ddb57f1 100644
--- a/specreduce/background.py
+++ b/specreduce/background.py
@@ -4,8 +4,6 @@
 from dataclasses import dataclass, field, InitVar
 
 import numpy as np
-from astropy import units as u
-from astropy.nddata import NDData
 from astropy.utils.decorators import deprecated_attribute
 from specutils import Spectrum1D
 
@@ -302,7 +300,9 @@ def one_sided(cls, image, trace_object, separation, **kwargs):
         kwargs['traces'] = [trace_object+separation]
         return cls(image=image, **kwargs)
 
-    def bkg_image(self, image=None, disp_axis: int = 1, crossdisp_axis: int | None = None) -> SRImage:
+    def bkg_image(self, image=None,
+                  disp_axis: int = 1,
+                  crossdisp_axis: int | None = None) -> SRImage:
         """
         Expose the background tiled to the dimension of ``image``.
 
@@ -319,7 +319,8 @@ def bkg_image(self, image=None, disp_axis: int = 1, crossdisp_axis: int | None =
         `~specutils.SRImage` object with same shape as ``image``.
         """
         image = self.image if image is None else as_image(image, disp_axis, crossdisp_axis)
-        return SRImage(np.tile(self._bkg_array,(image.shape[image.crossdisp_axis], 1)) * image.unit)
+        return SRImage(np.tile(self._bkg_array,
+                               (image.shape[image.crossdisp_axis], 1)) * image.unit)
 
     def bkg_spectrum(self, image=None) -> Spectrum1D:
         """
diff --git a/specreduce/core.py b/specreduce/core.py
index 76894b1..ec962a2 100644
--- a/specreduce/core.py
+++ b/specreduce/core.py
@@ -12,207 +12,8 @@
 __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`
-    """
-
-    # The '_valid_mask_treatment_methods' in the Background, Trace, and Extract
-    # classes is a subset of implemented methods.
-    implemented_mask_treatment_methods = 'filter', 'zero-fill', 'omit'
-
-    def _parse_image(self, image,
-                     disp_axis: int = 1,
-                     mask_treatment: str = 'filter') -> Spectrum1D:
-        """
-        Convert all accepted image types to a consistently formatted Spectrum1D object.
-
-        Parameters
-        ----------
-        image : `~astropy.nddata.NDData`-like or array-like
-            The image to be parsed. If None, defaults to class' own
-            image attribute.
-        disp_axis
-            The index of the image's dispersion axis. Should not be
-            changed until operations can handle variable image
-            orientations. [default: 1]
-        mask_treatment
-            Treatment method for the mask.
-
-        Returns
-        -------
-        Spectrum1D
-        """
-        # 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
-
-        return self._get_data_from_image(image, disp_axis=disp_axis,
-                                         mask_treatment=mask_treatment)
-
-    @staticmethod
-    def _get_data_from_image(image,
-                             disp_axis: int = 1,
-                             mask_treatment: str = 'filter') -> Spectrum1D:
-        """
-        Extract data array from various input types for `image`.
-
-        Parameters
-        ----------
-        image : array-like or Quantity
-            Input image from which data is extracted. This can be a 2D numpy
-            array, Quantity, or an NDData object.
-        disp_axis : int, optional
-            The dispersion axis of the image.
-        mask_treatment : str, optional
-            Treatment method for the mask:
-            - 'filter' (default): Return the unmodified input image and combined mask.
-            - 'zero-fill': Set masked values in the image to zero.
-            - 'omit': Mask all pixels along the cross dispersion axis if any value is masked.
-
-        Returns
-        -------
-        Spectrum1D
-        """
-        # This works only with 2D images.
-        crossdisp_axis = (disp_axis + 1) % 2
-
-        if isinstance(image, u.quantity.Quantity):
-            img = image.value
-        elif isinstance(image, np.ndarray):
-            img = image
-        else:  # NDData, including CCDData and Spectrum1D
-            img = image.data
-
-        mask = getattr(image, 'mask', None)
-
-        # next, handle masked and nonfinite data in image.
-        # A mask will be created from any nonfinite image data, and combined
-        # with any additional 'mask' passed in. If image is being parsed within
-        # a specreduce operation that has 'mask_treatment' options, this will be
-        # handled as well. Note that input data may be modified if a fill value
-        # is chosen to handle masked data. The returned image will always have
-        # `image.mask` even if there are no nonfinte or masked values.
-        img, mask = _ImageParser._mask_and_nonfinite_data_handling(image=img,
-                                                                   mask=mask,
-                                                                   mask_treatment=mask_treatment,
-                                                                   crossdisp_axis=crossdisp_axis)
-
-        # mask (handled above) 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 hasattr(image, 'uncertainty'):
-            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)
-
-        img = Spectrum1D(img * unit, spectral_axis=spectral_axis,
-                         uncertainty=uncertainty, mask=mask)
-        return img
-
-    @staticmethod
-    def _mask_and_nonfinite_data_handling(image, mask,
-                                          mask_treatment: str = 'filter',
-                                          crossdisp_axis: int = 0) -> tuple[np.ndarray, np.ndarray]:
-        """
-        Handle the treatment of masked and nonfinite data.
-
-        All operations in Specreduce can take in a mask for the data as
-        part of the input NDData. Additionally, any non-finite values in the
-        data that aren't in the user-supplied mask will be combined bitwise
-        with the input mask.
-
-        There are three options currently implemented for the treatment
-        of masked and nonfinite data - filter, omit, and zero-fill.
-        Depending on the step, all or a subset of these three options are valid.
-
-        Parameters
-        ----------
-        image : array-like
-            The input image data array that may contain nonfinite values.
-        mask : array-like or None
-            An optional mask array. Nonfinite values in the image will be added to this mask.
-        mask_treatment : str
-            Specifies how to handle masked data:
-            - 'filter' (default): Returns the unmodified input image and combined mask.
-            - 'zero-fill': Sets masked values in the image to zero.
-            - 'omit': Masks entire columns or rows if any value is masked.
-        crossdisp_axis : int
-            Axis along which to collapse the 2D mask into a 1D mask for treatment 'omit'.
-        """
-        if mask_treatment not in _ImageParser.implemented_mask_treatment_methods:
-            raise ValueError("`mask_treatment` must be one of "
-                             f"{_ImageParser.implemented_mask_treatment_methods}")
-
-        # make sure there is always a 'mask', even when all data is unmasked and finite.
-        if mask is not None:
-            # always mask any previously uncaught nonfinite values in image array
-            # combining these with the (optional) user-provided mask on `image.mask`
-            mask = np.logical_or(mask, ~np.isfinite(image))
-        else:
-            mask = ~np.isfinite(image)
-
-        # if mask option is the default 'filter' option,
-        # nothing needs to be done. input mask (combined with nonfinite data)
-        # remains with data as-is.
-
-        if mask_treatment == 'zero-fill':
-            # make a copy of the input image since we will be modifying it
-            image = deepcopy(image)
-
-            # if mask_treatment is 'zero_fill', set masked values to zero in
-            # image data and drop image.mask. note that this is done after
-            # _combine_mask_with_nonfinite_from_data, so non-finite values in
-            # data (which are now in the mask) will also be set to zero.
-            # set masked values to zero
-            image[mask] = 0.
-
-            # masked array with no masked values, so accessing image.mask works
-            # but we don't need the actual mask anymore since data has been set to 0
-            mask = np.zeros(image.shape, dtype=bool)
-
-        elif mask_treatment == 'omit':
-            # collapse 2d mask (after accounting for addl non-finite values in
-            # data) to a 1d mask, along dispersion axis, to fully mask columns
-            # that have any masked values.
-
-            # create a 1d mask along crossdisp axis - if any column has a single nan,
-            # the entire column should be masked
-            reduced_mask = np.logical_or.reduce(mask, axis=crossdisp_axis)
-
-            # back to a 2D mask
-            shape = (image.shape[0], 1) if crossdisp_axis == 0 else (1, image.shape[1])
-            mask = np.tile(reduced_mask, shape)
-
-        # check for case where entire image is masked.
-        if mask.all():
-            raise ValueError('Image is fully masked. Check for invalid values.')
-
-        return image, mask
-
-
 @dataclass
-class SpecreduceOperation(_ImageParser):
+class SpecreduceOperation:
     """
     An operation to perform as part of a spectroscopic reduction pipeline.
 
diff --git a/specreduce/extract.py b/specreduce/extract.py
index 1f35352..c027cc5 100644
--- a/specreduce/extract.py
+++ b/specreduce/extract.py
@@ -1,12 +1,11 @@
 # Licensed under a 3-clause BSD style license - see LICENSE.rst
 
-import warnings
 from dataclasses import dataclass, field, InitVar
 
 import numpy as np
 from astropy import units as u
 from astropy.modeling import Model, models, fitting
-from astropy.nddata import NDData, VarianceUncertainty
+from astropy.nddata import VarianceUncertainty
 from scipy.integrate import trapezoid
 from scipy.interpolate import RectBivariateSpline
 from specutils import Spectrum1D
@@ -18,7 +17,7 @@
 __all__ = ['BoxcarExtract', 'HorneExtract', 'OptimalExtract']
 
 
-def _get_boxcar_weights(center, hwidth, npix):
+def _get_boxcar_weights(center: float, hwidth: float, npix: int):
     """
     Compute weights given an aperture center, half width,
     and number of pixels.
@@ -28,14 +27,14 @@ def _get_boxcar_weights(center, hwidth, npix):
 
     Parameters
     ----------
-    center : float, required
+    center
         The index of the aperture's center pixel on the larger image's
         cross-dispersion axis.
 
-    hwidth : float, required
+    hwidth
         Half of the aperture's width in the cross-dispersion direction.
 
-    npix : float, required
+    npix
         The number of pixels in the larger image's cross-dispersion
         axis.
 
@@ -156,18 +155,22 @@ class BoxcarExtract(SpecreduceOperation):
     image: SRImage
     trace_object: Trace
     width: float = 5
-    disp_axis: int = 1
-    crossdisp_axis: int = 0
-    # TODO: should disp_axis and crossdisp_axis be defined in the Trace object?
+    disp_axis: InitVar[int] = 1
+    crossdisp_axis: InitVar[int | None] = None
     mask_treatment: str = 'filter'
     _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill')
 
+    def __post_init__(self, disp_axis, crossdisp_axis):
+        self.image = as_image(self.image,
+                              disp_axis=disp_axis,
+                              crossdisp_axis=crossdisp_axis)
+
     @property
     def spectrum(self):
         return self.__call__()
 
     def __call__(self, image=None, trace_object=None, width=None,
-                 disp_axis=None, crossdisp_axis=None,
+                 disp_axis=1, crossdisp_axis=None,
                  mask_treatment: str = 'filter') -> Spectrum1D:
         """
         Extract the 1D spectrum using the boxcar method.
@@ -208,8 +211,6 @@ def __call__(self, image=None, trace_object=None, width=None,
         """
         trace_object = trace_object if trace_object is not None else self.trace_object
         width = width if width is not None else self.width
-        disp_axis = disp_axis if disp_axis is not None else self.disp_axis
-        crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
 
         # Parse image, including masked/nonfinite data handling based on
         # choice of `mask_treatment`, which for BoxcarExtract can be filter, zero-fill, or
@@ -218,7 +219,7 @@ def __call__(self, image=None, trace_object=None, width=None,
         #                                mask_treatment=self.mask_treatment)
         # TODO: Mask treatment
         self.image = as_image(image if image is not None else self.image,
-                              self.disp_axis, self.crossdisp_axis)
+                              disp_axis, crossdisp_axis)
 
         # # _parse_image returns a Spectrum1D. convert this to a masked array
         # # for ease of calculations here (even if there is no masked data).
@@ -321,134 +322,28 @@ class HorneExtract(SpecreduceOperation):
     trace_object: Trace
     bkgrd_prof: Model = field(default=models.Polynomial1D(2))
     spatial_profile: str | dict = 'gaussian'  # can actually be str, dict
-    variance: np.ndarray | None = field(default=None)
-    mask: np.ndarray | None = field(default=None)
-    unit: u.Unit | None = field(default=None)
-    disp_axis: InitVar[int] = 1
-    crossdisp_axis: InitVar[int | None] = None
+    variance: InitVar[np.ndarray | None] = field(default=None)
+    mask: InitVar[np.ndarray | None] = field(default=None)
+    unit: InitVar[u.Unit | None] = field(default=None)
+    disp_axis: InitVar[int] = field(default=1)
+    crossdisp_axis: InitVar[int | None] = field(default=None)
+
+    def __post_init__(self, variance, mask, unit, disp_axis, crossdisp_axis):
+        self.image = SRImage(self.image,
+                        disp_axis=disp_axis,
+                        crossdisp_axis=crossdisp_axis,
+                        unit=unit,
+                        mask=mask,
+                        uncertainty=variance,
+                        uncertainty_type='var',
+                        ensure_var_uncertainty=False,
+                        require_uncertainty=False)
 
-    def __post_init__(self, disp_axis, crossdisp_axis):
-        self.image = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
 
     @property
     def spectrum(self):
         return self.__call__()
 
-    def _parse_image(self, image,
-                     variance=None, mask=None, unit=None, disp_axis=1):
-        """
-        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 in DN.
-        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(image, np.ndarray):
-            img = image
-        elif isinstance(image, u.quantity.Quantity):
-            img = image.value
-        else:  # NDData, including CCDData and Spectrum1D
-            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(image, 'mask', None) is not None:
-            mask = image.mask
-        elif mask is not None:
-            pass
-        else:
-            # if user provides no mask at all, don't mask anywhere
-            mask = np.zeros_like(img)
-
-        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(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 = image.uncertainty.array**2
-            elif image.uncertainty.uncertainty_type == 'ivar':
-                variance = 1 / image.uncertainty.array
-            else:
-                # other options are InverseUncertainty and UnknownUncertainty
-                raise ValueError("image NDData object has unexpected "
-                                 "uncertainty type. instead, try "
-                                 "VarianceUncertainty or StdDevUncertainty.")
-        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:
-            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 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(image, 'unit',
-                       u.Unit(unit) if unit is not None else 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=variance, mask=mask)
-
     def _fit_gaussian_spatial_profile(self, img, disp_axis, crossdisp_axis,
                                       or_mask, bkgrd_prof):
 
@@ -609,15 +504,29 @@ def __call__(self, image=None, trace_object=None,
             The final, Horne extracted 1D spectrum.
         """
         image = image if image is not None else self.image
+
+        if variance is not None:
+            uncertainty = VarianceUncertainty(variance)
+        else:
+            uncertainty = getattr(image, 'uncertainty', None)
+            if uncertainty is None:
+                raise ValueError('variance must be specified if the image data does '
+                                 'not contain uncertainty information.')
+
+        image = SRImage(image,
+                        disp_axis=disp_axis or self.image.disp_axis,
+                        crossdisp_axis=crossdisp_axis or self.image.crossdisp_axis,
+                        unit=unit or self.image.unit,
+                        mask=mask if mask is not None else self.image.mask,
+                        uncertainty=uncertainty,
+                        uncertainty_type='var',
+                        ensure_var_uncertainty=True,
+                        require_uncertainty=True)
+
         trace_object = trace_object if trace_object is not None else self.trace_object
-        disp_axis = disp_axis if disp_axis is not None else self.disp_axis
-        crossdisp_axis = crossdisp_axis if crossdisp_axis is not None else self.crossdisp_axis
         bkgrd_prof = bkgrd_prof if bkgrd_prof is not None else self.bkgrd_prof
         spatial_profile = (spatial_profile if spatial_profile is not None else
                            self.spatial_profile)
-        variance = variance if variance is not None else self.variance
-        mask = mask if mask is not None else self.mask
-        unit = unit if unit is not None else self.unit
 
         # figure out what 'spatial_profile' was provided
         # put this parsing into another method at some point, its a lot..
@@ -665,42 +574,22 @@ def __call__(self, image=None, trace_object=None,
         else:
             raise ValueError('``spatial_profile`` must either be string or dictionary.')
 
-        # parse image and replace optional arguments with updated values
-        self.image = SRImage(image,
-                             disp_axis=disp_axis,
-                             unit=unit,
-                             mask=mask,
-                             variance=variance,
-                             ensure_var_uncertainty=True,
-                             require_uncertainty=True)
-        #self.image = self._parse_image(image, variance, mask, unit, disp_axis)
-        variance = self.image.uncertainty.array
-        mask = self.image.mask
-        unit = self.image.unit
-
-        img = np.ma.masked_array(self.image.data, mask=mask)
-
-        # create separate mask including any previously uncaught non-finite
-        # values for purposes of calculating fit
-        or_mask = np.logical_or(img.mask,
-                                ~np.isfinite(self.image.data))
+        img = image.to_masked_array()
+        variance = image.uncertainty.array
+        mask = img.mask
+        unit = image.unit
 
         # If the trace is not flat, shift the rows in each column
         # so the image is aligned along the trace:
         if not isinstance(trace_object, FlatTrace):
-            img = _align_along_trace(
-                img,
-                trace_object.trace,
-                disp_axis=disp_axis,
-                crossdisp_axis=crossdisp_axis)
+            img = _align_along_trace(img, trace_object.trace)
 
         if self.spatial_profile == 'gaussian':
-
             # fit profile to average (mean) profile along crossdisp axis
             fit_ext_kernel = self._fit_gaussian_spatial_profile(img,
-                                                                disp_axis,
-                                                                crossdisp_axis,
-                                                                or_mask,
+                                                                DISP_AXIS,
+                                                                CROSSDISP_AXIS,
+                                                                mask,
                                                                 bkgrd_prof)
 
             # this is just creating an array of the trace to shift the mean
@@ -712,7 +601,7 @@ def __call__(self, image=None, trace_object=None,
                 mean_init_guess = trace_object.trace
             else:
                 mean_init_guess = np.broadcast_to(
-                    img.shape[crossdisp_axis] // 2, img.shape[disp_axis]
+                    img.shape[CROSSDISP_AXIS] // 2, img.shape[DISP_AXIS]
                 )
 
         else:  # interpolated_profile
@@ -725,7 +614,7 @@ def __call__(self, image=None, trace_object=None,
                                  ' be fit and subtracted from `img` beforehand.')
             # make sure n_bins doesnt exceed the number of (for now) finite
             # columns. update this when masking is fixed.
-            n_finite_cols = np.logical_or.reduce(or_mask, axis=crossdisp_axis)
+            n_finite_cols = np.logical_or.reduce(mask, axis=crossdisp_axis)
             n_finite_cols = np.count_nonzero(n_finite_cols.astype(int) == 0)
 
             # determine interpolation degree from input and make tuple if int
@@ -750,25 +639,26 @@ def __call__(self, image=None, trace_object=None,
                                  'must be less than the number of fully-finite '
                                  f'wavelength columns ({n_finite_cols}).')
 
-            interp_spatial_prof = self._fit_self_spatial_profile(img, disp_axis,
-                                                                 crossdisp_axis,
-                                                                 or_mask,
+            interp_spatial_prof = self._fit_self_spatial_profile(img,
+                                                                 DISP_AXIS,
+                                                                 CROSSDISP_AXIS,
+                                                                 mask,
                                                                  n_bins_interpolated_profile,
                                                                  kx, ky)
 
             # add private attribute to save fit profile. should this be public?
             self._interp_spatial_prof = interp_spatial_prof
 
-        col_mask = np.logical_or.reduce(or_mask, axis=crossdisp_axis)
-        nonf_col = [np.nan] * img.shape[crossdisp_axis]
+        col_mask = np.logical_or.reduce(mask, axis=CROSSDISP_AXIS)
+        nonf_col = [np.nan] * img.shape[CROSSDISP_AXIS]
 
         # array of 'x' values for each wavelength for extraction
-        nrows = img.shape[crossdisp_axis]
+        nrows = img.shape[CROSSDISP_AXIS]
         xd_pixels = np.arange(nrows)
 
         kernel_vals = []
         norms = []
-        for col_pix in range(img.shape[disp_axis]):
+        for col_pix in range(img.shape[DISP_AXIS]):
 
             # for now, skip columns with any non-finite values
             # NOTE: fit and other kernel operations should support masking again
@@ -807,11 +697,11 @@ def __call__(self, image=None, trace_object=None,
         norms = np.array(norms)
 
         # calculate kernel normalization
-        g_x = np.sum(kernel_vals**2 / variance, axis=crossdisp_axis)
+        g_x = np.sum(kernel_vals**2 / variance, axis=CROSSDISP_AXIS)
 
         # sum by column weights
         weighted_img = np.divide(img * kernel_vals, variance)
-        result = np.sum(weighted_img, axis=crossdisp_axis) / g_x
+        result = np.sum(weighted_img, axis=CROSSDISP_AXIS) / g_x
 
         # multiply kernel normalization into the extracted signal
         extraction = result * norms
@@ -820,7 +710,7 @@ def __call__(self, image=None, trace_object=None,
         return Spectrum1D(extraction * unit)
 
 
-def _align_along_trace(img, trace_array, disp_axis=1, crossdisp_axis=0):
+def _align_along_trace(img, trace_array):
     """
     Given an arbitrary trace ``trace_array`` (an np.ndarray), roll
     all columns of ``nddata`` to shift the NDData's pixels nearest
@@ -828,10 +718,6 @@ def _align_along_trace(img, trace_array, disp_axis=1, crossdisp_axis=0):
     NDData.
     """
     # TODO: this workflow does not support extraction for >2D spectra
-    if not (disp_axis == 1 and crossdisp_axis == 0):
-        # take the transpose to ensure the rows are the cross-disp axis:
-        img = img.T
-
     n_rows, n_cols = img.shape
 
     # indices of all columns, in their original order
diff --git a/specreduce/tests/test_extract.py b/specreduce/tests/test_extract.py
index 7200d2d..9aac149 100644
--- a/specreduce/tests/test_extract.py
+++ b/specreduce/tests/test_extract.py
@@ -10,6 +10,7 @@
 from specreduce.extract import (
     BoxcarExtract, HorneExtract, OptimalExtract, _align_along_trace
 )
+from specreduce.image import SRImage
 from specreduce.tracing import FitTrace, FlatTrace, ArrayTrace
 
 
@@ -122,15 +123,15 @@ def test_horne_image_validation(mk_test_img):
         ext = extract()
 
     # an NDData-type image can't have an empty uncertainty attribute
-    with pytest.raises(ValueError, match=r'.*NDData object lacks uncertainty'):
+    with pytest.raises(ValueError, match=r'.*variance must be specified.*'):
         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)
+    with pytest.raises(TypeError, match=r'.*does not support conversion*'):
+        err = UnknownUncertainty(np.ones_like(image.data))
+        im = SRImage(image, uncertainty=err)
+        ext = extract(image=im)
 
     # an array-type image must have the same dimensions as its variance argument
     with pytest.raises(ValueError, match=r'.*shapes must match.*'):
diff --git a/specreduce/tests/test_srimage.py b/specreduce/tests/test_srimage.py
index 90a7a2d..4179627 100644
--- a/specreduce/tests/test_srimage.py
+++ b/specreduce/tests/test_srimage.py
@@ -2,9 +2,10 @@
 import numpy as np
 import pytest
 
-from astropy.nddata import NDData, VarianceUncertainty
+from astropy.nddata import NDData, VarianceUncertainty  # noqa
 from specreduce.image import SRImage
 
+
 def create_image(imtype: str, nrow: int = 7, ncol: int = 11):
     img = np.tile((np.arange(ncol)), (nrow, 1)).astype('d')
     if imtype == 'ndarray':
@@ -16,8 +17,6 @@ def create_image(imtype: str, nrow: int = 7, ncol: int = 11):
     else:
         raise NotImplementedError('unkown image type')
 
-img = np.tile((np.arange(5, 15)), (7, 1)).astype('d')
-
 
 @pytest.mark.parametrize("imtype", ["ndarray", "quantity", 'nddata'])
 def test_init(imtype):
@@ -46,4 +45,4 @@ def test_init(imtype):
 
 def test_init_bad():
     with pytest.raises(ValueError, match='Unrecognized image type.'):
-        image = SRImage([[0.0, 1.0],[2.0, 3.0]])
\ No newline at end of file
+        image = SRImage([[0.0, 1.0],[2.0, 3.0]])  # noqa
\ No newline at end of file
diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index c5f18d9..f53141c 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -34,7 +34,8 @@ class Trace:
     crossdisp_axis: InitVar[int | None] = field(default=None, kw_only=True)
     trace: np.ma.MaskedArray | None = field(default=None, init=False, repr=False)
     mask_treatment: str | None = field(default=None, init=False, repr=False, kw_only=True)
-    _valid_mask_treatment_methods: tuple[str | None] = field(default=(None,), init=False, repr=False)
+    _valid_mask_treatment_methods: tuple[str | None] = field(default=(None,),
+                                                             init=False, repr=False)
 
     def __post_init__(self, disp_axis, crossdisp_axis):
         self.image = im = as_image(self.image, disp_axis=disp_axis, crossdisp_axis=crossdisp_axis)
@@ -265,7 +266,10 @@ class FitTrace(Trace):
     trace_model: Model = field(default=models.Polynomial1D(degree=1))
     peak_method: str = 'max'
     mask_treatment: str = 'filter'
-    _valid_mask_treatment_methods: tuple[str | None] = field(default=('filter', 'omit', 'zero-fill'), init=False, repr=False)
+    _valid_mask_treatment_methods: tuple[str | None] = field(default=('filter',
+                                                                      'omit',
+                                                                      'zero-fill'),
+                                                             init=False, repr=False)
     # for testing purposes only, save bin peaks if requested
     _save_bin_peaks_testing: bool = False
 
diff --git a/specreduce/utils/utils.py b/specreduce/utils/utils.py
index 8cdee45..325ff8d 100644
--- a/specreduce/utils/utils.py
+++ b/specreduce/utils/utils.py
@@ -1,6 +1,6 @@
 import numpy as np
 
-from specreduce.core import _ImageParser
+from specreduce.image import SRImage
 from specreduce.tracing import Trace, FlatTrace
 from specreduce.extract import _ap_weight_image, _align_along_trace
 
@@ -17,7 +17,7 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
     along the dispersion axis.
 
     If a single number is specified for `pixel`, then the profile at that pixel
-    (i.e wavelength) will be returned. If several pixels are specified in a list or
+    (i.e. wavelength) will be returned. If several pixels are specified in a list or
     array, then they will be averaged (median or mean, set by `statistic` which
     defaults to median). Alternatively, `pixel_range` can be specified as a tuple
     of integers specifying the minimum and maximum pixel range to average the
@@ -85,20 +85,7 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
 
     unit = getattr(image, 'unit', None)
 
-    # parse image, which will return a spectrum1D (note: this is not ideal,
-    # but will be addressed at some point)
-    parser = _ImageParser()
-    image = parser._parse_image(image, disp_axis=disp_axis)
-
-    # which we then need to make back into a masked array
-    # again this way of parsing the image is not ideal but
-    # thats just how it is for now.
-    image = np.ma.MaskedArray(image.data, mask=image.mask)
-
-    # transpose if disp_axis = 0 just for simplicity of calculations
-    # image is already copied so this won't modify input
-    if disp_axis == 0:
-        image = image.T
+    image = SRImage(image, disp_axis=disp_axis).to_masked_array()
 
     nrows = image.shape[crossdisp_axis]
     ncols = image.shape[disp_axis]
@@ -157,16 +144,13 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
     else:
         raise ValueError('`width` must be an integer, or None to use all '
                          'cross-dispersion pixels.')
-        width = int(width)
 
     # rectify trace, if _align_along_trace is True and trace is not flat
     aligned_trace = None
     if align_along_trace:
         if not isinstance(trace, FlatTrace):
             # note: img was transposed according to `crossdisp_axis`: disp_axis will always be 1
-            aligned_trace = _align_along_trace(image, trace.trace,
-                                               disp_axis=1,
-                                               crossdisp_axis=0)
+            aligned_trace = _align_along_trace(image, trace.trace)
 
             # new trace will be a flat trace at the center of the image
             trace_pos = nrows / 2
@@ -183,7 +167,7 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
 
     # now that we have figured out the mask for the window in cross-disp. axis,
     # select only the pixel(s) we want to include in measuring the avg. profile
-    pixel_mask = np.ones((image.shape))
+    pixel_mask = np.ones(image.shape)
     pixel_mask[:, pixels] = 0
 
     # combine these masks to isolate the rows and cols used to measure profile