def rand_alm(ps, ainfo=None, lmax=None, seed=None, dtype=np.complex128, m_major=True, return_ainfo=False):
"""This is a replacement for healpy.synalm. It generates the random
numbers in l-major order before transposing to m-major order in order
to allow generation of low-res and high-res maps that agree on large
scales. It uses 2/3 of the memory of healpy.synalm, and has comparable
speed."""
rtype = np.zeros([0],dtype=dtype).real.dtype
ps = np.asarray(ps)
if ainfo is None: ainfo = sharp.alm_info(min(lmax,ps.shape[-1]-1) or ps.shape[-1]-1)
if ps.ndim == 1:
wps = ps[None,None]
elif ps.ndim == 2:
wps = powspec.sym_expand(ps, scheme="diag")
elif ps.ndim == 3:
wps = ps
else:
raise ValuerError("power spectrum must be [nl], [nspec,nl] or [ncomp,ncomp,nl]")
ncomp = wps.shape[0]
ps12 = enmap.multi_pow(wps, 0.5)
# Draw random gaussian numbers in chunks to save memory
alm = np.empty([ncomp,ainfo.nelem],dtype=dtype)
aflat = alm.reshape(-1).view(rtype)
bsize = 0x10000
if seed != None: np.random.seed(seed)
for i in range(0, aflat.size, bsize):
aflat[i:i+bsize] = np.random.standard_normal(min(bsize,aflat.size-i))
# Transpose numbers to make them m-major.
if m_major: ainfo.transpose_alm(alm,alm)
# Scale alms by spectrum, taking into account which alms are complex
ainfo.lmul(alm, (ps12/2**0.5).astype(rtype), alm)
alm[:,:ainfo.lmax].imag = 0
alm[:,:ainfo.lmax].real *= 2**0.5
if ps.ndim == 1: alm = alm[0]
if return_ainfo: return alm, ainfo
else: return alm
def __init__(self, maps, inoise, ps, T=None):
self.d = maps
self.iN = inoise
self.hN = en.multi_pow(inoise, 0.5, axes=[1, 2])
self.iS = en.spec2flat(maps.shape[-3:], maps.wcs, ps, -1.0)
self.hS = en.spec2flat(maps.shape[-3:], maps.wcs, ps, -0.5)
self.ps = ps
self.b, self.x = None, None
# Prepare the preconditioner. It approximates the noise as the
# same in every pixel, and ignores the cmb-template coupling.
# See M(self,u) for details.
iN_white = np.array(np.sum(np.mean(np.mean(self.iN, -1), -1), 0))
# iN_white is now in pixel space, but the preconditioner needs it
# in harmonic space, which introduces a
#norm = np.prod((maps.box[1]-maps.box[0])/maps.shape[-2:])
#norm = 1./np.prod(maps.shape[-2:])
#iN_white /= norm
self.S_prec = en.multi_pow(self.iS + iN_white[:, :, None, None], -1)
# The template
self.set_template(T)
def solve(w, m):
if w.ndim == 2: return m / w[None]
elif w.ndim == 3: return m / w
elif w.ndim == 4: return enmap.map_mul(enmap.multi_pow(w, -1), m)
else: raise NotImplementedError("Only 2d, 3d or 4d weight maps understood")
p2d[modlmap<90]=0.
p2d[modlmap>4000]=0.
io.quickPlot2d(np.fft.fftshift(np.log10(p2d)),io.dout_dir+"p2d.png")
mg = enmap.MapGen(shape,wcs,p2d.reshape(1,1,modlmap.shape[0],modlmap.shape[1]))
imap = mg.get_map()
io.quickPlot2d(imap,io.dout_dir+"cmb.png")
# imap2 = mg.get_map()
# io.quickPlot2d(imap2,io.dout_dir+"cmb2.png")
dtype = np.complex128
rtype = np.zeros([0],dtype=dtype).real.dtype
print("alms...")
alm = curvedsky.map2alm(imap,lmax=4000)
powfac = 0.5
ps_data = enmap.multi_pow(alm*alm.conj()[None,None],powfac).astype(rtype)
print(alm)
print((alm.shape))
pspec = np.ones((8000))
pspec[:200] = 0.
pspec[4000:] = 0.
rand_alm,ainfo = curvedsky.rand_alm(ps=pspec,lmax=4000,return_ainfo=True)
#rand_map = curvedsky.rand_map(shape,wcs,pspec)
#rand_alm = curvedsky.map2alm(rand_map,lmax=4000)
ps_alm = enmap.multi_pow(rand_alm*rand_alm.conj()[None,None],powfac).astype(rtype)
new_sim_alm = rand_alm *np.nan_to_num(ps_data/ps_alm)
new_sim = curvedsky.alm2map(new_sim_alm,imap)#,ainfo=ainfo)
#new_sim = curvedsky.alm2map(alm,imap)#,ainfo=ainfo)