Updates on masks and windows scripts

Author: Merry Duparc

project/SO/pISO/python/get_windows.py

@@ -26,38 +26,25 @@
 apod_pts_source_degree = d["apod_pts_source_degree"]
 # the apodisation length of the final survey mask
 apod_survey_degree = d["apod_survey_degree"]
-# the threshold on the amount of cross linking to keep the data in
-cross_link_threshold = d["cross_link_threshold"]
-# pixel weight with inverse variance above n_ivar * median are set to ivar
-# this ensure that the window is not totally dominated by few pixels with too much weight
-n_med_ivar = d["n_med_ivar"]
 
 # we will skip the edges of the survey where the noise is very difficult to model, the default is to skip 0.5 degree for
 # constructing the kspace mask and 2 degree for the final survey mask, this parameter can be used to rescale the default
 rescale = d["edge_skip_rescale"]
 
 window_dir = d["window_dir"]
-skip_ivar = True
 
 pspy_utils.create_directory(window_dir)
 
-patch = None
-if "patch" in d:
-    patch = so_map.read_map(d["patch"])
-
-
-# here we list the different windows that need to be computed, we will then do a MPI loops over this list
-
-# Use this if you want to only compute one window
+# Use this if you want to only compute one window (for testing)
 # d["surveys"] = ['SO']
 # d["arrays_SO"] = ['i1_f090']
 
-
+# here we list the different windows that need to be computed, we will then do a MPI loops over this list
 n_wins, sv_list, ar_list = pspipe_list.get_arrays_list(d)
 
 log.info(f"number of windows to compute : {n_wins}")
 
-my_masks = {}
+my_masks: dict[so_map.so_map] = {}
 
 so_mpi.init(True)
 subtasks = so_mpi.taskrange(imin=0, imax=n_wins - 1)
@@ -78,13 +65,8 @@
 
     maps = d[f"maps_{sv}_{ar}"]
 
-    if not skip_ivar:
-        ivar_all = gal_mask.copy()
-        ivar_all.data[:] = 0
-
     # the first step is to iterate on the maps of a given array to identify pixels with zero values
     # we will form a first survey mask as the union of all these split masks
-    # we also compute the average ivar map
 
     log.info(f"[{task}] create survey mask from ivar maps")
 
@@ -94,15 +76,6 @@
         else:
             index = map.find("map.fits")
 
-        if not skip_ivar:
-            ivar_map = map[:index] + "ivar.fits"
-            ivar_map = so_map.read_map(ivar_map, geometry=template_geom)
-            my_masks["baseline"].data[ivar_map.data[:] == 0.0] = 0.0
-            ivar_all.data[:] += ivar_map.data[:]
-
-    if not skip_ivar:
-        ivar_all.data[:] /= np.max(ivar_all.data[:])
-
     if d[f"extra_mask_{sv}_{ar}"] is not None:
         log.info(f"[{task}] apply extra mask")
         extra_mask = so_map.read_map(d[f"extra_mask_{sv}_{ar}"], geometry=template_geom)
@@ -121,10 +94,6 @@
     log.info(f"[{task}] apply galactic mask")
     my_masks["baseline"].data *= gal_mask.data
 
