From dbd79b04706dcc8ddcb6e15ce83b9c642605e20c Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Mon, 22 Jan 2024 23:46:23 -0500
Subject: [PATCH] review comments

---
 specreduce/tracing.py | 24 +++++++++++++-----------
 1 file changed, 13 insertions(+), 11 deletions(-)

diff --git a/specreduce/tracing.py b/specreduce/tracing.py
index 4ef83b0..8fca13c 100644
--- a/specreduce/tracing.py
+++ b/specreduce/tracing.py
@@ -313,13 +313,14 @@ def _fit_trace(self, img):
         warn_bins = []
         for i in range(self.bins):
 
-            # binned column (sum) if bins < ncols. otherwise, just columns.
+            # 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)
 
             if self.peak_method == 'gaussian':
                 # if binned column is fully masked for peak_method='gaussian',
                 # the fit value for this bin should be nan, then continue to next
-                warn_msg = 'nan'
+
                 if z_i.mask.all():
                     warn_bins.append(i)
                     y_bins[i] = np.nan
@@ -355,26 +356,23 @@ def _fit_trace(self, img):
             if z_i.mask.all():  # all-masked bins when peak_method is 'centroid' or 'max'
                 warn_bins.append(i)
 
-            elif self.peak_method == 'centroid':
+            if self.peak_method == 'centroid':
                 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)
 
-                # warning message for fully masked bin
                 # NOTE this reflects current behavior, should eventually be changed
                 # to set to nan by default (or zero fill / interpoate option once
-                # available) and warn accordingly.
-                warn_msg = f'largest bin index ({str(z_i_cumsum.shape[0])})'
+                # available)
 
             elif self.peak_method == 'max':
                 # TODO: implement smoothing with provided width
                 y_bins[i] = ilum2[z_i.argmax()]
 
-                # warning message for fully masked bin
-                # NOTE: should eventually be changed to set to nan by default
-                # (or zero fill / interpoate option onceavailable) and warn.
-                warn_msg = 'zero'
+                # NOTE: a fully masked should eventually be changed to set to
+                # nan by default (or zero fill / interpoate option once available)
+
 
         # warn about fully-masked bins
         if len(warn_bins) > 0:
@@ -383,10 +381,14 @@ def _fit_trace(self, img):
             if len(warn_bins) > 20:
                 warn_bins = warn_bins[0: 10] + ['...'] + [warn_bins[-1]]
 
+            # warning message printed out depends on `peak_method`
+            warn_msgs = {'gaussian': 'nan', 'max': 'zero', 
+                         'centroid': f'largest bin index ({str(img.shape[0])})'}
+
             warnings.warn(f"All pixels in {'bins' if len(warn_bins) else 'bin'} "
                           f"{', '.join([str(x) for x in warn_bins])}"
                           " are fully masked. Setting bin"
-                          f" peak{'s' if len(warn_bins) else ''} to {warn_msg}.")
+                          f" peak{'s' if len(warn_bins) else ''} to {warn_msgs[self.peak_method]}.")
 
         # recenter bin positions
         x_bins = (x_bins[:-1] + x_bins[1:]) / 2