change the implementation of the filter to make it more modular, also include option for sigurd filtering method

Author: Thibaut Louis

pspy/so_map_preprocessing.py

@@ -1,56 +1,125 @@
 import numpy as np, pylab as plt
 from pspy import flat_tools, pspy_utils
+from pixell import enmap, utils
 
-def kspace_filter(map, vk_mask=None, hk_mask=None, window=None, normalize=False):
-    filtered_map = map.copy()
-
-    ft = flat_tools.fft_from_so_map(map, normalize=normalize)
-    if window is not None:
-        ft = flat_tools.get_ffts(map, window, normalize=normalize)
-
+def build_std_filter(shape, wcs, vk_mask, hk_mask, dtype=np.float64):
+    ly, lx  = enmap.laxes(shape, wcs, method="auto")
+    ly = ly.astype(dtype)
+    lx = lx.astype(dtype)
+    filter = enmap.ones(shape[-2:], wcs, dtype)
     if vk_mask is not None:
-        id_vk = np.where((ft.lx > vk_mask[0]) & (ft.lx < vk_mask[1]))
+        id_vk = np.where((lx > vk_mask[0]) & (lx < vk_mask[1]))
     if hk_mask is not None:
-        id_hk = np.where((ft.ly > hk_mask[0]) & (ft.ly < hk_mask[1]))
-
-    if map.ncomp == 1:
-        if vk_mask is not None:
-            ft.kmap[: , id_vk] = 0.
-        if hk_mask is not None:
-            ft.kmap[id_hk , :] = 0.
-        
-    if map.ncomp == 3:
-        for i in range(3):
-            if vk_mask is not None:
-                ft.kmap[i, : , id_vk] = 0.
-            if hk_mask is not None:
-                ft.kmap[i, id_hk , :] = 0.
-
-    filtered_map.data[:] = ft.map_from_fft(normalize=normalize)
+        id_hk = np.where((ly > hk_mask[0]) & (ly < hk_mask[1]))
+    filter[:, id_vk] = 0
+    filter[id_hk, :] = 0
+    return filter
+
+def apply_std_filter(imap, filter):
+    filtered_map = imap.copy()
+    if imap.ncomp == 1:
+        filtered_map.data = enmap.ifft(filter*enmap.fft(imap.data)).real
+    else:
+        for i in range(imap.ncomp):
+            filtered_map.data[i] = enmap.ifft(filter*enmap.fft(imap.data[i])).real
+
     return filtered_map
 
+def build_sigurd_filter(shape, wcs, lbounds, dtype=np.float64):
+    ly, lx  = enmap.laxes(shape, wcs, method="auto")
+    ly = ly.astype(dtype)
+    lx = lx.astype(dtype)
+    lbounds = np.asarray(lbounds)
+    if lbounds.ndim < 2:
+        lbounds = np.broadcast_to(lbounds, (1,2))
+    if lbounds.ndim > 2 or lbounds.shape[-1] != 2:
+        raise ValueError("lbounds must be [:,{ly,lx}]")
+    filter = enmap.ones(shape[-2:], wcs, dtype)
+    # Apply the filters
+    for i , (ycut, xcut) in enumerate(lbounds):
+        filter *= 1-(np.exp(-0.5*(ly/ycut)**2)[:,None]*np.exp(-0.5*(lx/xcut)**2)[None,:])
+    return filter
 
-def analytical_tf(map, binning_file, lmax, vk_mask=None, hk_mask=None):
-    import time
-    t = time.time()
-    ft = flat_tools.fft_from_so_map(map)
-    ft.create_kspace_mask(vertical_stripe=vk_mask, horizontal_stripe=hk_mask)
-    lmap = ft.lmap
-    kmask = ft.kmask
-    twod_index = lmap.copy()
+def apply_sigurd_filter(imap, ivar, filter, tol=1e-4, ref=0.9):
+    """Filter enmap imap with the given 2d fourier filter while
+    weithing spatially with ivar"""
+    
+    filtered_map = imap.copy()
+    rhs    = enmap.ifft((1-filter)*enmap.fft(imap.data*ivar)).real
+    div    = enmap.ifft((1-filter)*enmap.fft(ivar)).real
+    # Avoid division by very low values
+    div    = np.maximum(div, np.percentile(ivar[::10,::10],ref*100)*tol)
+    filtered_map.data = imap.data - rhs/div
+    return filtered_map
 
+def analytical_tf(map, filter, binning_file, lmax):
+    lmap  = enmap.modlmap(map.data.shape, map.data.wcs, method="auto")
+    twod_index = lmap.copy()
     bin_lo, bin_hi, bin_c, bin_size = pspy_utils.read_binning_file(binning_file, lmax)
     nbins = len(bin_lo)
-    print(lmap.shape)
     tf = np.zeros(nbins)
-    print("init tf", time.time()-t)
     for ii in range(nbins):
