From c7fe46fb26556abae5d3f68aea0172ca5d0a0009 Mon Sep 17 00:00:00 2001
From: Clare Shanahan <cshanahan@stsci.edu>
Date: Mon, 25 Mar 2024 00:31:51 -0400
Subject: [PATCH] added test for non-flat trace

---
 specreduce/tests/test_utils.py | 47 ++++++++++++++++++++++++++++++++--
 specreduce/utils/utils.py      | 10 +++++---
 2 files changed, 52 insertions(+), 5 deletions(-)

diff --git a/specreduce/tests/test_utils.py b/specreduce/tests/test_utils.py
index 7265e44..2407e9d 100644
--- a/specreduce/tests/test_utils.py
+++ b/specreduce/tests/test_utils.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 from astropy.modeling import fitting, models
-from specreduce.tracing import FlatTrace
+from specreduce.tracing import FitTrace, FlatTrace
 from specreduce.utils.utils import measure_cross_dispersion_profile
 from specutils import Spectrum1D
 from astropy.nddata import NDData
@@ -48,7 +48,7 @@ def test_measure_cross_dispersion_profile(self, pixel):
         Basic test for `measure_cross_dispersion_profile`. Parametrized over
         different options for `pixel` to test using all wavelengths, a single
         wavelength, and a set of wavelengths, as well as different input types
-        (plain array, quantity, Spectrum1D, and NDData)
+        (plain array, quantity, Spectrum1D, and NDData).
         """
 
         # test a few input formats
@@ -85,6 +85,49 @@ def test_measure_cross_dispersion_profile(self, pixel):
             assert fit_model.mean.value == mean
             assert fit_model.stddev.value == stddev
 
+    @pytest.mark.filterwarnings("ignore:Model is linear in parameters")
+    def test_cross_dispersion_profile_non_flat_trace(self):
+        """
+        Test measure_cross_dispersion_profile with a non-flat trace.
+        Tests with 'align_along_trace' set to both True and False,
+        to account for the changing center of the trace and measure
+        the true profile shape, or to 'blur' the profile, respectivley.
+        """
+
+        image = mk_img_non_flat_trace()
+
+        # fit the trace
+        trace_fit = FitTrace(image)
+
+        # when not aligning along trace and using the entire image
+        # rows for the window, the center of the profile should follow
+        # 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, 
+                                                       trace=trace_fit,
+                                                       width=None,
+                                                       pixel=pixel,
+                                                       statistic='mean')
+            peak_loc = (np.where(profile == max(profile))[0][0])
+            assert peak_loc == peak_locs[i]
+
+        # when align_along_trace = True, the shape of the profile should
+        # 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 
+        # be at pixel 20 or 21.
+        for i, pixel in enumerate(range(0, image.shape[1], 7)):
+            profile = measure_cross_dispersion_profile(image, 
+                                                       trace=trace_fit,
+                                                       width=None,
+                                                       pixel=pixel,
+                                                       align_along_trace=True,
+                                                       statistic='mean')
+            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)
diff --git a/specreduce/utils/utils.py b/specreduce/utils/utils.py
index f900300..3690647 100644
--- a/specreduce/utils/utils.py
+++ b/specreduce/utils/utils.py
@@ -175,9 +175,13 @@ def measure_cross_dispersion_profile(image, trace=None, crossdisp_axis=0,
             trace = FlatTrace(aligned_trace, trace_pos)
 
     # create a weight image based on the trace and 'width' to mask around trace
-    wimg = _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image.shape)
-    # invert mask to include, not exclude, pixels around trace
-    wimg = (1 - wimg).astype(int)
+
+    if width == nrows:
+        wimg = np.zeros(image.shape)
+    else:
+        wimg = _ap_weight_image(trace, width, disp_axis, crossdisp_axis, image.shape)
+        # invert mask to include, not exclude, pixels around trace
+        wimg = (1 - wimg).astype(int)
 
     # now that we have figured out the mask for the window in cross-disp. axis,
     # select only the pixel(s) we want to include in measuring the avg. profile