-    if patch is not None:
-        log.info(f"[{task}] apply patch mask")
-
-        my_masks["baseline"].data *= patch.data
     # with this we can create the k space mask this will only be used for applying the kspace filter
     log.info(
         f"[{task}] appodize kspace mask with {rescale:.2f} apod and write it to disk"
@@ -174,35 +143,18 @@
 
         my_masks[mask_type].write_map(f"{window_dir}/window_{sv}_{ar}_{mask_type}.fits")
 
-        # we also make a version of the windows taking into account the ivar of the maps
-        if not skip_ivar:
-            log.info(f"[{task}] include ivar ")
-
-            mask_type_w = mask_type + "_ivar"
-
-            my_masks[mask_type_w] = my_masks[mask_type].copy()
-            id = np.where(ivar_all.data[:] * my_masks[mask_type_w].data[:] != 0)
-            med = np.median(ivar_all.data[id])
-            ivar_all.data[ivar_all.data[:] > n_med_ivar * med] = n_med_ivar * med
-            my_masks[mask_type_w].data[:] *= ivar_all.data[:]
-
-            my_masks[mask_type_w].data = my_masks[mask_type_w].data.astype(np.float32)
-            my_masks[mask_type_w].write_map(
-                f"{window_dir}/window_{sv}_{ar}_{mask_type_w}.fits"
-            )
-
     for mask_type, mask in my_masks.items():
         log.info(f"[{task}] downgrade and plot {mask_type} ")
-        mask = mask.downgrade(4)
-        mask.plot(file_name=f"{window_dir}/window_{sv}_{ar}_{mask_type}")
+        mask.downgrade(4).plot(file_name=f"{window_dir}/window_{sv}_{ar}_{mask_type}")
 
     if f"{sv}_{ar}" in d["plot_windowed_maps"]:
-        pspy_utils.create_directory(f"{window_dir}/windowed maps")
-        maps_to_plot = so_map.read_map(d[f"maps_{sv}_{ar}"][0])
-        maps_to_plot.data *= my_masks["baseline"].data
-        maps_to_plot.downgrade(2).calibrate(
-            cal=d[f"cal_{sv}_{ar}"], pol_eff=d[f"pol_eff_{sv}_{ar}"]
-        ).plot(
-            file_name=f"{window_dir}/windowed maps/{sv}_{ar}",
-            color_range=(300, 100, 100),
-        )
+        pspy_utils.create_directory(f"{window_dir}/windowed_maps")
+        for s, sv_ar_split in enumerate(d[f"maps_{sv}_{ar}"]):
+            maps_to_plot = so_map.read_map(sv_ar_split)
+            maps_to_plot.data *= my_masks["baseline"].data
+            maps_to_plot.downgrade(4).calibrate(
+                cal=d[f"cal_{sv}_{ar}"], pol_eff=d[f"pol_eff_{sv}_{ar}"]
+            ).plot(
+                file_name=f"{window_dir}/windowed_maps/{sv}_{ar}_split{s}",
+                color_range=(300, 100, 100),
+            )

project/SO/pISO/python/masks/get_xtra_mask.py

@@ -3,7 +3,7 @@
 Optionnally plots all ivar and maps
 """
 
-from pspy import so_dict, so_map, so_window
+from pspy import so_dict, so_map, so_window, so_mpi
 from pspipe_utils import log
 
 from pixell import enmap, enplot      
@@ -26,9 +26,8 @@
     os.makedirs(maps_ivar_dir, exist_ok=True)
 
 # get reasonable ivar, top 90% of nonzero values seems to work decently
-# TODO: make dynamic
-ivar_smooth_deg = 0.2
-ivar_quantile = d['xtra_mask_ivar_quantile']
+ivar_smooth_deg = d['ivar_smooth_deg']
+ivar_quantile = d['ivar_quantile']
 
 mask_intersect = True # intersection of every single mask
 mask_union = False # union of every single mask
@@ -42,29 +41,39 @@
 else:
     additional_mask = True
 
-for sv in d['surveys']:
-    for m in d[f'arrays_{sv}']:
-        
+n_ar = len(d['arrays_SO'])
+
+log.info(f"number of arrays for the mpi loop : {n_ar}")
+so_mpi.init(True)
+subtasks = so_mpi.taskrange(imin=0, imax=n_ar-1)
+
+sv = 'SO'   # we want to run tthise script only on SO
+
+for task in subtasks:
+        m = d[f'arrays_{sv}'][task]
         ivar_mask = True # intersection of masks just for this map (over splits)
 
         map_fns = d[f'maps_{sv}_{m}']
         for i, map_fn in enumerate(map_fns):
 
             map_dir_fn, map_base_fn = os.path.split(map_fn)
             if d[f"src_free_maps_{sv}"] == True:
-                ivar_base_fn = map_base_fn.replace('map_srcfree', 'ivar')
+                ivar_fn = map_fn.replace('_map_srcfree', '_ivar')
             else:
-                ivar_base_fn = map_base_fn.replace('map', 'ivar')
-            ivar_fn = os.path.join(map_dir_fn, ivar_base_fn)   
+                ivar_fn = map_fn.replace('_map', '_ivar')   
 
             # mask is based on the smoothed ivar map
             # only the pixels were the original ivar were nonzero though
-            ivar = enmap.read_map(ivar_fn) 
-            ivar_smooth = enmap.smooth_gauss(ivar, np.deg2rad(ivar_smooth_deg))
-            ivar_set_mask = ivar_smooth > np.quantile(ivar_smooth[ivar > 0], ivar_quantile)
-            ivar_set_mask *= ivar > 0
-            ivar_set_mask = np.logical_and(ivar_set_mask, additional_mask)
+            ivar = enmap.read_map(ivar_fn).astype(bool) 
+            while len(ivar.shape) > 2:
+                ivar = ivar[0].astype(np.float32)
 
+            # ivar_set_mask = ivar_smooth > np.quantile(ivar_smooth[ivar > 0], ivar_quantile)
+            # ivar_set_mask = np.logical_and(ivar_set_mask, additional_mask)
+
+            ivar_smooth = enmap.smooth_gauss(ivar, np.deg2rad(ivar_smooth_deg))
+            ivar_set_mask = enmap.zeros(ivar.shape, ivar.wcs, dtype=bool)
+            ivar_set_mask[ivar_smooth > np.quantile(ivar_smooth[ivar > 0], ivar_quantile)] = 1
             # check that inside of ivar_mask, there are no zero ivar
             assert np.all(ivar[ivar_set_mask] > 0), \
                 f'{sv}, {m}, set{i} has zero ivar inside ivar_mask'
@@ -76,11 +85,11 @@
                 map = enmap.read_map(map_fn)
                 p = enplot.plot(map, downgrade=8, ticks=1, colorbar=True, range=[500, 100, 100])
                 map_plot_fn = os.path.basename(map_fn)
-                enplot.write(os.path.join(maps_ivar_dir, map_plot_fn), p)
+                enplot.write(os.path.join(maps_ivar_dir, map_plot_fn)[:-5], p)
 
                 p = enplot.plot(ivar, downgrade=8, ticks=1, colorbar=True)
                 ivar_plot_fn = os.path.basename(ivar_fn)
-                enplot.write(os.path.join(maps_ivar_dir, ivar_plot_fn), p)
+                enplot.write(os.path.join(maps_ivar_dir, ivar_plot_fn)[:-5], p)
 
             log.info(f'{sv}, {m}, set{i} survey solid angle : {so_window.get_survey_solid_angle(so_map.from_enmap(ivar_set_mask)):.5f}')
 
@@ -91,8 +100,8 @@
             # save plot of mask
             if save_plot_mask:
                 p = enplot.plot(ivar_set_mask, downgrade=8, ticks=1, colorbar=True)
-                enplot.write(ivar_set_mask_fn, p)
-    
+                enplot.write(ivar_set_mask_fn[:-5], p)
+
             # also build xtra masks that are the union and intersection
             # of all the xtra masks
             mask_intersect = np.logical_and(mask_intersect, ivar_set_mask)
@@ -108,16 +117,16 @@
         # save plot of mask
         if save_plot_mask:
             p = enplot.plot(ivar_mask, downgrade=8, ticks=1, colorbar=True)
-            enplot.write(ivar_mask_fn, p)
+            enplot.write(ivar_mask_fn[:-5], p)
 
 log.info(f'All intersection survey solid angle : {so_window.get_survey_solid_angle(so_map.from_enmap(mask_intersect)):.5f}')
 log.info(f'All union survey solid angle : {so_window.get_survey_solid_angle(so_map.from_enmap(mask_union)):.5f}')
 
-# save union and intersect masks
-enmap.write_map(os.path.join(mask_dir, f'xtra_mask_intersect.fits'), mask_intersect.astype(np.float32))
-enmap.write_map(os.path.join(mask_dir, f'xtra_mask_union.fits'), mask_union.astype(np.float32))
+# plot and save union and intersect masks
+p = enplot.plot(mask_intersect, downgrade=8, ticks=1, colorbar=True)
+enplot.write(os.path.join(mask_dir, f'xtra_mask_intersect'), p)
+enmap.write_map(os.path.join(mask_dir, f'xtra_mask_intersect.fits'), mask_intersect.astype(np.float32))\
 
-# save plot of union and intersect masks
-if save_plot_mask:
-    p = enplot.plot([mask_intersect, mask_union], downgrade=8, colorbar=True, ticks=1)
-    enplot.write(os.path.join(mask_dir, 'xtra_mask_intersect_and_union'), p)
+p = enplot.plot(mask_union, downgrade=8, ticks=1, colorbar=True)
+enplot.write(os.path.join(mask_dir, f'xtra_mask_union'), p)
+enmap.write_map(os.path.join(mask_dir, f'xtra_mask_union.fits'), mask_union.astype(np.float32))