def lens(ulensed,convergence):
posmap = ulensed.posmap()
kmask = maps.mask_kspace(ulensed.shape,ulensed.wcs,lmin=10,lmax=8000)
phi,_ = lensing.kappa_to_phi(enmap.enmap(maps.filter_map(enmap.enmap(convergence,wcs),kmask),wcs),ulensed.modlmap(),return_fphi=True)
grad_phi = enmap.grad(phi)
pos = posmap + grad_phi
alpha_pix = ulensed.sky2pix(pos, safe=False)
lensed = enlensing.displace_map(ulensed, alpha_pix, order=5)
return lensed
mstats = stats.Stats(comm)
np.random.seed(rank)
# QE
tellmin = modlmap[modlmap>2].min(); tellmax = 8000; kellmin = tellmin ; kellmax = 8096
tmask = maps.mask_kspace(shape,wcs,lmin=tellmin,lmax=tellmax)
kmask = maps.mask_kspace(shape,wcs,lmin=kellmin,lmax=kellmax)
qest = lensing.qest(shape,wcs,theory,noise2d=kbeam*0.+(noise_uK_rad)**2.,beam2d=kbeam,kmask=tmask,kmask_K=kmask,pol=False,grad_cut=2000,unlensed_equals_lensed=False)
for i,task in enumerate(my_tasks):
if (i+1)%10==0 and rank==0: print(i+1)
unlensed = mg.get_map()
noise_map = ng.get_map()
lensed = maps.filter_map(enlensing.displace_map(unlensed, alpha_pix, order=lens_order),kbeam)
stamp = lensed + noise_map
if task==0: io.plot_img(stamp,pout_dir+"cmb_noisy.png")
recon = qest.kappa_from_map("TT",stamp)
cents, recon1d = binner.bin(recon)
mstats.add_to_stats("recon1d",recon1d)
mstats.add_to_stack("recon",recon)
mstats.get_stats()
mstats.get_stacks()
if rank==0:
stack = mstats.stacks['recon']
pos = thetaMap.posmap() + alpha
pix = thetaMap.sky2pix(pos, safe=False)
gradfit = 2.*np.pi/180./60.
from scipy.ndimage.interpolation import rotate
N = 100
avgrot = 0.
for i in range(N):
map = enmap.rand_map(shape, wcs, ps)/TCMB
gmap = enmap.grad(map)
angle = np.arctan2(gmap[0][np.where(thetaMap<gradfit)].mean(),gmap[1][np.where(thetaMap<gradfit)].mean())*180./np.pi
lensedTQU = lensing.displace_map(map, pix,order=5)+0.j
diffmap = (lensedTQU-map).real
rotmap = rotate(diffmap, angle,reshape=False)
avgrot += rotmap
print((i,angle))
#io.highResPlot2d(map,out_dir+"unl"+str(i)+".png")
#io.highResPlot2d(rotmap,out_dir+"rot"+str(i)+".png")
avgrot /= N
io.highResPlot2d(avgrot,out_dir+"avgrot.png")
# Stats
Nsims = args.Nclusters
Njobs = Nsims
num_each,each_tasks = mpi.mpi_distribute(Njobs,numcores)
if rank==0: print ("At most ", max(num_each) , " tasks...")
my_tasks = each_tasks[rank]
mstats = stats.Stats(comm)
np.random.seed(rank)
for i,task in enumerate(my_tasks):
if (i+1)%10==0 and rank==0: print(i+1)
# Sim
unlensed = mg.get_map()
noise_map = ng.get_map()
lensed = enlensing.displace_map(unlensed, alpha_pix, order=lens_order)
tot_beamed = maps.filter_map(lensed,kbeam) #+ fg_true
stamp = tot_beamed + noise_map
if task==0:
io.plot_img(unlensed,pout_dir + "0_unlensed.png")
io.plot_img(lensed,pout_dir + "1_lensed.png")
io.plot_img(fg,pout_dir + "2_fg.png")
io.plot_img(stamp,pout_dir + "3_tot.png")
# Bayesian
totlnlikes = []
for k,kamp in enumerate(bkamps):
lnlike = maps.get_lnlike(cinvs[k],stamp) + logdets[k]
totlnlike = lnlike #+ lnprior[k]
totlnlikes.append(totlnlike)
fkappa = lens_func(fmodrmap)
phi, _ = lensing.kappa_to_phi(fkappa, fmodlmap, return_fphi=True)
grad_phi = enmap.grad(phi)
pos = enmap.posmap(fshape, fwcs) + grad_phi
alpha_pix = enmap.sky2pix(fshape, fwcs, pos, safe=False)
lmax = cmodlmap.max()
ells = np.arange(0, lmax, 1)
cls = theory.uCl('TT', ells)
ps = cls.reshape((1, 1, ells.size))
mg = maps.MapGen(cshape, cwcs, ps)
unlensed = mg.get_map()
hunlensed = enmap.enmap(resample.resample_fft(unlensed.copy(), fshape), fwcs)
hlensed = enlensing.displace_map(hunlensed,
alpha_pix,
order=lens_order,
mode=mode)
lensed = enmap.enmap(resample.resample_fft(hlensed, cshape), cwcs)
rkappa = lens_func(rmodrmap)
phi, _ = lensing.kappa_to_phi(rkappa, rmodlmap, return_fphi=True)
grad_phi = enmap.grad(phi)
pos = enmap.posmap(rshape, rwcs) + grad_phi
ralpha_pix = enmap.sky2pix(rshape, rwcs, pos, safe=False)
hunlensed = enmap.enmap(resample.resample_fft(unlensed.copy(), rshape), rwcs)
enmap.write_map("unlensed_0.001arc.fits", hunlensed)
hlensed = enlensing.displace_map(hunlensed,
ralpha_pix,
order=lens_order,
phi, _ = lensing.kappa_to_phi(kappa, gmodlmap, return_fphi=True)
grad_phi = enmap.grad(phi)
posmap = enmap.posmap(gshape, gwcs)
pos = posmap + grad_phi
alpha_pix = enmap.sky2pix(gshape, gwcs, pos, safe=False)
kbeam = maps.gauss_beam(args.beam, gmodlmap)
mstats = stats.Stats()
for i in range(args.Nclusters):
if (i + 1) % 100 == 0: print(i + 1)
unlensed = mg.get_map()
noise_map = ng.get_map()
lensed = maps.filter_map(
enlensing.displace_map(unlensed.copy(), alpha_pix, order=lens_order),
kbeam)
fdownsampled = enmap.enmap(resample.resample_fft(lensed, bshape), bwcs)
stamp = fdownsampled + noise_map
#cutout = lensed + noise_map
cutout = stamp[int(bshape[0] / 2. - shape[0] / 2.):int(bshape[0] / 2. +
shape[0] / 2.),
int(bshape[0] / 2. - shape[0] / 2.):int(bshape[0] / 2. +
shape[0] / 2.)]
# print(cinvs[k].shape,cutout.shape)
totlnlikes = []
for k, kamp in enumerate(kamps):
lnlike = maps.get_lnlike(cinvs[k], cutout) + logdets[k]
def get_lensed(self, unlensed, order=3, mode="spline", border="cyclic"):
return lensing.displace_map(unlensed,
self._displace_pix,
order=order,
mode=mode,
border=border)