-        id = np.where( (lmap >= bin_lo[ii]) & (lmap <= bin_hi[ii]))
+        id = np.where((lmap >= bin_lo[ii]) & (lmap <= bin_hi[ii]))
         twod_index *= 0
         twod_index[id] = 1
         bin_area = np.sum(twod_index)
-        masked_bin_area= np.sum(twod_index * kmask)
+        masked_bin_area= np.sum(twod_index * filter)
         tf[ii] = masked_bin_area / bin_area
-    print("loop", time.time()-t)
 
     return bin_c, tf
+
+
+
+#def analytical_tf_old(map, binning_file, lmax, vk_mask=None, hk_mask=None):
+#    import time
+#    t = time.time()
+#    ft = flat_tools.fft_from_so_map(map)
+#    ft.create_kspace_mask(vertical_stripe=vk_mask, horizontal_stripe=hk_mask)
+#    lmap = ft.lmap
+#    kmask = ft.kmask
+#    twod_index = lmap.copy()
+
+#    bin_lo, bin_hi, bin_c, bin_size = pspy_utils.read_binning_file(binning_file, lmax)
+#    nbins = len(bin_lo)
+#    print(lmap.shape)
+#    tf = np.zeros(nbins)
+#    print("init tf", time.time()-t)
+#    for ii in range(nbins):
+#        id = np.where( (lmap >= bin_lo[ii]) & (lmap <= bin_hi[ii]))
+#        twod_index *= 0
+#        twod_index[id] = 1
+#        bin_area = np.sum(twod_index)
+#        masked_bin_area= np.sum(twod_index * kmask)
+#        tf[ii] = masked_bin_area / bin_area
+#    print("loop", time.time()-t)
+
+#    return bin_c, tf
+
+
+#def kspace_filter_old(map, vk_mask=None, hk_mask=None, window=None, normalize=False):
+#    filtered_map = map.copy()
+
+ #   ft = flat_tools.fft_from_so_map(map, normalize=normalize)
+ #   if window is not None:
+ #       ft = flat_tools.get_ffts(map, window, normalize=normalize)
+
+ #   if vk_mask is not None:
+ #       id_vk = np.where((ft.lx > vk_mask[0]) & (ft.lx < vk_mask[1]))
+ #   if hk_mask is not None:
+ #       id_hk = np.where((ft.ly > hk_mask[0]) & (ft.ly < hk_mask[1]))
+
+ #   if map.ncomp == 1:
+ #       if vk_mask is not None:
+ #           ft.kmap[: , id_vk] = 0.
+ #       if hk_mask is not None:
+ #           ft.kmap[id_hk , :] = 0.
+    
+ #   if map.ncomp == 3:
+ #      for i in range(3):
+ #           if vk_mask is not None:
+ #               ft.kmap[i, : , id_vk] = 0.
+ #           if hk_mask is not None:
+ #               ft.kmap[i, id_hk , :] = 0.
+
+ #   filtered_map.data[:] = ft.map_from_fft(normalize=normalize)
+ #   return filtered_map

tutorials/tuto_kspace_filter.py

@@ -18,10 +18,8 @@
 apo_radius_degree = 2
 niter = 0
 lmax = 2000
-nsims = 4
+nsims = 30
 
-vk_mask = [-90, 90]
-hk_mask = [-50, 50]
 
 test_dir = "result_kspace"
 try:
@@ -45,6 +43,21 @@
 window = (window,window)
 mbb_inv, Bbl = so_mcm.mcm_and_bbl_spin0and2(window, binning_file, lmax=lmax, type="Dl", niter=niter)
 
+# we build two different filters and compute the analytical 1d binned spectra they correspond to
+
+vk_mask = [-90, 90]
+hk_mask = [-50, 50]
+filter_std = so_map_preprocessing.build_std_filter(template.data.shape, template.data.wcs, vk_mask, hk_mask, dtype=np.float64)
+
+lbounds = [4000, 5]
+filter_sig = so_map_preprocessing.build_sigurd_filter(template.data.shape, template.data.wcs, lbounds, dtype=np.float64)
+
+
+lb, analytic_tf_std = so_map_preprocessing.analytical_tf(template, filter_std, binning_file, lmax)
+lb, analytic_tf_sig = so_map_preprocessing.analytical_tf(template, filter_sig, binning_file, lmax)
+
+
+
 
 # Let's run nsims, and apply the transfer function, we will compute both the SPHT and the FFT of the initial and filtered maps.
 # Then we compute the associated 1d and 2d power spectra, we will store the value of the resulting
@@ -56,23 +69,27 @@
 spectra_2d = ["II", "IQ", "IU", "QI", "QQ", "QU", "UI", "UQ", "UU"]
 Db_list = {}
 p2d_list = {}
