Have WFSS simulations optionally use pickle files of sources

CHANGES.rst

@@ -21,9 +21,16 @@ Updated the MOVING_TARGET_POSITION table from detector -> sci. (#757)
 
 Galaxy and Extended Source Position Angles
 ------------------------------------------
+
 Fixed a bug that was causing incorrect position angles of galaxy and extended sources when PAV3 was non-zero. (#767)
 
 
+WFSS
+----
+
+Save "seed_cube" pickle files of point source, galaxy, and extended sources. These are then picked up by the disperser. This will allow overlapping sources to be treated correctly when everything is disersed. (#779)
+
+
 2.2.1
 =====
 

mirage/seed_image/catalog_seed_image.py

@@ -14,6 +14,7 @@
 import logging
 import os
 import copy
+import pickle
 import re
 import shutil
 from yaml.scanner import ScannerError
@@ -53,7 +54,8 @@
 from ..utils.constants import CRDS_FILE_TYPES, MEAN_GAIN_VALUES, SERSIC_FRACTIONAL_SIGNAL, \
                               SEGMENTATION_MIN_SIGNAL_RATE, SUPPORTED_SEGMENTATION_THRESHOLD_UNITS, \
                               LOG_CONFIG_FILENAME, STANDARD_LOGFILE_NAME, TSO_MODES, NIRISS_GHOST_GAP_FILE, \
-                              NIRISS_GHOST_GAP_URL, NIRCAM_SW_GRISMTS_APERTURES, NIRCAM_LW_GRISMTS_APERTURES
+                              NIRISS_GHOST_GAP_URL, NIRCAM_SW_GRISMTS_APERTURES, NIRCAM_LW_GRISMTS_APERTURES, \
+                              DISPERSED_MODES
 from ..utils.flux_cal import fluxcal_info, sersic_fractional_radius, sersic_total_signal
 from ..utils.timer import Timer
 from ..psf.psf_selection import get_gridded_psf_library, get_psf_wings
@@ -2956,6 +2958,9 @@ def make_point_source_image(self, pointSources, segment_number=None, ptsrc_segma
         # Create the empty image
         psfimage = np.zeros(self.output_dims)
 
+        # Create empty seed cube for possible WFSS dispersion
+        seed_cube = {}
+
         if ptsrc_segmap is None:
             # Create empty segmentation map
             ptsrc_segmap = segmap.SegMap()
@@ -3015,12 +3020,21 @@ def make_point_source_image(self, pointSources, segment_number=None, ptsrc_segma
                 self.timer.stop()
                 continue
 
+            self.logger.info("******************************* %s" % (self.basename))
+
             try:
                 psf_to_add = scaled_psf[l1:l2, k1:k2]
                 psfimage[j1:j2, i1:i2] += psf_to_add
 
                 # Add source to segmentation map
                 ptsrc_segmap.add_object_threshold(psf_to_add, j1, i1, entry['index'], self.segmentation_threshold)
+
+                if self.params['Inst']['mode'] in DISPERSED_MODES:
+                    # Add source to seed cube file
+                    stamp = np.zeros(psf_to_add.shape)
+                    flag = psf_to_add >= self.segmentation_threshold
+                    stamp[flag] = entry['index']
+                    seed_cube[entry['index']] = [i1, j1, psf_to_add*1, stamp*1]
             except IndexError:
                 # In here we catch sources that are off the edge
                 # of the detector. These may not necessarily be caught in
@@ -3042,6 +3056,10 @@ def make_point_source_image(self, pointSources, segment_number=None, ptsrc_segma
                     self.logger.info(('Working on source #{}. Estimated time remaining to add all point sources to the stamp image: {} minutes. '
                                       'Projected finish time: {}'.format(i, time_remaining, finish_time)))
 
+        if self.params['Inst']['mode'] in DISPERSED_MODES:
+            # Save the seed cube file of point sources
+            pickle.dump(seed_cube, open("%s_star_seed_cube.pickle" % (self.basename), "wb"), protocol=pickle.HIGHEST_PROTOCOL)
+
         return psfimage, ptsrc_segmap
 
     def create_psf_stamp(self, x_location, y_location, psf_dim_x, psf_dim_y,
@@ -4199,6 +4217,9 @@ def make_galaxy_image(self, file):
         # final output image
         yd, xd = self.output_dims
 
+        # Seed cube for disperser
+        seed_cube = {}
+
         # create the final galaxy countrate image
         galimage = np.zeros((yd, xd))
 
@@ -4295,6 +4316,13 @@ def make_galaxy_image(self, file):
                     # Add source to segmentation map
                     segmentation.add_object_threshold(stamp_to_add, j1, i1, entry['index'], self.segmentation_threshold)
 
+                    if self.params['Inst']['mode'] in DISPERSED_MODES:
+                        # Add source to the seed cube
+                        stamp = np.zeros(stamp_to_add.shape)
+                        flag = stamp_to_add >= self.segmentation_threshold
+                        stamp[flag] = entry['index']
+                        seed_cube[entry['index']] = [i1, j1, stamp_to_add*1, stamp*1]
+
                 else:
                     pass
                     # print("Source located entirely outside the field of view. Skipping.")
@@ -4313,6 +4341,10 @@ def make_galaxy_image(self, file):
                     self.logger.info(('Working on galaxy #{}. Estimated time remaining to add all galaxies to the stamp image: {} minutes. '
                                       'Projected finish time: {}'.format(entry_index, time_remaining, finish_time)))
 
+        if self.params['Inst']['mode'] in DISPERSED_MODES:
+            # Save the seed cube file of galaxy sources
+            pickle.dump(seed_cube, open("%s_galaxy_seed_cube.pickle" % (self.basename), "wb"), protocol=pickle.HIGHEST_PROTOCOL)
+
         return galimage, segmentation.segmap, ghost_sources_from_galaxies
 
     def calc_x_position_angle(self, position_angle):
@@ -4791,6 +4823,9 @@ def make_extended_source_image(self, extSources, extStamps, extConvolutions):
         yd, xd = self.output_dims
         extimage = np.zeros(self.output_dims)
 
+        # Prepare seed cube for extended sources
+        seed_cube = {}
+
         # Create corresponding segmentation map
         segmentation = segmap.SegMap()
         segmentation.xdim = xd
@@ -4883,8 +4918,20 @@ def make_extended_source_image(self, extSources, extStamps, extConvolutions):
                 # Add source to segmentation map
                 segmentation.add_object_threshold(stamp_to_add, j1, i1, entry['index'],
                                                   self.segmentation_threshold)
+
+                if self.params['Inst']['mode'] in DISPERSED_MODES:
+                    # Add source to seed cube
+                    stamp = np.zeros(stamp_to_add.shape)
+                    flag = stamp_to_add >= self.segmentation_threshold
+                    stamp[flag] = entry['index']
+                    seed_cube[entry['index']] = [i1, j1, stamp_to_add*1, stamp*1]
+
                 self.n_extend += 1
 
+        if self.params['Inst']['mode'] in DISPERSED_MODES:
+            # Save the seed cube
+            pickle.dump(seed_cube, open("%s_extended_seed_cube.pickle" % (self.basename), "wb"), protocol=pickle.HIGHEST_PROTOCOL)
+
         if self.n_extend == 0:
             self.logger.info("No extended sources present within the aperture.")
         else:

mirage/utils/constants.py

@@ -25,6 +25,8 @@
                        "wfss": "NIS_WFSS", "soss": "NIS_SOSS"},
             "fgs": {"imaging": "FGS_IMAGE"}}
 
+DISPERSED_MODES = ['wfss', 'ts_grism']
+
 # Number of detector resets prior to the start of an exposure
 NUM_RESETS_BEFORE_EXP = {"nircam": {"full": 0, "sub": 1},
                          "niriss": {"full": 0, "sub": 1},