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
# Noise model
noise_uK_rad = args.noise*np.pi/180./60.
normfact = np.sqrt(np.prod(enmap.pixsize(shape,wcs)))
kbeam = maps.gauss_beam(args.beam,modlmap)
# Simulate
lmax = int(modlmap.max()+1)
ells = np.arange(0,lmax,1)
ps = theory.uCl('TT',ells).reshape((1,1,lmax))
~ps_noise = np.array([(noise_uK_rad)**2.]*ells.size).reshape((1,1,ells.size))
mg = maps.MapGen(shape,wcs,ps)
ng = maps.MapGen(shape,wcs,ps_noise)
kamp_true = args.Amp
kappa = lensing.nfw_kappa(kamp_true*1e15,modrmap,cc,overdensity=200.,critical=True,atClusterZ=True)
phi,_ = lensing.kappa_to_phi(kappa,modlmap,return_fphi=True)
grad_phi = enmap.grad(phi)
posmap = enmap.posmap(shape,wcs)
pos = posmap + grad_phi
alpha_pix = enmap.sky2pix(shape,wcs,pos, safe=False)
lens_order = 5
if rank==0: print("Starting sims...")
# 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)
#lens_func = lambda x: 0.2 * np.exp(-x**2./sigma**2./2.)
rshape, rwcs = maps.rect_geometry(width_arcmin=5., px_res_arcmin=0.001)
fshape, fwcs = maps.rect_geometry(width_arcmin=20., px_res_arcmin=0.1)
cshape, cwcs = maps.rect_geometry(width_arcmin=20., px_res_arcmin=0.5)
rmodrmap = enmap.modrmap(rshape, rwcs)
fmodrmap = enmap.modrmap(fshape, fwcs)
cmodrmap = enmap.modrmap(cshape, cwcs)
rmodlmap = enmap.modlmap(rshape, rwcs)
fmodlmap = enmap.modlmap(fshape, fwcs)
cmodlmap = enmap.modlmap(cshape, cwcs)
print(fshape, cshape)
mass = 2.e14
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,
amin = 0.18
amax = 0.22
num_amps = 10
kamps = np.linspace(amin, amax, num_amps)
cinvs = []
logdets = []
for k, kamp in enumerate(kamps):
kappa_template = lensing.nfw_kappa(kamp * 1e15,
bmodrmap,
cc,
overdensity=200.,
critical=True,
atClusterZ=True)
phi, _ = lensing.kappa_to_phi(kappa_template, bmodlmap, return_fphi=True)
grad_phi = enmap.grad(phi)
pos = posmap + grad_phi
alpha_pix = enmap.sky2pix(bshape, bwcs, pos, safe=False)
#if k==0: io.plot_img(kappa_template)
with bench.show("lensing cov"):
Scov = lensing.lens_cov(Ucov,
alpha_pix,
lens_order=lens_order,
kbeam=kbeam,
bshape=shape)
Tcov = Scov + Ncov + 5000 # !!!
with bench.show("covwork"):