diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py index a609677..3421d00 100644 --- a/specreduce/tests/test_tracing.py +++ b/specreduce/tests/test_tracing.py @@ -2,6 +2,8 @@ import pytest from astropy.modeling import models +from astropy.nddata import NDData +import astropy.units as u from specreduce.utils.synth_data import make_2d_trace_image from specreduce.tracing import Trace, FlatTrace, ArrayTrace, FitTrace @@ -148,4 +150,14 @@ def test_fit_trace(): with pytest.raises(ValueError, match=r'image is fully masked'): FitTrace(img_all_nans) - # could try to catch warning thrown for all-nan bins + # test that warning is raised when many bins are masked + mask = np.zeros(img.shape) + mask[:, 75:76] = 1 + nddat = NDData(data=img, mask=mask, unit=u.DN) + warn_bins = np.arange(20) + warn_bins = ', '.join([str(x) for x in warn_bins]) + msg = f"All pixels in bin 75 are masked. Falling to trace value from all-bin fit." + with pytest.warns(UserWarning, match=msg): + FitTrace(nddat) + + # could try to catch warning thrown for all-nan bins \ No newline at end of file diff --git a/specreduce/tracing.py b/specreduce/tracing.py index 0780cc4..f7838f3 100644 --- a/specreduce/tracing.py +++ b/specreduce/tracing.py @@ -349,13 +349,13 @@ def __post_init__(self): # warn about fully-masked bins if len(warn_bins) > 0: - plural = 's' if len(warn_bins) == 1 else '' + plural = 's' if len(warn_bins) > 1 else '' # if there are a ton of bins, we don't want to print them all out + print(warn_bins) if len(warn_bins) > 20: warn_bins = warn_bins[0: 10] + ['...'] + warn_bins[10: 19] - - warnings.warn(f"All pixels in bin{plural} {', '.join([str(x) for x in warn_bins])} are masked. Falling " - 'to trace value from all-bin fit.') + warnings.warn(f"All pixels in bin{plural} {', '.join([str(x) for x in warn_bins])}" + " are masked. Falling to trace value from all-bin fit.") # recenter bin positions x_bins = (x_bins[:-1] + x_bins[1:]) / 2