initialize spectrogram starfield and flat

spectractor/extractor/extractor.py

@@ -1352,10 +1352,7 @@ def bgd_model_func(x, y):
     spectrum.chromatic_psf = s
 
     # Extract the spectrogram edges
-    right_edge = image.data.shape[1]
-    data = np.copy(image.data)
-    err = np.copy(image.err)
-    Ny, Nx = data.shape
+    Ny, Nx = image.data.shape
     x0 = int(image.target_pixcoords[0])
     y0 = int(image.target_pixcoords[1])
     ymax = min(Ny, y0 + int(s.table['Dy_disp_axis'].max()) + ws[1] + 1)  # +1 to  include edges
@@ -1365,7 +1362,7 @@ def bgd_model_func(x, y):
     lambda_min_index = int(np.argmin(np.abs(lambdas[::np.sign(spectrum.order)] - parameters.LAMBDA_MIN)))
     lambda_max_index = int(np.argmin(np.abs(lambdas[::np.sign(spectrum.order)] - parameters.LAMBDA_MAX)))
     xmin = max(0, int(s.table['Dx'][lambda_min_index] + x0))
-    xmax = min(right_edge, int(s.table['Dx'][lambda_max_index] + x0) + 1)  # +1 to  include edges
+    xmax = min(Nx, int(s.table['Dx'][lambda_max_index] + x0) + 1)  # +1 to  include edges
     # Position of the order 0 in the spectrogram coordinates
     target_pixcoords_spectrogram = [image.target_pixcoords[0] - xmin, image.target_pixcoords[1] - ymin]
     s.y0 = target_pixcoords_spectrogram[1]
@@ -1378,13 +1375,18 @@ def bgd_model_func(x, y):
                     f"\n{s.table[['amplitude', 'x_c', 'y_c', 'Dx', 'Dy', 'Dy_disp_axis']]}")
 
     # Create spectrogram
-    data = data[ymin:ymax, xmin:xmax]
-    err = err[ymin:ymax, xmin:xmax]
-
-    spectrum.spectrogram_data = data
-    spectrum.spectrogram_err = err
+    spectrum.spectrogram_data = np.copy(image.data[ymin:ymax, xmin:xmax])
+    spectrum.spectrogram_err = np.copy(image.err[ymin:ymax, xmin:xmax])
+    if image.starfield is not None:
+        spectrum.spectrogram_starfield = np.copy(image.starfield[ymin:ymax, xmin:xmax])
+    else:
+        spectrum.spectrogram_starfield = None  # np.zeros_like(spectrum.spectrogram_data)
+    if image.flat is not None:
+        spectrum.spectrogram_flat = np.copy(image.flat[ymin:ymax, xmin:xmax])
+    else:
+        spectrum.spectrogram_flat = None  # np.ones_like(spectrum.spectrogram_data)
 
-    Ny, Nx = data.shape
+    Ny, Nx = spectrum.spectrogram_data.shape
     my_logger.info(f'\n\tExtract spectrogram: crop raw image [{xmin}:{xmax},{ymin}:{ymax}] (size ({Nx}, {Ny}))')
 
     # Extract the non-rotated background
@@ -1448,7 +1450,8 @@ def bgd_model_func(x, y):
         gs_kw = dict(width_ratios=[3, 0.08], height_ratios=[1, 1])
         fig, ax = plt.subplots(2, 2, sharex='none', figsize=(16, 6), gridspec_kw=gs_kw)
         xx = np.arange(s.table['Dx'].size)
-        plot_image_simple(ax[1, 0], data=data, scale="symlog", title='', units=image.units, aspect='auto', cax=ax[1, 1])
+        plot_image_simple(ax[1, 0], data=spectrum.spectrogram_data, scale="symlog", title='',
+                          units=image.units, aspect='auto', cax=ax[1, 1])
         ax[1, 0].plot(xx, target_pixcoords_spectrogram[1] + s.table['Dy_disp_axis'], label='Dispersion axis', color="r")
         ax[1, 0].scatter(xx, target_pixcoords_spectrogram[1] + s.table['Dy'],
                          c=s.table['lambdas'], edgecolors='None', cmap=from_lambda_to_colormap(s.table['lambdas']),

spectractor/extractor/images.py

@@ -44,7 +44,7 @@ class Image(object):
         Rotated image 2D uncertainty array in self.units units.
     flat: array
         Flat 2D array without units and median of 1.
-    star_field: array
+    starfield: array
         Star field simulation, no units needed but better in ADU/s.
     target_pixcoords_rotated: array
         Target position [x,y] in the rotated image in pixels.
@@ -159,9 +159,7 @@ def __init__(self, file_name, *, target_label="", disperser_label="",
         self.humidity = 0
 
         self.flat = None
-        self.star_field = None
-
-        self.imgs = [self.data, self.err, self.flat, self.star_field]
+        self.starfield = None
 
         if parameters.CALLING_CODE != 'LSST_DM' and file_name != "":
             self.load_image(file_name)

spectractor/extractor/spectrum.py

@@ -243,6 +243,8 @@ def __init__(self, file_name="", image=None, order=1, target=None, config="", fa
         self.spectrogram_err = None
         self.spectrogram_residuals = None
         self.spectrogram_fit = None
+        self.spectrogram_flat = None
+        self.spectrogram_starfield = None
         self.spectrogram_x0 = None
         self.spectrogram_y0 = None
         self.spectrogram_xmin = None
@@ -700,6 +702,10 @@ def save_spectrum(self, output_file_name, overwrite=False):
                 hdus[extname].data = self.spectrogram_fit
             elif extname == "S_RES":
                 hdus[extname].data = self.spectrogram_residuals
+            elif extname == "S_FLAT" and self.spectrogram_flat is not None:
+                hdus[extname].data = self.spectrogram_flat
+            elif extname == "S_STAR" and self.spectrogram_starfield is not None:
+                hdus[extname].data = self.spectrogram_starfield
             elif extname == "PSF_TAB":
                 hdus[extname] = fits.table_to_hdu(self.chromatic_psf.table)
             elif extname == "LINES":
@@ -1080,6 +1086,10 @@ def load_spectrum_latest(self, input_file_name):
                 self.spectrogram_bgd_rms = hdu_list["S_BGD_ER"].data
                 self.spectrogram_fit = hdu_list["S_FIT"].data
                 self.spectrogram_residuals = hdu_list["S_RES"].data
+                if "S_FLAT" in hdu_list.keys():
+                    self.spectrogram_flat = hdu_list["S_FLAT"].data
+                if "S_STAR" in hdu_list.keys():
+                    self.spectrogram_starfield = hdu_list["S_STAR"].data
                 self.chromatic_psf.init_from_table(Table.read(hdu_list["PSF_TAB"]),
                                                    saturation=self.spectrogram_saturation)
                 self.lines.table = Table.read(hdu_list["LINES"], unit_parse_strict="silent")