update spectral extraction with new model types

jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.py

@@ -13,6 +13,7 @@
 from jdaviz.core.custom_traitlets import IntHandleEmpty, FloatHandleEmpty
 from jdaviz.core.marks import PluginLine
 
+from astropy.modeling import models
 from astropy.nddata import NDData, StdDevUncertainty, VarianceUncertainty, UnknownUncertainty
 from specutils import Spectrum1D
 from specreduce import tracing
@@ -21,6 +22,10 @@
 
 __all__ = ['SpectralExtraction']
 
+_model_cls = {'Spline': models.Spline1D,
+              'Polynomial': models.Polynomial1D,
+              'Legendre': models.Legendre1D}
+
 
 @tray_registry('spectral-extraction', label="Spectral Extraction")
 class SpectralExtraction(PluginTemplateMixin):
@@ -100,6 +105,7 @@ class SpectralExtraction(PluginTemplateMixin):
     trace_type_selected = Unicode().tag(sync=True)
 
     trace_pixel = FloatHandleEmpty(0).tag(sync=True)
+    trace_order = IntHandleEmpty(3).tag(sync=True)
 
     trace_peak_method_items = List().tag(sync=True)
     trace_peak_method_selected = Unicode().tag(sync=True)
@@ -189,7 +195,8 @@ def __init__(self, *args, **kwargs):
         self.trace_type = SelectPluginComponent(self,
                                                 items='trace_type_items',
                                                 selected='trace_type_selected',
-                                                manual_options=['Flat', 'Auto'])
+                                                manual_options=['Flat', 'Polynomial',
+                                                                'Legendre', 'Spline'])
 
         self.trace_peak_method = SelectPluginComponent(self,
                                                        items='trace_peak_method_items',
@@ -276,7 +283,8 @@ def __init__(self, *args, **kwargs):
     @property
     def user_api(self):
         return PluginUserApi(self, expose=('interactive_extract',
-                                           'trace_dataset', 'trace_type', 'trace_peak_method',
+                                           'trace_dataset', 'trace_type',
+                                           'trace_order', 'trace_peak_method',
                                            'trace_pixel', 'trace_bins', 'trace_window',
                                            'import_trace',
                                            'export_trace',
@@ -377,7 +385,7 @@ def _update_plugin_marks(self, *args):
                 sp1d = self.export_extract_spectrum(add_data=False)
             except Exception as e:
                 # NOTE: ignore error, but will be raised when clicking ANY of the export buttons
-                # NOTE: KosmosTrace or manual background are often giving a
+                # NOTE: AutoTrace or manual background are often giving a
                 # "background regions overlapped" error from specreduce
                 self.ext_specreduce_err = repr(e)
                 self.marks['extract'].clear()
@@ -442,7 +450,7 @@ def marks(self):
         return self._marks
 
     @observe('trace_dataset_selected', 'trace_type_selected',
-             'trace_trace_selected', 'trace_offset',
+             'trace_trace_selected', 'trace_offset', 'trace_order',
              'trace_pixel', 'trace_peak_method_selected',
              'trace_bins', 'trace_window', 'active_step')
     def _interaction_in_trace_step(self, event={}):
@@ -575,7 +583,7 @@ def import_trace(self, trace):
         if isinstance(trace, tracing.FlatTrace):
             self.trace_type_selected = 'Flat'
             self.trace_pixel = trace.trace_pos
-        elif isinstance(trace, tracing.KosmosTrace):
+        elif isinstance(trace, tracing.AutoTrace):
             self.trace_type_selected = 'Auto'
             self.trace_pixel = trace.guess
             self.trace_window = trace.window
@@ -616,12 +624,14 @@ def export_trace(self, add_data=False, **kwargs):
             trace = tracing.FlatTrace(self.trace_dataset.selected_obj.data,
                                       self.trace_pixel)
 
-        elif self.trace_type_selected == 'Auto':
-            trace = tracing.KosmosTrace(self.trace_dataset.selected_obj.data,
-                                        guess=self.trace_pixel,
-                                        bins=int(self.trace_bins),
-                                        window=self.trace_window,
-                                        peak_method=self.trace_peak_method_selected.lower())
+        elif self.trace_type_selected in _model_cls:
+            trace_model = _model_cls[self.trace_type_selected](degree=self.trace_order)
+            trace = tracing.AutoTrace(self.trace_dataset.selected_obj.data,
+                                      guess=self.trace_pixel,
+                                      bins=int(self.trace_bins),
+                                      window=self.trace_window,
+                                      peak_method=self.trace_peak_method_selected.lower(),
+                                      trace_model=trace_model)
 
         else:
             raise NotImplementedError(f"trace_type={self.trace_type_selected} not implemented")

jdaviz/configs/specviz2d/plugins/spectral_extraction/spectral_extraction.vue

@@ -75,6 +75,20 @@
           ></v-select>
         </v-row>
 
+        <v-row v-if="trace_type_selected!=='Flat'">
+          <v-text-field
+            label="Order"
+            type="number"
+            v-model.number="trace_order"
+            :rules="[() => trace_order!=='' || 'This field is required',
+                     () => trace_order>=0 || 'Order must be positive',
+                     () => (trace_type_selected!=='Spline' || (trace_order > 0 && trace_order <= 5)) || 'Spline order must be between 1 and 5']"
+            hint="Order of the trace model."
+            persistent-hint
+          >
+          </v-text-field>
+        </v-row>
+
         <v-row>
           <v-text-field
             label="Pixel"
@@ -87,7 +101,7 @@
           </v-text-field>
         </v-row>
 
-        <v-row v-if="trace_type_selected==='Auto'">
+        <v-row v-if="trace_type_selected!=='Flat'">
           <v-text-field
             label="Bins"
             type="number"