Merge pull request #239 from TheSkyentist/ctx

Support for Context Maps

grizli/aws/aws_drizzler.py

@@ -531,7 +531,7 @@ def segmentation_figure(label, cat, segfile):
     th.close()
 
 
-def drizzle_images(label='macs0647-jd1', ra=101.9822125, dec=70.24326667, pixscale=0.1, size=10, wcs=None, pixfrac=0.33, kernel='square', theta=0, half_optical_pixscale=True, filters=['f160w', 'f140w', 'f125w', 'f105w', 'f110w', 'f098m', 'f850lp', 'f814w', 'f775w', 'f606w', 'f475w', 'f555w', 'f600lp', 'f390w', 'f350lp'], skip=None, remove=True, rgb_params=RGB_PARAMS, master='grizli-v1-19.12.04', aws_bucket='s3://grizli/CutoutProducts/', scale_ab=21, thumb_height=2.0, sync_fits=True, subtract_median=True, include_saturated=True, include_ir_psf=False, oversample_psf=False, show_filters=['visb', 'visr', 'y', 'j', 'h'], combine_similar_filters=True, single_output=True, aws_prep_dir=None, make_segmentation_figure=False, get_dict=False, dryrun=False, thumbnail_ext='png', **kwargs):
+def drizzle_images(label='macs0647-jd1', ra=101.9822125, dec=70.24326667, pixscale=0.1, size=10, wcs=None, pixfrac=0.33, kernel='square', theta=0, half_optical_pixscale=True, filters=['f160w', 'f140w', 'f125w', 'f105w', 'f110w', 'f098m', 'f850lp', 'f814w', 'f775w', 'f606w', 'f475w', 'f555w', 'f600lp', 'f390w', 'f350lp'], skip=None, remove=True, rgb_params=RGB_PARAMS, master='grizli-v1-19.12.04', aws_bucket='s3://grizli/CutoutProducts/', scale_ab=21, thumb_height=2.0, sync_fits=True, subtract_median=True, include_saturated=True, include_ir_psf=False, oversample_psf=False, show_filters=['visb', 'visr', 'y', 'j', 'h'], combine_similar_filters=True, single_output=True, aws_prep_dir=None, make_segmentation_figure=False, get_dict=False, dryrun=False, thumbnail_ext='png', write_ctx=False, **kwargs):
     """
     label='cp561356'; ra=150.208875; dec=1.850241667; size=40; filters=['f160w','f814w', 'f140w','f125w','f105w','f606w','f475w']
 
@@ -799,7 +799,7 @@ def drizzle_images(label='macs0647-jd1', ra=101.9822125, dec=70.24326667, pixsca
             continue
 
         elif status is not None:
-            sci, wht, var, outh, filt_dict[filt], wcs_tab = status
+            sci, wht, var, ctx, outh, filt_dict[filt], wcs_tab = status
 
             if subtract_median:
                 #med = np.median(sci[sci != 0])
@@ -825,6 +825,10 @@ def drizzle_images(label='macs0647-jd1', ra=101.9822125, dec=70.24326667, pixsca
             pyfits.writeto('{0}-{1}_drz_wht.fits'.format(label, filt),
                            data=wht, header=outh, overwrite=True,
                            output_verify='fix')
+            if write_ctx:
+                pyfits.writeto('{0}-{1}_drz_ctx.fits'.format(label, filt),
+                            data=ctx, header=outh, overwrite=True,
+                            output_verify='fix')
 
             has_filts.append(filt)
 

grizli/aws/visit_processor.py

@@ -2031,7 +2031,7 @@ def query_exposures(ra=53.16, dec=-27.79, size=1., pixel_scale=0.1, theta=0,  fi
     return header, wcs, SQL, res
 
 
-def cutout_mosaic(rootname='gds', product='{rootname}-{f}', ra=53.1615666, dec=-27.7910651, size=5*60, theta=0., filters=['F160W'], ir_scale=0.1, ir_wcs=None, res=None, half_optical=True, kernel='point', pixfrac=0.33, make_figure=True, skip_existing=True, clean_flt=True, gzip_output=True, s3output='s3://grizli-v2/HST/Pipeline/Mosaic/', split_uvis=True, extra_query='', extra_wfc3ir_badpix=True, fix_niriss=True, scale_nanojy=10, verbose=True, weight_type='jwst_var', rnoise_percentile=99, get_dbmask=True, niriss_ghost_kwargs={}, snowblind_kwargs=None, scale_photom=True, context='jwst_0989.pmap', calc_wcsmap=False, make_exptime_map=False, expmap_sample_factor=4, keep_expmap_small=True, **kwargs):
+def cutout_mosaic(rootname='gds', product='{rootname}-{f}', ra=53.1615666, dec=-27.7910651, size=5*60, theta=0., filters=['F160W'], ir_scale=0.1, ir_wcs=None, res=None, half_optical=True, kernel='point', pixfrac=0.33, make_figure=True, skip_existing=True, clean_flt=True, gzip_output=True, s3output='s3://grizli-v2/HST/Pipeline/Mosaic/', split_uvis=True, extra_query='', extra_wfc3ir_badpix=True, fix_niriss=True, scale_nanojy=10, verbose=True, weight_type='jwst_var', rnoise_percentile=99, get_dbmask=True, niriss_ghost_kwargs={}, snowblind_kwargs=None, scale_photom=True, context='jwst_0989.pmap', calc_wcsmap=False, make_exptime_map=False, expmap_sample_factor=4, keep_expmap_small=True, write_ctx=False, **kwargs):
     """
     Make mosaic from exposures defined in the exposure database
     