-for spec in spectra:
-    Db_list["standard", spec] = []
-    Db_list["filtered", spec] = []
-for spec in spectra_2d:
-    p2d_list["standard", spec] = []
-    p2d_list["filtered", spec] = []
-    
+
+runs = ["sig_filtered", "std_filtered", "no_filter"]
+
+for run in runs:
+    for (spec, spec2d) in zip(spectra, spectra_2d):
+        Db_list[run, spec] = []
+        p2d_list[run, spec2d] = []
+
 
 for iii in range(nsims):
     print("sim number %03d" % iii)
-    for run in ["standard", "filtered"]:
+    for run in runs:
         cmb = template.synfast(clfile)
 
         if iii == 0: cmb.plot(file_name="%s/cmb"%(test_dir))
-        if run == "filtered":
-            cmb = so_map_preprocessing.kspace_filter(cmb, vk_mask=vk_mask, hk_mask=hk_mask, normalize="phys")
+        if run == "std_filtered":
+            cmb = so_map_preprocessing.apply_std_filter(cmb, filter_std)
             if iii == 0: cmb.plot(file_name="%s/cmb_filter"%(test_dir))
+        elif run == "sig_filtered":
+            cmb = so_map_preprocessing.apply_sigurd_filter(cmb, binary.data, filter_sig, tol=1e-4, ref=0.9)
+            if iii == 0: cmb.plot(file_name="%s/cmb_sigurd_filter"%(test_dir))
 
         alm_cmb = sph_tools.get_alms(cmb, window, niter, lmax)
         fft_cmb = flat_tools.get_ffts(cmb, window, lmax)
@@ -97,48 +114,54 @@
 
 
 # First the 1d part with the associated transfer functions
-# we also compute a simple analytical version of the transfer function
 
-lb, analytic_tf = so_map_preprocessing.analytical_tf(cmb, binning_file, lmax, vk_mask=vk_mask, hk_mask=hk_mask)
 
 mean = {}
 std = {}
 for spec in spectra:
-
-    mean["standard", spec]= np.mean(Db_list["standard", spec], axis=0)
-    mean["filtered", spec]= np.mean(Db_list["filtered", spec], axis=0)
-    std["standard", spec]= np.std(Db_list["standard", spec], axis=0)
-    std["filtered", spec]= np.std(Db_list["filtered", spec], axis=0)
-
-    plt.errorbar(lb, mean["standard", spec], std["standard", spec], fmt=".", label="standard")
-    plt.errorbar(lb, mean["filtered", spec], std["filtered", spec], fmt=".", label="filtered")
+    for run in runs:
+        mean[run, spec]= np.mean(Db_list[run, spec], axis=0)
+        std[run, spec]= np.std(Db_list[run, spec], axis=0)
+        plt.errorbar(lb, mean[run, spec], std[run, spec], fmt=".", label=run)
+        
     plt.legend()
     plt.savefig("%s/spectra_%s.png" % (test_dir, spec))
     plt.clf()
     plt.close()
 
 
     if spec in ["TT", "EE", "TE", "BB"] :
-        plt.plot(lb, analytic_tf)
-        plt.plot(lb, mean["filtered", spec]/mean["standard", spec], ".")
+        plt.plot(lb, analytic_tf_std, label = "std analytic")
+        plt.plot(lb, analytic_tf_sig, label = "sigurd analytic")
+
+        plt.errorbar(lb, mean["std_filtered", spec]/mean["no_filter", spec], label="std filter", fmt=".")
+        plt.errorbar(lb, mean["sig_filtered", spec]/mean["no_filter", spec], label="sigurd filter", fmt=".")
+
+        plt.legend()
         plt.savefig("%s/tf_%s.png" % (test_dir, spec))
         plt.clf()
         plt.close()
 
 
-plt.plot(lb, analytic_tf)
+plt.plot(lb, analytic_tf_std, label = "std analytic")
+plt.plot(lb, analytic_tf_sig, label = "sigurd analytic")
+
 for spec in ["TT", "EE"] :
-    plt.plot(lb, mean["filtered", spec]/mean["standard", spec], ".", label = spec)
+    plt.plot(lb, mean["std_filtered", spec]/mean["no_filter", spec], ".", label = spec + "std")
+    plt.plot(lb, mean["sig_filtered", spec]/mean["no_filter", spec], ".", label = spec + "sigurd")
+
+plt.legend()
 plt.savefig("%s/tf_TT_and_EE.png" % (test_dir))
 plt.clf()
 plt.close()
 
 # Now the 2d part with the plot of the effect of the filter
 
-mean["standard"] = p2d_dict.copy()
-mean["filtered"] = p2d_dict.copy()
+
+for run in runs:
+    mean[run] = p2d_dict.copy()
 for spec in spectra_2d:
-    for run in ["standard", "filtered"]:
+    for run in runs:
         mean[run].powermap[spec] = np.mean(p2d_list[run, spec], axis=0)
-for run in ["standard", "filtered"]:
+for run in runs:
     mean[run].plot(power_of_ell=2, png_file="%s/%s" % (test_dir, run))