Merge pull request #3 from hpparvi/masking_options_rebase

Modified FitTrace to exclude masked pixels from the fit and fixed the corresponding test.
cshanahan1 · Dec 13, 2024 · 354b98f · 354b98f
2 parents 4d8ece7 + edf299e
commit 354b98f
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/specreduce/tests/test_tracing.py b/specreduce/tests/test_tracing.py
@@ -373,8 +373,8 @@ def test_fit_trace_fully_masked_cols(self, mask_treatment):
     @pytest.mark.parametrize("peak_method,expected",
                              [("max", [5., 3., 5., 5., 7., 5.,
                                        np.nan, 5., 5., 5., 2., 5.]),
-                              ("gaussian", [5., 1.696159, 5., 5., 7.80744334,
-                                            5., np.nan, 5., 5., 5., 1.28216332, 5.]),
+                              ("gaussian", [5., 5., 5., 5., 5.,
+                                            5., np.nan, 5., 5., 5., 1.576, 5.]),
                               ("centroid", [4.27108332, 2.24060342, 4.27108332,
                                             4.27108332, 6.66827608, 4.27108332,
                                             np.nan, 4.27108332, 4.27108332,

diff --git a/specreduce/tracing.py b/specreduce/tracing.py
@@ -352,7 +352,7 @@ def _fit_trace(self, img):
         yy = np.arange(img.shape[self._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.mean(axis=self._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
 
@@ -393,9 +393,9 @@ def _fit_trace(self, img):
         warn_bins = []
         for i in range(self.bins):
 
-            # binned columns, summed along disp. axis.
+            # binned columns, averaged 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]].mean(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
@@ -424,7 +424,7 @@ def _fit_trace(self, img):
                 offset_init_i = models.Const1D(np.ma.median(z_i))
 
                 profile_i = g1d_init_i + offset_init_i
-                popt_i = fitter(profile_i, ilum2, z_i)
+                popt_i = fitter(profile_i, ilum2[~z_i.mask], z_i.data[~z_i.mask])
 
                 # if gaussian fits off chip, then fall back to previous answer
                 if not ilum2.min() <= popt_i.mean_0 <= ilum2.max():