@@ -2127,6 +2127,9 @@ def cutout_mosaic(rootname='gds', product='{rootname}-{f}', ra=53.1615666, dec=-
     calc_wcsmap : bool
         See `~grizli.utils.drizzle_from_visit`
     
+    write_ctx : bool
+        Optionally output context map
+    
     Returns
     -------
     Image mosaic files with the requested WCS and filters
@@ -2318,7 +2321,7 @@ def cutout_mosaic(rootname='gds', product='{rootname}-{f}', ra=53.1615666, dec=-
             **kwargs,
         )
 
-        outsci, outwht, outvar, header, flist, wcs_tab = _
+        outsci, outwht, outvar, outctx, header, flist, wcs_tab = _
 
         wcs_tab.write('{0}_wcs.csv'.format(visit['product']), overwrite=True)
 
@@ -2377,6 +2380,11 @@ def cutout_mosaic(rootname='gds', product='{rootname}-{f}', ra=53.1615666, dec=-
                 overwrite=True
             )
 
+        if write_ctx:
+            pyfits.writeto('{0}_{1}_ctx.fits'.format(visit['product'], drz),
+                        data=outctx, header=header, 
+                        overwrite=True)
+
         if make_exptime_map:
             matched_exptime_map(
                 sci_file,
@@ -2727,7 +2735,7 @@ def show_epochs_filter(ra, dec, size=4, filters=['F444W-CLEAR'], cleanup=True, v
     return fig
 
 
-def make_mosaic(jname='', ds9=None, skip_existing=True, ir_scale=0.1, half_optical=False, pad=16, kernel='point', pixfrac=0.33, sync=True, ir_wcs=None, weight_type='jwst_var', **kwargs):
+def make_mosaic(jname='', ds9=None, skip_existing=True, ir_scale=0.1, half_optical=False, pad=16, kernel='point', pixfrac=0.33, sync=True, ir_wcs=None, weight_type='jwst_var', write_ctx=False, **kwargs):
     """
     Make mosaics from all exposures in a group of associations
     """
@@ -2912,7 +2920,7 @@ def make_mosaic(jname='', ds9=None, skip_existing=True, ir_scale=0.1, half_optic
             **kwargs
         )
 
-        outsci, outwht, outvar, header, flist, wcs_tab = _
+        outsci, outwht, outvar,outctx, header, flist, wcs_tab = _
 
         pyfits.writeto(
             groups[f]['product']+'_drz_sci.fits',
@@ -2935,7 +2943,12 @@ def make_mosaic(jname='', ds9=None, skip_existing=True, ir_scale=0.1, half_optic
                 header=header,
                 overwrite=True
             )
-
+
+        if write_ctx:
+            pyfits.writeto(groups[f]['product']+'_drz_ctx.fits',
+                        data=outctx, header=header, 
+                        overwrite=True)
+
     if ds9 is not None:
         auto_script.field_rgb(base, HOME_PATH=None, xsize=12, output_dpi=300, 
               ds9=ds9, scl=2, suffix='.rgb', timestamp=True, mw_ebv=0)

grizli/multifit.py

@@ -1002,7 +1002,7 @@ def make_stack(self, id, size=20, target='grism', skip=True, fcontam=1., scale=1
         return hdu, fig
 
 
-    def drizzle_grism_models(self, root='grism_model', kernel='square', scale=0.1, pixfrac=1, make_figure=True, fig_xsize=10):
+    def drizzle_grism_models(self, root='grism_model', kernel='square', scale=0.1, pixfrac=1, make_figure=True, fig_xsize=10, write_ctx=False):
         """
         Make model-subtracted drizzled images of each grism / PA
 
@@ -1020,6 +1020,9 @@ def drizzle_grism_models(self, root='grism_model', kernel='square', scale=0.1, p
         pixfrac : float
             Drizzle "pixfrac".
 
+        write_ctx : bool
+            Save context image as well.
+
         """
         try:
             from .utils import drizzle_array_groups
@@ -1071,12 +1074,21 @@ def drizzle_grism_models(self, root='grism_model', kernel='square', scale=0.1, p
                     pixfrac=pixfrac
                 )
 
-                outsci, _, _, _, header, outputwcs = out
+                outsci, _, _, outctx, header, outputwcs = out
                 header['FILTER'] = g
                 header['PA'] = pa
+
+                # Add invidual FLTs to header
+                for i, ix in enumerate(idx):
+                    header[f'FLT{str(i+1).zfill(5)}'] = self.FLTs[ix].grism_file
+
                 pyfits.writeto(outfile, data=outsci, header=header,
                                overwrite=True, output_verify='fix')
 
+                if write_ctx:
+                    pyfits.writeto(outfile.replace('sci', 'ctx'), data=outctx,
+                                header=header, overwrite=True, output_verify='fix')
+
                 # Model-subtracted
                 outfile = '{0}-{1}-{2}_grism_clean.fits'.format(root, g.lower(),
                                                               pa)

grizli/prep.py

@@ -3901,14 +3901,13 @@ def subtract_visit_angle_averages(visit, threshold=1.8, detection_background=Tru
             weight_type=weight_type,
         )
 
-        drz_h = drz[3]
+        data, wht, ctx, var, drz_h, files, info = drz
 
         drz_h['PA_APER'] = pa_aper, 'PA_APER used for reference'
         drz_h['BKGANGLE'] = (','.join([f'{a:.1f}' for a in angles]), 
                              'Background angles')
         drz_h['NITANGLE'] = niter, 'Number of iterations for angle subtraction'
 
-        data, wht = drz[0], drz[1]
         err = 1/np.sqrt(wht)
         mask = (wht > 0)
 
@@ -4001,7 +4000,7 @@ def subtract_visit_angle_averages(visit, threshold=1.8, detection_background=Tru
             drz_hdu.append(pyfits.ImageHDU(data=smask*1, 
                                            header=drz_h,
                                           )
-                           )            
+                           )
         else:
             drz_hdu[0].data += med_model.astype(np.float32)
             drz_hdu[1].data = _cleaned.astype(np.float32)
@@ -4396,6 +4395,7 @@ def process_direct_grism_visit(direct={},
                                miri_skyfile=None,
                                use_skyflats=True,
                                do_pure_parallel_wcs=True,
+                               write_ctx=False,
                                ):
     """Full processing of a direct (+grism) image visit.
     
@@ -5023,14 +5023,17 @@ def process_direct_grism_visit(direct={},
             calc_wcsmap=False,
         )
 
-        _sci, _wht, _var, _hdr, _files, _info = _
+        _sci, _wht, _var, _ctx, _hdr, _files, _info = _
         _drcfile = glob.glob(f"{direct['product']}_dr*sci.fits")[0]
         _whtfile = _drcfile.replace('_sci.fits','_wht.fits')
+        _ctxfile = _drcfile.replace('_sci.fits','_ctx.fits')
         pyfits.PrimaryHDU(data=_sci, header=_hdr).writeto(_drcfile,
                                                           overwrite=True)
         pyfits.PrimaryHDU(data=_wht, header=_hdr).writeto(_whtfile,
                                                           overwrite=True)
-
+        if write_ctx:
+            pyfits.PrimaryHDU(data=_ctx, header=_hdr).writeto(_ctxfile,
+                                                            overwrite=True)
         # Remake catalog
         #cat = make_drz_catalog(root=direct['product'], threshold=thresh)
         if 'footprints' in direct:

grizli/utils.py

@@ -8235,7 +8235,7 @@ def drizzle_from_visit(
             # outvar = varnum / outwht
             outvar *= pscale_ratio**-2
 
-        return outsci, outwht, outvar, header, flist, wcs_tab
+        return outsci, outwht, outvar, outctx, header, flist, wcs_tab
 
 
 def drizzle_array_groups(
@@ -8251,6 +8251,7 @@ def drizzle_array_groups(
     calc_wcsmap=False,
     verbose=True,
     data=None,
+    first_uniqid=1,
 ):
     """
     Drizzle array data with associated wcs
@@ -8299,6 +8300,9 @@ def drizzle_array_groups(
 
         If not provided, new arrays will be created.
 
+    first_uniqid : int, optional
+        First `uniqid` value to use for the drizzle for contex maps
+
     Returns
     -------
     outsci, outwht, outctx : `~numpy.ndarray`
@@ -8399,8 +8403,13 @@ def drizzle_array_groups(
 
     needs_var = (outvar is not None) & (var_list is not None)
 
-    # Do drizzle
+    # Number of input arrays
     N = len(sci_list)
+
+    # Drizzlepac cannot support >31 input images
+    if first_uniqid + N > 31 and verbose:
+        msg = "Warning: Too many input images for context map, will wrap around"
+        log_comment(LOGFILE, msg, verbose=verbose, show_date=True)
     for i in range(N):
         if verbose:
             # log.info('Drizzle array {0}/{1}'.format(i+1, N))
@@ -8425,7 +8434,7 @@ def drizzle_array_groups(
             "cps",
             1,
             wcslin_pscale=wcs_list[i].pscale,
-            uniqid=1,
+            uniqid=((first_uniqid - 1 + i) % 32) + 1,
             pixfrac=pixfrac,
             kernel=kernel,
             fillval="0",