From 1d45f370bda4b63a59bc47137373293a65308e5a Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Mon, 25 Mar 2024 00:34:41 -0400
Subject: [PATCH] .

---
 specreduce/tests/test_utils.py | 14 ++++++--------
 specreduce/utils/utils.py      |  6 +++---
 2 files changed, 9 insertions(+), 11 deletions(-)

diff --git a/specreduce/tests/test_utils.py b/specreduce/tests/test_utils.py
index 2407e9d..9b64170 100644
--- a/specreduce/tests/test_utils.py
+++ b/specreduce/tests/test_utils.py
@@ -20,13 +20,13 @@ def mk_gaussian_img(nrows=20, ncols=16, mean=10, stddev=4):
 
     return col_model(index_arr.T)
 
+
 def mk_img_non_flat_trace(nrows=40, ncols=100, amp=10, stddev=2):
     """
     Makes an image with a gaussian source that has a non-flat trace dispersed
     along the x axis.
     """
     spec2d = np.zeros((nrows, ncols))
-    spec2dvar = np.ones((nrows, ncols))
 
     for ii in range(spec2d.shape[1]):
         mgaus = models.Gaussian1D(amplitude=amp,
@@ -34,14 +34,13 @@ def mk_img_non_flat_trace(nrows=40, ncols=100, amp=10, stddev=2):
                                   stddev=stddev)
         rg = np.arange(0, spec2d.shape[0], 1)
         gaus = mgaus(rg)
-        spec2d[:,ii] = gaus
+        spec2d[:, ii] = gaus
 
     return spec2d
 
 
 class TestMeasureCrossDispersionProfile():
 
-
     @pytest.mark.parametrize('pixel', [None, 1, [1, 2, 3]])
     def test_measure_cross_dispersion_profile(self, pixel):
         """
@@ -59,7 +58,7 @@ def test_measure_cross_dispersion_profile(self, pixel):
         images.append(dat)  # test unitless
         images.append(dat * u.DN)
         images.append(NDData(dat * u.DN))
-        images.append(Spectrum1D(flux = dat * u.DN))
+        images.append(Spectrum1D(flux=dat * u.DN))
 
         for img in images:
 
@@ -104,7 +103,7 @@ def test_cross_dispersion_profile_non_flat_trace(self):
         # the shape of the trace
         peak_locs = [9, 10, 12, 13, 15, 16, 17, 19, 20, 22, 23, 24, 26, 27, 29]
         for i, pixel in enumerate(range(0, image.shape[1], 7)):
-            profile = measure_cross_dispersion_profile(image, 
+            profile = measure_cross_dispersion_profile(image,
                                                        trace=trace_fit,
                                                        width=None,
                                                        pixel=pixel,
@@ -116,10 +115,10 @@ def test_cross_dispersion_profile_non_flat_trace(self):
         # not change since (there is some wiggling around though due to the
         # fact that the trace is rolled to the nearest integer value. this can
         # be smoothed with an interpolation option later on, but it is 'rough'
-        # for now). In this test case, the peak positions will all either 
+        # for now). In this test case, the peak positions will all either
         # be at pixel 20 or 21.
         for i, pixel in enumerate(range(0, image.shape[1], 7)):
-            profile = measure_cross_dispersion_profile(image, 
+            profile = measure_cross_dispersion_profile(image,
                                                        trace=trace_fit,
                                                        width=None,
                                                        pixel=pixel,
@@ -128,7 +127,6 @@ def test_cross_dispersion_profile_non_flat_trace(self):
             peak_loc = (np.where(profile == max(profile))[0][0])
             assert peak_loc in [20, 21]
 
-
     def test_errors_warnings(self):
         img = mk_gaussian_img(nrows=10, ncols=10)
         with pytest.raises(ValueError,
diff --git a/specreduce/utils/utils.py b/specreduce/utils/utils.py
index 3690647..5489c34 100644
--- a/specreduce/utils/utils.py
+++ b/specreduce/utils/utils.py
@@ -52,8 +52,8 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
     crossdisp_axis : int, optional
         The index of the image's cross-dispersion axis. [default: 0]
     width : tuple of int or None
-        Width around 'trace' to calculate profile. If None, then all rows in the cross-dispersion axis
-        will be used. [default: None]
+        Width around 'trace' to calculate profile. If None, then all rows in the
+        cross-dispersion axis will be used. [default: None]
     pixel: int, list of int, or None
         Pixel value(s) along the dispersion axis to return cross-dispersion profile.
         If several specified in list, then the average (method set by `statistic`)
@@ -165,7 +165,7 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
     aligned_trace = None
     if align_along_trace:
         if not isinstance(trace, FlatTrace):
-            # note: image was transposed according to `crossdisp_axis`, so disp_axis will always be 1
+            # 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)