diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py index 96853e2..ec372b8 100644 --- a/specreduce/tests/test_tracing.py +++ b/specreduce/tests/test_tracing.py @@ -417,21 +417,11 @@ def test_mask_treatment_filter(self, peak_method, expected): 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): + @pytest.mark.parametrize("peak_method", ["max", "gaussian", "centroid"]) + def test_mask_treatment_zero_fill(self, peak_method): """ - 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. + Test to ensure mask_treatment=`zero_fill` for FitTrace raises a `ValueError`. + Parametrized over different `peak_method` options. """ # Make an image with some nonfinite values. @@ -441,27 +431,12 @@ def test_mask_treatment_zero_fill(self, peak_method, expected): # 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) + image2 = NDData(dat, mask=~np.isfinite(image1)) 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) + with pytest.raises(ValueError): + FitTrace(imgg, peak_method=peak_method, mask_treatment='zero-fill') @pytest.mark.filterwarnings("ignore:All pixels in bins") @pytest.mark.parametrize("peak_method,expected", diff --git a/specreduce/tracing.py b/specreduce/tracing.py index d18b7fc..e82292b 100644 --- a/specreduce/tracing.py +++ b/specreduce/tracing.py @@ -3,6 +3,7 @@ import warnings from copy import deepcopy from dataclasses import dataclass, field +from typing import Literal import numpy as np from astropy.modeling import Model, fitting, models @@ -257,30 +258,27 @@ class FitTrace(Trace, _ImageParser): ``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. + The method for handling masked or non-finite data. Choice of `filter` or + `omit`. 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). + For both 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 - window: int = None + bins: int | None = None + guess: float | None = None + window: int | None = None trace_model: Model = field(default=models.Polynomial1D(degree=1)) - peak_method: str = 'max' - _crossdisp_axis = 0 - _disp_axis = 1 - mask_treatment: str = 'filter' - _valid_mask_treatment_methods = ('filter', 'omit', 'zero-fill') + peak_method: Literal['gaussian', 'centroid', 'max'] = 'max' + _crossdisp_axis: int = 0 + _disp_axis: int = 1 + mask_treatment: Literal['filter', 'omit'] = 'filter' + _valid_mask_treatment_methods = ('filter', 'omit') # for testing purposes only, save bin peaks if requested _save_bin_peaks_testing: bool = False @@ -288,6 +286,10 @@ def __post_init__(self): # Parse image, including masked/nonfinite data handling based on # choice of `mask_treatment`. returns a Spectrum1D + if self.mask_treatment not in self._valid_mask_treatment_methods: + raise ValueError("`mask_treatment` must be one of " + f"{self._valid_mask_treatment_methods}") + self.image = self._parse_image(self.image, disp_axis=self._disp_axis, mask_treatment=self.mask_treatment)