From 3b60b71041282963b1d4d7585f3b03e798a63411 Mon Sep 17 00:00:00 2001
From: Kyle Conroy <kyleconroy@gmail.com>
Date: Thu, 11 Aug 2022 08:34:35 -0400
Subject: [PATCH] more advanced input parameter estimation

* to handle cases where extraction or background regions otherwise would default to outside of the image
---
 .../spectral_extraction.py                    | 23 +++++++++++++++----
 .../tests/test_spectral_extraction.py         | 11 +++++++++
 2 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.py b/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.py
index aa91e75015..be021bafa9 100644
--- a/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.py
+++ b/jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.py
@@ -148,8 +148,17 @@ def _trace_dataset_selected(self, msg=None):
         width = self.trace_dataset.selected_obj.shape[0]
         # estimate the pixel number by taking the median of the brightest pixel index in each column
         brightest_pixel = int(np.median(np.argmax(self.trace_dataset.selected_obj.flux, axis=0)))
-        # default width will be 10% of cross-dispersion "height"
-        default_width = int(np.ceil(width / 10))
+        # do not allow to be an edge pixel
+        if brightest_pixel < 1:
+            brightest_pixel = 1
+        if brightest_pixel > width - 1:
+            brightest_pixel = width - 1
+        distance_from_edge = min(brightest_pixel, width-brightest_pixel)
+        # default width will be 10% of cross-dispersion "height",
+        # but no larger than distance from the edge
+        default_bg_width = int(np.ceil(width / 10))
+        default_width = min(default_bg_width, distance_from_edge * 2)
+
         if self.trace_pixel == 0:
             self.trace_pixel = brightest_pixel
         if self.trace_window == 0:
@@ -157,9 +166,15 @@ def _trace_dataset_selected(self, msg=None):
         if self.bg_trace_pixel == 0:
             self.bg_trace_pixel = brightest_pixel
         if self.bg_separation == 0:
-            self.bg_separation = default_width * 2
+            if default_bg_width * 2 > distance_from_edge:
+                self.bg_type_selected = 'OneSided'
+                # we want positive separation if brightest_pixel near bottom
+                sign = 1 if (brightest_pixel < width / 2) else -1
+                self.bg_separation = sign * default_bg_width * 2
+            else:
+                self.bg_separation = default_bg_width * 2
         if self.bg_width == 0:
-            self.bg_width = default_width
+            self.bg_width = default_bg_width
         if self.ext_width == 0:
             self.ext_width = default_width
 
diff --git a/jdaviz/configs/specviz2d/plugins/spectral_extraction/tests/test_spectral_extraction.py b/jdaviz/configs/specviz2d/plugins/spectral_extraction/tests/test_spectral_extraction.py
index 78cb469105..98c61e526a 100644
--- a/jdaviz/configs/specviz2d/plugins/spectral_extraction/tests/test_spectral_extraction.py
+++ b/jdaviz/configs/specviz2d/plugins/spectral_extraction/tests/test_spectral_extraction.py
@@ -82,3 +82,14 @@ def test_plugin(specviz2d_helper):
     pext.bg_results_label = 'should not be created'
     pext.vue_create_bg()
     assert 'should not be created' not in [d.label for d in specviz2d_helper.app.data_collection]
+
+
+@pytest.mark.remote_data
+def test_spectrum_on_top(specviz2d_helper):
+    fn = download_file('https://mast.stsci.edu/api/v0.1/Download/file/?uri=mast:jwst/product/jw01529-c1002_t002_miri_p750l_s2d.fits', cache=True)  # noqa
+
+    specviz2d_helper.load_data(spectrum_2d=fn)
+
+    pext = specviz2d_helper.app.get_tray_item_from_name('spectral-extraction')
+    assert pext.bg_type_selected == 'OneSided'
+    assert pext.bg_separation < 0