def get_lensed_cmb(self, seed=None, kappa=None, save_output=True, verbose=True, overwrite=False, dtype=np.float64):
try:
assert(seed is not None)
assert(not overwrite)
if verbose:
print(f"trying to load saved lensed cmb. sim idx: {seed}")
lmaps = enmap.empty((3,)+self.shape, self.wcs, dtype=dtype)
for i, polidx in enumerate(['T','Q','U']):
fname = self.get_output_file_name("lensed_cmb", seed, polidx=polidx)
lmaps[i] = enmap.read_map(fname)
except:
ualm = self._generate_unlensed_cmb_alm(seed=seed)
if kappa is None:
kappa = self._get_kappa(seed, dtype=np.float64)
lmax = 10000
l = np.arange(lmax+1)
l_fact = 1/((l*(l+1))/2)
l_fact[0] = 0
kalm = curvedsky.map2alm(kappa, lmax=lmax)
kalm = curvedsky.map2alm(kappa, lmax=lmax); del kappa
kalm = hp.almxfl(kalm, l_fact)
tshape, twcs = enmap.fullsky_geometry(res=1*utils.arcmin)
print("start lensing")
lmaps = lensing.lens_map_curved((3,)+tshape, twcs, kalm, ualm)[0]
lalms = curvedsky.map2alm(lmaps, lmax=lmax, spin=0)
lmaps = curvedsky.alm2map(lalms, enmap.zeros((3,)+self.shape, self.wcs), spin=0)
print("finish lensing")
if save_output:
for i, polidx in enumerate(['T','Q','U']):
fname = self.get_output_file_name("lensed_cmb", seed, polidx=polidx)
os.makedirs(os.path.dirname(fname), exist_ok=True)
enmap.write_map(fname, lmaps[i].astype(np.float32))
return lmaps.astype(dtype)
def test_b_sign(self):
"""
We generate a random IQU map with geometry such that cdelt[0]<0
We transform this to TEB with map2harm and map2alm followed by
scalar harm2map and alm2map and use these as reference T,E,B maps.
We flip the original map along the RA direction.
We transform this to TEB with map2harm and map2alm followed by
scalar harm2map and alm2map and use these as comparison T,E,B maps.
We compare these maps.
"""
ells,cltt,clee,clbb,clte = np.loadtxt(DATA_PREFIX+"cosmo2017_10K_acc3_lensedCls.dat",unpack=True)
ps_cmb = np.zeros((3,3,ells.size))
ps_cmb[0,0] = cltt
ps_cmb[1,1] = clee
ps_cmb[2,2] = clbb
ps_cmb[1,0] = clte
ps_cmb[0,1] = clte
np.random.seed(100)
# Curved-sky is fine
lmax = 1000
alm = curvedsky.rand_alm_healpy(ps_cmb,lmax=lmax)
shape,iwcs = enmap.fullsky_geometry(res=np.deg2rad(10./60.))
wcs = enmap.empty(shape,iwcs)[...,::-1].wcs
shape = (3,) + shape
imap = curvedsky.alm2map(alm,enmap.empty(shape,wcs))
oalm = curvedsky.map2alm(imap.copy(),lmax=lmax)
rmap = curvedsky.alm2map(oalm,enmap.empty(shape,wcs),spin=0)
imap2 = imap.copy()[...,::-1]
oalm = curvedsky.map2alm(imap2.copy(),lmax=lmax)
rmap2 = curvedsky.alm2map(oalm,enmap.empty(shape,wcs),spin=0)
assert np.all(np.isclose(rmap[0],rmap2[0]))
assert np.all(np.isclose(rmap[1],rmap2[1]))
assert np.all(np.isclose(rmap[2],rmap2[2]))
# Flat-sky
px = 2.0
N = 300
shape,iwcs = enmap.geometry(pos=(0,0),res=np.deg2rad(px/60.),shape=(300,300))
shape = (3,) + shape
a = enmap.zeros(shape,iwcs)
a = a[...,::-1]
wcs = a.wcs
seed = 100
imap = enmap.rand_map(shape,wcs,ps_cmb,seed=seed)
kmap = enmap.map2harm(imap.copy())
rmap = enmap.harm2map(kmap,spin=0) # reference map
imap = imap[...,::-1]
kmap = enmap.map2harm(imap.copy())
rmap2 = enmap.harm2map(kmap,spin=0)[...,::-1] # comparison map
assert np.all(np.isclose(rmap[0],rmap2[0]))
assert np.all(np.isclose(rmap[1],rmap2[1],atol=1e0))
assert np.all(np.isclose(rmap[2],rmap2[2],atol=1e0))
def get_power(map_list, ivar_list, a, b, mask, N=20):
"""
Calculate the average coadded flattened power spectrum P_{ab} used to generate simulation for the splits.
Inputs:
map_list: list of source free splits
ivar_list: list of the inverse variance maps splits
a: 0,1,2 for I,Q,U respectively
b:0,1,2 for I,Q,U, respectively
N: window to smooth the power spectrum by in the rolling average.
mask: apodizing mask
Output:
1D power spectrum accounted for w2 from 0 to 10000
"""
pmap = enmap.pixsizemap(map_list[0].shape, map_list[0].wcs)
cl_ab = []
n = len(map_list)
#calculate the coadd maps
if a != b:
coadd_a = coadd_mapnew(map_list, ivar_list, a)
coadd_b = coadd_mapnew(map_list, ivar_list, b)
else:
coadd_a = coadd_mapnew(map_list, ivar_list, a)
for i in range(n):
print(i)
if a != b:
d_a = map_list[i][a] - coadd_a
noise_a = d_a * np.sqrt(ivar_eff(i, ivar_list) / pmap) * mask
alm_a = cs.map2alm(noise_a, lmax=10000)
d_b = map_list[i][b] - coadd_b
noise_b = d_b * np.sqrt(ivar_eff(i, ivar_list) / pmap) * mask
alm_b = cs.map2alm(noise_b, lmax=10000)
cls = hp.alm2cl(alm_a, alm_b)
cl_ab.append(cls)
else:
d_a = map_list[i][a] - coadd_a
noise_a = d_a * np.sqrt(ivar_eff(i, ivar_list) / pmap) * mask
print("generating alms")
alm_a = cs.map2alm(noise_a, lmax=10000)
cls = hp.alm2cl(alm_a)
cl_ab.append(cls)
cl_ab = np.array(cl_ab)
sqrt_ivar = np.sqrt(ivar_eff(0, ivar_list) / pmap)
mask_ivar = sqrt_ivar * 0 + 1
mask_ivar[sqrt_ivar <= 0] = 0
mask = mask * mask_ivar
mask[mask <= 0] = 0
w2 = np.sum((mask**2) * pmap) / np.pi / 4.
power = 1 / n / (n - 1) * np.sum(cl_ab, axis=0)
ls = np.arange(len(power))
power[~np.isfinite(power)] = 0
power = rolling_average(power, N)
bins = np.arange(len(power))
power = maps.interp(bins, power)(ls)
return power / w2
def get_spectra(emap1, emap2=None, lmax=5000):
atol = np.min(np.array(emap1.pixshape()/eutils.arcmin))
alm1 = curvedsky.map2alm(emap1, lmax=lmax, atol=atol).astype(np.complex128)
alm2 = alm1 if emap2 is None else curvedsky.map2alm(emap2, lmax=lmax, atol=atol).astype(np.complex128)
cl = hp.alm2cl(alm1, alm2, lmax=lmax)
l = np.arange(len(cl))
return (l, cl)
def phi2kappa(phiMap, LMAX):
phiAlm = curvedsky.map2alm(phiMap, lmax=LMAX)
ells = np.arange(LMAX - 1)
kappaAlm = healpy.sphtfunc.almxfl(phiAlm, ells * (ells + 1) / 2.)
kappaMap = curvedsky.alm2map(kappaAlm, phiMap.copy())
return kappaMap
def get_spectra(tube, omap, nside, res, white, w2):
zeros = []
if nside is not None:
mlmax = nside * 2
alms = [
hp.map2alm(omap[i, 0, ...], lmax=mlmax, pol=True) / np.sqrt(w2)
for i in range(2)
]
else:
mlmax = 4096 * 2 / res
alms = [
cs.map2alm(omap[i, 0, ...], lmax=mlmax, spin=[0, 2]) / np.sqrt(w2)
for i in range(2)
]
for i in range(2):
for j in range(i, 2):
for p in range(3):
for q in range(p + 1, 3):
zeros.append(hp.alm2cl(alms[i][p], alms[j][q]))
for i in range(1, 3):
zeros.append(hp.alm2cl(alms[0][i], alms[1][i]))
c11 = [hp.alm2cl(alms[0][i], alms[0][i]) for i in range(3)]
c22 = [hp.alm2cl(alms[1][i], alms[1][i]) for i in range(3)]
cross = hp.alm2cl(alms[0][0], alms[1][0])
if white or tube[0] == "S":
zeros.append(cross)
c12 = []
else:
c12 = [cross]
ls = np.arange(c11[0].size)
return ls, c11, c22, c12, zeros
def map2alm_spin(self,imap,lmax,spin_alm,spin_transform):
dmap = -irot2d(np.stack((imap,imap.conj())),spin=spin_alm).real
if self.hpix:
res = hp.map2alm_spin(dmap,lmax=lmax,spin=spin_transform)
return res
else:
return cs.map2alm(enmap.enmap(dmap,imap.wcs),spin=spin_transform,lmax=lmax)
def rand_map(shape, wcs, ps_lensinput, doRotation=False, ps_rot=None, lmax=None,
maplmax=None, dtype=np.float64, seed=None, phi_seed=None,
oversample=2.0, spin=[0,2], output="l", geodesic=True, verbose=False,
delta_theta=None):
from pixell import curvedsky, sharp
ctype = np.result_type(dtype,0j)
# Restrict to target number of components
oshape = shape[-3:]
if len(oshape) == 2: shape = (1,)+tuple(shape)
ncomp = shape[-3]
ps_lensinput = ps_lensinput[:1+ncomp,:1+ncomp]
# First draw a random lensing field, and use it to compute the undeflected positions
if verbose: print("Generating alms")
if phi_seed is None:
alm, ainfo = curvedsky.rand_alm(ps_lensinput, lmax=lmax, seed=seed, dtype=ctype, return_ainfo=True)
else:
# We want separate seeds for cmb and phi. This means we have to do things a bit more manually
wps, ainfo = curvedsky.prepare_ps(ps_lensinput, lmax=lmax)
alm = np.empty([1+ncomp,ainfo.nelem],ctype)
curvedsky.rand_alm_white(ainfo, alm=alm[:1], seed=phi_seed)
curvedsky.rand_alm_white(ainfo, alm=alm[1:], seed=seed)
ps12 = enmap.multi_pow(wps, 0.5)
ainfo.lmul(alm, (ps12/2**0.5).astype(dtype), alm)
alm[:,:ainfo.lmax].imag = 0
alm[:,:ainfo.lmax].real *= 2**0.5
del wps, ps12
phi_alm, cmb_alm = alm[0], alm[1:]
if doRotation:
# generate a gaussian random field of rotational field
alpha_shape = (1, shape[-2], shape[-1])
alpha_map = curvedsky.rand_map(alpha_shape, wcs, ps_rot, lmax=lmax)
# convert alm to enmap object
cmb_map = enmap.empty((3,shape[-2],shape[-1]), wcs)
curvedsky.alm2map(cmb_alm, cmb_map)
# rotate tqu map by a rotational field
# Q_map = cmb_map[1]
# U_map = cmb_map[2]
cmb_map_rot = enmap.rotate_pol(cmb_map, alpha_map)
# convert tqu map back to the alm
cmb_alm = curvedsky.map2alm(cmb_map_rot, lmax=lmax)
del alpha_map, cmb_map, cmb_map_rot
del alm
# Truncate alm if we want a smoother map. In taylens, it was necessary to truncate
# to a lower lmax for the map than for phi, to avoid aliasing. The appropriate lmax
# for the cmb was the one that fits the resolution. FIXME: Can't slice alm this way.
#if maplmax: cmb_alm = cmb_alm[:,:maplmax]
return lens_map_curved(shape=shape, wcs=wcs, phi_alm=phi_alm,
cmb_alm=cmb_alm, phi_ainfo=ainfo, maplmax=maplmax,
dtype=dtype, oversample=oversample, spin=spin,
output=output, geodesic=geodesic, verbose=verbose,
delta_theta=delta_theta)
def tqu2teb(tmap, qmap, umap, lmax):
atol = np.min(np.array(tmap.pixshape() / eutils.arcmin))
tqu = np.zeros((3, ) + tmap.shape)
tqu[0], tqu[1], tqu[2] = (tmap, qmap, umap)
tqu = enmap.enmap(tqu, tmap.wcs)
alm = curvedsky.map2alm(tqu, lmax=lmax, atol=atol)
teb = curvedsky.alm2map(alm[:, None], tqu.copy()[:, None], spin=0)[:, 0]
del tqu
return (teb[0], teb[1], teb[2]) # tmap, emap, bmap
def tqu2teb(tqu, LMAX, wantCl=False, wantAlmAndCl=False):
alm = curvedsky.map2alm(tqu, lmax=LMAX)
teb = curvedsky.alm2map(alm[:, None], tqu.copy()[:, None], spin=0)[:, 0]
if wantCl:
cls = healpy.sphtfunc.alm2cl(alm)
return teb, cls
if wantAlmAndCl:
cls = healpy.sphtfunc.alm2cl(alm)
return teb, alm, cls
else:
return teb
def check_simulation(a, b, map_list, sim_list, ivar_list, mask):
"""
Check whether simulated power spectrum P_{ab} is consistent with data. Returns list of (split_sim-coadd,split_data-coadd)
weighted by the mask*effective_ivar.
"""
shape = ivar_list[0].shape
wcs = ivar_list[0].wcs
pmap = enmap.pixsizemap(shape, wcs)
sim_coadd = []
data_coadd = []
for i in range(len(sim_list)):
dsim = sim_list[i] - coadd_map(sim_list, ivar_list)
dsim = dsim * mask * ivar_eff(i, ivar_list) / pmap
testalm = cs.map2alm(dsim, lmax=10000)
testalm = testalm.astype(np.complex128)
testcl = hp.alm2cl(testalm)
sim_coadd.append(testcl)
if a == b:
for i in range(len(map_list)):
dataco = map_list[i][a] - coadd_mapnew(map_list, ivar_list, a)
dataco = dataco * mask * ivar_eff(i, ivar_list) / pmap
testalm = cs.map2alm(dataco, lmax=10000)
testalm = testalm.astype(np.complex128)
testcl = hp.alm2cl(testalm)
data_coadd.append(testcl)
else:
for i in range(len(map_list)):
data_a = map_list[i][a] - coadd_mapnew(map_list, ivar_list, a)
data_a = data_a * mask * ivar_eff(i, ivar_list) / pmap
data_b = map_list[i][b] - coadd_mapnew(map_list, ivar_list, b)
data_b = data_b * mask * ivar_eff(i, ivar_list) / pmap
testalm_a = cs.map2alm(data_a, lmax=10000)
testalm_a = testalm_a.astype(np.complex128)
testalm_b = cs.map2alm(data_b, lmax=10000)
testalm_b = testalm_b.astype(np.complex128)
testcl = hp.alm2cl(testalm_a, testalm_b)
data_coadd.append(testcl)
sim_coadd = np.array(sim_coadd)
data_coadd = np.array(data_coadd)
return (sim_coadd, data_coadd)
ikappa.wcs = wcs
assert Xmap.shape == shape
assert ikappa.shape == shape
### DO FULL SKY RECONSTRUCTION
print("Calculating unnormalized full-sky kappa...")
lcltt = theory.lCl('TT', range(lmax))
with bench.show("reconstruction"):
ukappa_alm = qe.qe_tt_simple(Xmap,
lcltt=lcltt,
nltt_deconvolved=0.,
lmin=2,
lmax=lmax)
# alms of input kappa
ik_alm = cs.map2alm(ikappa, lmax=mlmax).astype(np.complex128)
# if use_saved:
# p2k = ells*(ells+1.)/2.
# ik_alm = hp.almxfl(ik_alm,p2k)
# alms of reconstruction
kappa_alm = hp.almxfl(ukappa_alm, Al).astype(np.complex128)
# cross and auto powers
cri = hp.alm2cl(ik_alm, kappa_alm)
crr = hp.alm2cl(kappa_alm)
cii = hp.alm2cl(ik_alm)
ls = np.arange(len(crr))
cri[ls < 2] = np.nan
crr[ls < 2] = np.nan
box = np.deg2rad([[-1, -1], [10, 10]])
#shape,wcs = enmap.fullsky_geometry(res=np.deg2rad(res/60.))
# shape,wcs = enmap.geometry(pos=box,res=np.deg2rad(res/60.))
hmap = hp.read_map(hfile)
ofunc = lambda shape, wcs: reproject.enmap_from_healpix_interp(
hmap, shape, wcs, rot=None, interpolate=False)
cmap = populate(shape, wcs, ofunc, maxpixy=2000,
maxpixx=2000) #enmap.downgrade(populate(shape,wcs,ofunc),2)
# enmap.write_map("cmap.fits",cmap)
# io.plot_img(cmap,"fmap.png")
cmap = cmap.submap(box)
io.hplot(cmap, "cmap", grid=False)
sys.exit()
pcls = hp.alm2cl(curvedsky.map2alm(cmap, lmax=6000))
hcls = hp.anafast(hmap, lmax=6000)
pells = range(len(pcls))
hells = range(len(hcls))
pl = io.Plotter()
pl.add(pells, (pcls - hcls) / pcls, label='car')
#pl.add(pells,pcls*pells,label='car')
#pl.add(hells,hcls*hells,label='healpix')
pl._ax.set_xlim(200, 4000)
pl._ax.set_ylim(-0.2, 0.1)
pl.done("cls.png")
ellvals = np.arange(lmax)
#factor to change kappa into phi
func = 2. / (ellvals * (ellvals + 1.))
func[0] = 0
phi_alm = healpy.sphtfunc.almxfl(kappa_alms, func)
cmb_powers = get_cmb_powerspectra.websky_cmb_spectra()['unlensed_scalar']
cmb_alm, cmb_ainfo = curvedsky.rand_alm(cmb_powers, lmax = lmax, seed = 1, return_ainfo = True)
unlensed_map, lensed_map = lensing.lens_map_curved((3,) + shape, wcs, phi_alm, cmb_alm, ainfo, output = 'ul')
cmb_alm_lensed = curvedsky.map2alm(lensed_map, ainfo = cmb_ainfo)
healpy.fitsfunc.write_alm(p['outdir'] + 'lensed_alm.fits',
np.complex128(cmb_alm_lensed), overwrite = True)
healpy.fitsfunc.write_alm(p['outdir'] + 'unlensed_alm.fits',
np.complex128(cmb_alm), overwrite = True)
degs = 3
crange = [-300, 300]
unl_cutout = unlensed_map[0, shape[0] / 2 - int(np.rint(degs * 60 / p['PIX_SIZE_ARCMIN'])) :
shape[0] / 2 + int(np.rint(degs * 60 / p['PIX_SIZE_ARCMIN'])),
def get_sim(healpix, homogeneous, white, scale_to_rms=None):
if not (healpix):
mask = enmap.read_map(
"/scratch/r/rbond/msyriac/data/depot/actlens/car_mask_lmax_3000_apodized_2.0_deg.fits"
)[0]
shape, wcs = mask.shape, mask.wcs
nside = None
w2 = wfactor(2, mask)
map2alm = lambda x, lmax: cs.map2alm(x, lmax=lmax)
else:
shape = None
wcs = None
mask = hp.read_map(
"/scratch/r/rbond/msyriac/data/depot/solenspipe/lensing_mask_nside_2048_apodized_4.0.fits"
)
nside = 2048
w2 = wfactor(2, mask, equal_area=True)
map2alm = lambda x, lmax: hp.map2alm(x, lmax=lmax)
tube = 'LT2'
with bench.show("init"):
nsim = mapsims.SONoiseSimulator \
( \
nside=nside,
shape=shape,
wcs=wcs,
ell_max=None,
return_uK_CMB=True,
sensitivity_mode="baseline",
apply_beam_correction=False,
apply_kludge_correction=True,
homogeneous=homogeneous,
no_power_below_ell=None,
rolloff_ell=50,
survey_efficiency=0.2,
full_covariance=True,
LA_years=5,
LA_noise_model="SOLatV3point1",
elevation=50,
SA_years=5,
SA_one_over_f_mode="pessimistic",
sky_fraction=None,
cache_hitmaps=True,
boolean_sky_fraction=False,
)
with bench.show("sim"):
omap = nsim.simulate(
tube,
output_units="uK_CMB",
seed=None,
nsplits=1,
mask_value=0,
atmosphere=not (white),
hitmap=None,
white_noise_rms=scale_to_rms,
)
# io.hplot(omap[0][0][0],f'{out_path}/mapsims_sim_homogeneous_{homogeneous}_white_{white}_scale_to_rms_{scale_to_rms}',downgrade=4,grid=True,ticks=20)
with bench.show("nells"):
ell, ps_T, ps_P, fsky, wnoise_power, hitmaps = nsim.get_noise_properties(
tube,
nsplits=1,
hitmap=None,
white_noise_rms=scale_to_rms,
atmosphere=not (white))
with bench.show("ivar"):
ivar = nsim.get_inverse_variance(tube,
output_units="uK_CMB",
hitmap=None,
white_noise_rms=scale_to_rms)
# Calculate raw power spectrum
imap = omap[0][0][0]
imap = imap * mask
alm = map2alm(imap, lmax=4000)
cls = hp.alm2cl(alm) / w2
ls = np.arange(len(cls))
pl = io.Plotter('Cell')
pl.add(ls, cls)
pl.add(ell, ps_T[0], ls='--')
pl.done(
f'{out_path}/mapsims_nells_homogeneous_{homogeneous}_white_{white}_scale_to_rms_{scale_to_rms}_healpix_{healpix}.png'
)
# Calculate whitened map power spectrum
imap = np.nan_to_num(omap[0][0][0]) * np.sqrt(
ivar[0] / nsim.pixarea_map) * mask
# io.hplot(imap,f'{out_path}/mapsims_wsim_homogeneous_{homogeneous}_white_{white}_scale_to_rms_{scale_to_rms}',downgrade=4,grid=True,ticks=20)
alm = map2alm(imap, lmax=4000)
cls = hp.alm2cl(alm) / w2
ls = np.arange(len(cls))
pl = io.Plotter('Cell')
pl.add(ls, cls)
pl.add(ell, ps_T[0] / wnoise_power[0])
pl.done(
f'{out_path}/mapsims_wnells_homogeneous_{homogeneous}_white_{white}_scale_to_rms_{scale_to_rms}_healpix_{healpix}.png'
)
solution,
deproject=deproject,
data_comb=dcomb,
version="v1.1.1",
sim_index=None))
mask = enmap.read_map(
tutils.get_generic_fname(tdir,
region,
solution,
deproject=deproject,
data_comb=dcomb,
version="v1.1.1",
sim_index=None,
mask=True))
malm = cs.map2alm(mask, lmax=lmax)
ialm = cs.map2alm(imap, lmax=lmax)
malm = malm.astype(np.complex128, copy=False)
malm = rotate_alm(malm, 2000.0, 2000.0, 'C', 'G')
ialm = ialm.astype(np.complex128, copy=False)
ialm = rotate_alm(ialm, 2000.0, 2000.0, 'C', 'G')
imask = maps.binary_mask(hp.alm2map(malm, nside))
hmap = hmap + hp.alm2map(ialm, nside) * imask
hmask = hmask + imask
io.mollview(hmap, f'{froot}test_rot_{deproject}.png')
io.mollview(hmask, f'{froot}test_rot_mask_{deproject}.png')
def load_alms_base(self,
set_idx,
sim_idx,
cache=True,
fg_override=None,
ret_alm=False,
fgflux="15mjy",
alm_file_postfix=''):
# note: beam is set to false
print("loading alm base")
#cmb_file = os.path.join(self.signal_path, 'fullsky%s_alm_set%02d_%05d%s.fits' %(self.cmb_type, set_idx, 0, alm_file_postfix))
alm_signal = self.load_cmb_alm(set_idx, sim_idx, alm_file_postfix)
# check if rotation is needed
if self.apply_rotation:
if not np.any(self.alpha_map):
print(
"[Warning] want to apply rotation but alpha map is not provided..."
)
else:
alpha_map = self.alpha_map
print("applying rotation field")
# get wcs from rot_map: doesn't matter which wcs is used because eventually
# it is converted back to alm
wcs = alpha_map.wcs
shape = (3, ) + alpha_map.shape[-2:]
# generate cmb map from alm_signal
cmb_map = enmap.empty(shape, wcs)
curvedsky.alm2map(alm_signal, cmb_map)
# apply the rotation field
cmb_map = enmap.rotate_pol(cmb_map, 2 * alpha_map)
# convert back to alm
lmax = hp.Alm.getlmax(alm_signal.shape[-1])
alm_signal = curvedsky.map2alm(cmb_map, lmax=lmax)
del cmb_map
alm_signal = np.tile(alm_signal, (len(self.freqs), 1, 1))
if fg_override is not None:
print("[Warning] override the default foreground flag....")
add_foregrounds = fg_override if fg_override is not None else self.add_foregrounds
if add_foregrounds:
print("adding fgs to the base")
alm_fg90_150 = self.load_alm_fg(set_idx, sim_idx, fgflux=fgflux)
lmax_sg = hp.Alm.getlmax(alm_signal.shape[-1])
alm_out = np.zeros((len(self.freqs), 3, len(alm_fg90_150[-1])),
dtype=np.complex128)
lmax_fg = hp.Alm.getlmax(alm_fg90_150.shape[-1])
for idx, freq in enumerate(self.freqs):
freq_idx = 1 if freq == 'f150' else 0
alm_out[idx, 0, :] = alm_fg90_150[freq_idx, :].copy()
for m in range(lmax_sg + 1):
lmin = m
lmax = lmax_sg
idx_ssidx = hp.Alm.getidx(lmax_sg, lmin, m)
idx_seidx = hp.Alm.getidx(lmax_sg, lmax, m)
idx_fsidx = hp.Alm.getidx(lmax_fg, lmin, m)
idx_feidx = hp.Alm.getidx(lmax_fg, lmax, m)
alm_out[..., idx_fsidx:idx_feidx +
1] = alm_out[..., idx_fsidx:idx_feidx +
1] + alm_signal[...,
idx_ssidx:idx_seidx + 1]
alm_signal = alm_out.copy()
del alm_out, alm_fg90_150
if cache:
self.alms_base[self.get_base_alm_idx(set_idx,
sim_idx)] = alm_signal.copy()
if ret_alm: return alm_signal
del alm_signal
def comp_cltt(alex):
if alex:
lmax = 5000
res = 1.0
camb_root = "/gpfs01/astro/workarea/msyriac/data/act/theory/cosmo2017_10K_acc3"
sim_name = "fullskyLensedMapUnaberrated_T_00002.fits"
#sim_name = "fullskyLensedMap_T_00000.fits"
ksim_name = "kappaMap_00002.fits"
sim_location = "/gpfs01/astro/workarea/msyriac/data/sims/dw/"
else:
lmax = 3000
res = 1.5
camb_root = "/astro/u/msyriac/repos/quicklens/quicklens/data/cl/planck_wp_highL/planck_lensing_wp_highL_bestFit_20130627"
sim_name = "lensed_map_seed_1.fits"
ksim_name = "kappa_map_seed_1.fits"
sim_location = "/gpfs01/astro/workarea/msyriac/data/depot/falafel/"
theory = cosmology.loadTheorySpectraFromCAMB(camb_root,
get_dimensionless=False)
Xmap = enmap.read_map(sim_location + sim_name)
ikappa = enmap.read_map(sim_location + ksim_name)
if alex:
shape, wcs = enmap.fullsky_geometry(res=np.deg2rad(res / 60.),
proj="car")
Xmap.wcs = wcs
ikappa.wcs = wcs
calm = cs.map2alm(Xmap, lmax=lmax).astype(np.complex128)
kalm = cs.map2alm(ikappa, lmax=lmax).astype(np.complex128)
cls = hp.alm2cl(calm)
clskk = hp.alm2cl(kalm)
lsc = np.arange(0, len(cls), 1)
print(cls[:5])
print(clskk[:5])
cls[lsc < 2] = np.nan
clskk[lsc < 2] = np.nan
bin_edges = np.logspace(np.log10(2), np.log10(lmax), 40)
binner = stats.bin1D(bin_edges)
ls = np.arange(0, lmax, 1)
cltt = theory.lCl('TT', ls)
clkk = theory.gCl('kk', ls)
cents, btt = binner.binned(ls, cltt)
cents, bcc = binner.binned(lsc, cls)
cents, bttkk = binner.binned(ls, clkk)
cents, bcckk = binner.binned(lsc, clskk)
# plot
pl = io.Plotter(yscale='log',
xscale='log',
xlabel='$\\ell$',
ylabel='$C_{\\ell}$')
pl.add(ls, ls**2. * theory.lCl('TT', ls), lw=3, color='k')
pl.add(cents, cents**2. * btt, ls="none", marker="x")
pl.add(cents, cents**2. * bcc, ls="none", marker="o")
pl.done(io.dout_dir + "fullsky_sim_test_cmb_alex_" + str(alex) + ".png")
pl = io.Plotter(yscale='log', xscale='log', xlabel='$L$', ylabel='$C_L$')
pl.add(ls, theory.gCl('kk', ls), lw=3, color='k')
pl.add(cents, bttkk, ls="none", marker="x")
pl.add(cents, bcckk, ls="none", marker="o")
pl.done(io.dout_dir + "fullsky_sim_test_kk_alex_" + str(alex) + ".png")
pl = io.Plotter(xscale='log', xlabel='$L$', ylabel='$\\delta C_L / C_L$')
pl.add(cents, (bcc - btt) / btt, ls="-", marker="o", label="cmb")
pl.add(cents, (bcckk - bttkk) / bttkk, ls="-", marker="o", label="kk")
pl.hline()
pl._ax.set_ylim(-0.2, 0.2)
pl.done(io.dout_dir + "fullsky_sim_test_alex_diff_" + str(alex) + "_2.png")
def get_maps(self, rot_angle1, rot_angle2, compts=None, use_sht=True, ret_alm=True, transfer=None,
load_processed=False, save_processed=False, flux_cut=None):
if compts is None: compts = self.compts
shape, wcs = self.geometry
nshape = (len(compts),) + shape[-2:]
ret = enmap.zeros(nshape, wcs)
if load_processed and not ret_alm:
for i, compt_idx in enumerate(compts):
input_file = self.get_fits_path(self.processed_dir, rot_angle1, rot_angle2, compt_idx)
print("loading", input_file)
temp = enmap.read_map(input_file)
ret[i, ...] = enmap.extract(temp, shape, wcs).copy()
del temp
return ret
else:
for i, compt_idx in enumerate(compts):
if "pts" not in compt_idx:
input_file = self.get_fits_path(self.input_dir, rot_angle1, rot_angle2, compt_idx)
print("loading", input_file)
alm = np.complex128(hp.read_alm(input_file, hdu=(1)))
ret[i, ...] = curvedsky.alm2map(alm, enmap.zeros(nshape[1:], wcs))
else:
input_file = self.get_fits_path(self.input_dir, rot_angle1, rot_angle2, compt_idx,
fits_type="enmap")
print("loading", input_file)
temp = enmap.read_map(input_file)
ret[i, ...] = enmap.extract(temp, shape, wcs).copy()
del temp
alms = None
if transfer is not None:
l, f = transfer
interp_func = scipy.interpolate.interp1d(l, f, bounds_error=False, fill_value=0.)
if use_sht:
l_intp = np.arange(self.lmax + 1)
f_int = interp_func(l_intp)
alms = curvedsky.map2alm(ret, lmax=self.lmax, spin=0)
for i in range(len(compts)):
alms[i] = hp.almxfl(alms[i], f_int)
ret = curvedsky.alm2map(alms, ret, spin=0)
else:
ftmap = enmap.fft(ret)
f_int = interp_func(enmap.modlmap(shape, wcs).ravel())
ftmap = ftmap * np.reshape(f_int, (shape[-2:]))
ret = enmap.ifft(ftmap).real;
del ftmap
if save_processed:
raise NotImplemented()
if flux_cut is not None:
flux_map = flux_cut / enmap.pixsizemap(shape, wcs)
flux_map *= 1e-3 * jysr2thermo(148)
for i, compt_idx in enumerate(compts):
if "pts" not in compt_idx: continue
loc = np.where(ret[i] > flux_map)
ret[i][loc] = 0.
del flux_map
if ret_alm and alms is None:
alms = curvedsky.map2alm(ret, lmax=self.lmax, spin=0)
return ret if not ret_alm else (ret, alms)
def map2alm(self,imap,lmax):
if self.hpix:
return hp.map2alm(imap,lmax=lmax,iter=self.iter)
else:
return cs.map2alm(imap,lmax=lmax)
test_nell_standard = False
if test_nell_standard:
shape, wcs = enmap.fullsky_geometry(res=2.0 * utils.arcmin)
noise = 10.0
lknee = 3000
alpha = -4
pmap = maps.psizemap(shape, wcs)
ivar = maps.ivar(shape, wcs, noise, ipsizemap=pmap)
imap = maps.modulated_noise_map(ivar,
lknee=lknee,
alpha=alpha,
lmax=4000,
lmin=50)
alm = cs.map2alm(imap * np.sqrt(ivar / pmap), lmax=4000)
cls = hp.alm2cl(alm)
ls = np.arange(len(cls))
N_ell = maps.atm_factor(ls, lknee, alpha) + 1.
N_ell[ls < 50] = 0
pl = io.Plotter('Cell')
pl.add(ls, cls)
pl.add(ls, N_ell, ls='--')
pl.done(f'{out_path}/mapsims_nells_test.png')
def get_sim(healpix, homogeneous, white, scale_to_rms=None):
if not (healpix):
mask = enmap.read_map(
"/scratch/r/rbond/msyriac/data/depot/actlens/car_mask_lmax_3000_apodized_2.0_deg.fits"
# aberration.beta, modulation = False)
print('rank', rank, 'doing cmbSet', cmbSet, 'iii' , iii, \
', iMin', iMin, ', iMax', iMax, 'calling aberration.boost_map', time.time() - start)
lTquMap, AForMod = aberration.boost_map(lTquMap,
aberration.dir_equ,
aberration.beta,
modulation=None,
return_modulation=True)
mapList += [unaberrated]
mapNameList += ['fullskyLensedUnabberatedCMB']
for mi, mmm in enumerate(mapList):
print(iii, ' calling curvedsky.map2alm for ', mapNameList[mi])
alm = curvedsky.map2alm(mmm, lmax=p['LMAX_WRITE'])
cmbDir = p['dataDir']
print('writing to disk')
filename = cmbDir + "/%s_alm_%s%05d.fits" \
% ( mapNameList[mi],
('set%02d_' % cmbSet if 'CMB' in mapNameList[mi] else '' ) ,
iii)
healpy.fitsfunc.write_alm(filename,
np.complex64(alm),
overwrite=True)
# cls[iii] += [healpy.sphtfunc.alm2cl(alm)]
def _generate_foreground_148GHz(self, seed=None, verbose=True, input_kappa=None, post_processes=[],
flux_cut=7, polfix=True):
if input_kappa is None:
if verbose: print("making input gaussian kappa")
gaussian_kappa = self._generate_gaussian_kappa(seed=seed)
else:
if input_kappa.ndim == 2: input_kappa = input_kappa[np.newaxis, ...]
gaussian_kappa = input_kappa
Ny, Nx = self.shape
Ny_pad, Nx_pad = self.shape_padded
ny, nx = self.stamp_shape[-2:]
taper_width = self.taper_width
nbatchy, nbatchx = self.num_batch
xgrid = self._get_xgrid()
taper = self._get_taper()
batch_idxes = np.array_split(np.arange(nbatchy), min(nbatchy, self.nprocess))
xin = np.arange(Nx + 1)
if verbose: print("start sampling")
global _get_sampled
def _get_sampled(batch_idxes):
retysidx = batch_idxes[0] * (ny)
retyeidx = (batch_idxes[-1] + 1) * (ny)
ret = np.zeros((retyeidx - retysidx, Nx_pad))
for i, yidx in enumerate(batch_idxes):
ysidx = yidx * (ny - taper_width)
yeidx = min(ysidx + ny, Ny)
yoffset = yidx * taper_width
for ycidx in np.arange(ysidx, yeidx):
if ycidx >= Ny:
continue
yocidx = ycidx + yoffset
yrcidx = yocidx - retysidx
yvals = np.append(gaussian_kappa[0, ycidx, :], gaussian_kappa[0, ycidx, 0])
fit = scipy.interpolate.CubicSpline(xin, yvals, bc_type="periodic")
xval = xgrid[yocidx, :] % Nx
ret[yrcidx, :] = fit(xval)
return ret
with Pool(len(batch_idxes)) as p:
gaussian_kappa = p.map(_get_sampled, batch_idxes)
del _get_sampled, xin
gc.collect()
gaussian_kappa = np.vstack(gaussian_kappa)
gaussian_kappa = gaussian_kappa[np.newaxis, ...]
if verbose: print("end sampling")
if polfix:
if verbose: print("pol fix")
gaussian_kappa[:, :1 * ny, :] = np.flip(gaussian_kappa[:, 1 * ny:2 * ny, :], 1)
gaussian_kappa[:, -1 * ny:, :] = np.flip(gaussian_kappa[:, -2 * ny:-1 * ny, :], 1)
gaussian_kappa = self.normalizer(gaussian_kappa)
def process_ml(input_imgs, batch_maker):
input_imgs = batch_maker(input_imgs)
nsample = input_imgs.shape[0]
output_imgs = np.zeros((nsample, 5, ny, nx))
ctr = 0
nitr = int(np.ceil(input_imgs.shape[0] / self.nbatch))
for batch in np.array_split(np.arange(input_imgs.shape[0]), nitr):
input_tensor = torch.autograd.Variable(self.Tensor(input_imgs[batch].copy()))
ret = self.pixgan_generator(input_tensor).detach()
ret = torch.cat((input_tensor, ret), 1)
ret = self.forse_generator(ret).detach()
output_imgs[batch] = ret.data.to(device="cpu").numpy()
if verbose and ctr % 20 == 0:
print(f"batch {ctr}/{nitr} completed")
ctr += 1
return output_imgs
def post_process(output_imgs, unbatch_maker):
output_imgs = unbatch_maker(output_imgs)
output_imgs = output_imgs[0, ...]
return output_imgs
if verbose: print("make the primary images")
batch_maker = transforms.Batch((1, ny, nx))
unbatch_maker = transforms.UnBatch((1, Ny_pad, Nx_pad))
processed = process_ml(gaussian_kappa, batch_maker);
del gaussian_kappa
processed = post_process(processed, unbatch_maker)
del batch_maker, unbatch_maker
torch.cuda.empty_cache()
gc.collect()
for post_process in post_processes:
processed = post_process(processed)
processed = self.unnormalizer(processed)
loc = np.where(processed[3:5] < 0)
processed[3:5][loc] = 0.
reprojected = np.zeros((5, Ny, Nx), dtype=np.float32)
for compt_idx in range(0, 5):
if verbose: print(f"reprojecting images {compt_idx}")
method = "interp" if compt_idx < 3 else "nearest"
global _get_reprojected
def _get_reprojected(batch_idxes, method=method):
retysidx = batch_idxes[0] * (ny - taper_width)
retyeidx = min(batch_idxes[-1] * (ny - taper_width) + ny, Ny)
ret = np.zeros((retyeidx - retysidx, Nx))
for yidx in batch_idxes:
ysidx = yidx * (ny - taper_width)
yeidx = min(ysidx + ny, Ny)
yoffset = taper_width * yidx
for xidx in range(nbatchx):
xosidx = xidx * nx
xoeidx = xosidx + nx
for j, ycidx in enumerate(np.arange(ysidx, yeidx)):
if ycidx >= Ny:
continue
yrcidx = ycidx - retysidx
yocidx = ycidx + yoffset
yvals = processed[compt_idx, yocidx, xosidx:xoeidx]
xvals = xgrid[yocidx, xosidx:xoeidx]
xmin = int(np.ceil(xvals[0]))
xmax = int(np.floor(xvals[-1]))
xin = np.arange(xmin, xmax + 1)[:Nx]
yvals = yvals * taper[j, :]
if method == "nearest":
fit = scipy.interpolate.interp1d(xvals, yvals, assume_sorted=True, kind="nearest")
elif method == "interp":
fit = scipy.interpolate.interp1d(xvals, yvals, assume_sorted=True)
else:
assert (False)
ret[yrcidx, xin % Nx] += fit(xin)
return ((retysidx, retyeidx), ret)
with Pool(len(batch_idxes)) as p:
storage = p.map(_get_reprojected, batch_idxes)
for idxes, ring in storage:
reprojected[compt_idx, idxes[0]:idxes[1], :] += ring
del storage, _get_reprojected
gc.collect()
## weight correction for diff
reprojected[:3] = reprojected[:3] / self._get_weight()[0]
reprojected[3:5] = reprojected[3:5] / self._get_weight()[1]
reprojected = enmap.enmap(reprojected.astype(np.float32), wcs=self.wcs)
#return reprojected
## apply transfer functions
kmap = enmap.fft(reprojected[:3])
for j, compt_idx in enumerate(self.fg_compts[:3]):
if verbose: print(f"applying the transfer functions to {compt_idx}")
xtransf = self.transf_2dspec[compt_idx]['px']
ytransf = self.transf_2dspec[compt_idx]['py']
kmap[j] = kmap[j] * np.outer(ytransf, xtransf)
reprojected[j] = enmap.ifft(kmap[j]).real
alm = curvedsky.map2alm(reprojected[j].astype(np.float64), lmax=10000)
alm = hp.almxfl(alm, self.transf_1dspec[compt_idx])
reprojected[j] = curvedsky.alm2map(alm, reprojected[j])
del kmap
def boxcox(arr, lamb):
return ((arr + 1) ** lamb - 1) / lamb
reprojected[3:5] *= 1 / self._get_jysr2thermo(mode="car")
reprojected[3] *= 1.1
loc = np.where(reprojected[3] > 1)
reprojected[3][loc] = reprojected[3][loc] ** 0.63;
del loc
reprojected[4] = boxcox(reprojected[4], 1.25)
loc = np.where(reprojected[3:5] > flux_cut)
reprojected[3:5][loc] = 0.;
del loc
reprojected[3:5] *= self._get_jysr2thermo(mode="car")
gc.collect()
return reprojected.astype(np.float32)
"tsz",
'cmb',
"joint",
beam=True)
dbfile = tutils.get_generic_fname(tdir,
region,
"tsz",
'cib',
"joint",
beam=True)
imask = enmap.read_map(
tutils.get_generic_fname(tdir, region, "tsz", None, "joint",
mask=True))
dmask = hp.alm2map(
cs.map2alm(imask, lmax=lmax).astype(np.complex128), nside)
dmask[dmask < 0.5] = 0
dmask[dmask > 0.5] = 1
jmask = dmask * mask
io.mollview(jmask, os.environ['WORK'] + '/jmask_%s.png' % region)
fsky = jmask.sum() / jmask.size
print(fsky * 41252.)
delta = hdelta * jmask
galm = hp.map2alm(delta, lmax=lmax)
for bfile, imap in zip([ybfile, cbfile, dbfile], [ymap, cmap, dmap]):
yalm = hp.map2alm(
hp.alm2map(cs.map2alm(imap, lmax=lmax).astype(np.complex128),
nside=nside) * jmask,
importlib.reload(plot_ps)
lmin = 2
lmax = 1000
ls = np.arange(0, lmax)
nside = 400
def t(ls):
return ls * (ls + 1)
shape, wcs = enmap.fullsky_geometry(res=20 * utils.arcmin, proj='car')
# generate unlensed map using input unlensd ps
unlensed_ps = load_data.unlensed(lmin, lmax, 'TT').spectra()
unlensed_map = curvedsky.rand_map(shape, wcs, unlensed_ps)
unlensed_alm = curvedsky.map2alm(unlensed_map, lmax=lmax)
# generate deflection potential using input
input_cldd = load_data.input_cldd(lmin, lmax).spectra()
if 1:
phi_lm = fitsfunc.read_alm('./fullskyPhi_alm_00000.fits')
phi_lm = phi_lm.astype(np.cdouble)
cl_phi = sphtfunc.alm2cl(a)[0:lmax]
cl_dd = cl_phi * t(ls)
if 1:
plot_ps.plot(lmin, lmax, [input_cldd, cl_dd], ['1', '2'])
#find spin +2 and -2 transform of t_alm
ells = np.arange(0, lmax)
#multiply by sqrt(ls)
alms = qe.almxfl(np.stack((t_alm, t_alm)),
np.sqrt((ells - 1.) * ells * (ells + 1.) * (ells + 2.)))
print(np.shape(alms))
# return real space spin pm 2 components
rmap = qe.alm2map_spin(alms, 0, 2, omap)
prodmap = rmap * rmapT
realsp2 = prodmap[0] #spin +2 real space
realsm2 = prodmap[0] #spin -2 real space
res1 = cs.map2alm(enmap.enmap(
-qe.irot2d(np.stack((realsp2, realsp2.conj())), spin=0).real, omap.wcs),
spin=2,
lmax=solint.mlmax)
res2 = cs.map2alm(enmap.enmap(
-qe.irot2d(np.stack((realsm2, realsm2.conj())), spin=0).real, omap.wcs),
spin=2,
lmax=solint.mlmax)
#spin 2 ylm
ttalmsp2 = res1[0]
ttalmsm2 = res2[1]
#multiply by sqrt(L(L+1) as a way of multiplying by exponential
qe.almxfl(np.stack((ttalmsm2, ttalmsm2)),
np.sqrt((ells - 1.) * ells * (ells + 1.) * (ells + 2.))) + qe.almxfl(
np.stack((ttalmsp2, ttalmsp2)),
np.sqrt((ells - 1.) * ells * (ells + 1.) * (ells + 2.)))
def get_poisson_srcs_alms(self, set_idx, sim_num, patch, alm_shape, oshape,
owcs):
def deltaTOverTcmbToJyPerSr(freqGHz, T0=2.726):
"""
@brief the function name is self-eplanatory
@return the converstion factor
stolen from Flipper -- van engelen
"""
kB = 1.380658e-16
h = 6.6260755e-27
c = 29979245800.
nu = freqGHz * 1.e9
x = h * nu / (kB * T0)
cNu = 2 * (kB * T0)**3 / (h**2 *
c**2) * x**4 / (4 * (np.sinh(x / 2.))**2)
cNu *= 1e23
return cNu
TCMB_uk = 2.72e6
if oshape[0] > 3:
#then this is a multichroic array, and sadly we only have this at 150 GHz for now
raise Exception('get_poisson_srcs_alms only implemented for 150 GHz so far ' \
+ '(that is the model we currently have for radio sources) ')
else:
freq_ghz = 148
#ideally this RNG stuff would be defined in a central place to
#avoid RNG collisions. Old version is currently commented out at top of
#simgen.py
templ = self.get_template(patch, shape=oshape, wcs=owcs)
templ[:] = 0
seed = seedgen.get_poisson_seed(set_idx, sim_num)
np.random.seed(seed=seed)
#Wasn't sure how to codify this stuff outside this routine - hardcoded for now
S_min_Jy = .001
S_max_Jy = .015
tucci = np.loadtxt(
os.path.join(os.path.dirname(os.path.abspath(__file__)),
'../data/ns_148GHz_modC2Ex.dat'))
S = tucci[:, 0]
dS = S[1:] - S[0:-1]
dS = np.append(dS, [0.])
dNdS = tucci[:, 1]
mean_numbers_per_patch = dNdS * enmap.area(templ.shape, templ.wcs) * dS
numbers_per_fluxbin = np.random.poisson(mean_numbers_per_patch)
#note pixel areas not constant for pixell maps
pixel_areas = enmap.pixsizemap(templ.shape, templ.wcs)
for si, fluxval in enumerate(S[S <= S_max_Jy]):
xlocs = np.random.randint(templ.shape[-1],
size=numbers_per_fluxbin[si])
ylocs = np.random.randint(templ.shape[-2],
size=numbers_per_fluxbin[si])
#add the value in jy / sr, i.e. divide by the solid angle of a pixel.
templ[0, ylocs, xlocs] += fluxval / pixel_areas[ylocs, xlocs]
map_factor = TCMB_uk / deltaTOverTcmbToJyPerSr(freq_ghz)
templ *= map_factor
#GET ALMs
output = curvedsky.map2alm(templ[0], lmax=hp.Alm.getlmax(alm_shape[0]))
return output
def hdowngrade(imap, shape, wcs, lmax):
return (cs.alm2map(cs.map2alm(imap, lmax=lmax),
enmap.empty(shape, wcs, dtype=imap.dtype)))
def get_maps(self, rot_angle1, rot_angle2, compts=None, use_sht=True, ret_alm=True, transfer=None,
load_processed=False, save_processed=False, flux_cut=None):
if compts is None: compts = self.compts
shape, wcs = self.geometry
nshape = (len(compts),) + shape[-2:]
ret = enmap.zeros(nshape, wcs)
if load_processed and not ret_alm:
for i, compt_idx in enumerate(compts):
input_file = self.get_fits_path(self.processed_dir, rot_angle1, rot_angle2, compt_idx)
print("loading", input_file)
temp = enmap.read_map(input_file)
ret[i, ...] = enmap.extract(temp, shape, wcs).copy()
del temp
return ret
else:
for i, compt_idx in enumerate(compts):
input_file = self.get_fits_path(self.input_dir, rot_angle1, rot_angle2, compt_idx)
print("loading", input_file)
alm = np.complex128(hp.read_alm(input_file, hdu=(1)))
ret[i, ...] = curvedsky.alm2map(alm, enmap.zeros(nshape[1:], wcs))
del alm
if compt_idx in self.highflux_cats:
print("adding high flux cats")
hiflux_cat = np.load(self.get_highflux_cat_path(compt_idx))
hiflux_cat[:, :2] = car2hp_coords(hiflux_cat[:, :2])
mat_rot, _, _ = hp.rotator.get_rotation_matrix(
(rot_angle1 * utils.degree * -1, rot_angle2 * utils.degree, 0))
uvec = hp.ang2vec(hiflux_cat[:, 0], hiflux_cat[:, 1])
rot_vec = np.inner(mat_rot, uvec).T
temppos = hp.vec2ang(rot_vec)
rot_pos = np.zeros(hiflux_cat[:, :2].shape)
rot_pos[:, 0] = temppos[0]
rot_pos[:, 1] = temppos[1]
rot_pos = hp2car_coords(rot_pos)
del temppos
rot_pix = np.round(enmap.sky2pix(nshape[-2:], wcs, rot_pos.T).T).astype(np.int)
loc = np.where((rot_pix[:, 0] >= 0) & (rot_pix[:, 0] < nshape[-2]) & (rot_pix[:, 1] >= 0.) & (
rot_pix[:, 1] < nshape[-1]))
hiflux_cat = hiflux_cat[loc[0], 2]
rot_pix = rot_pix[loc[0], :]
hiflux_map = enmap.zeros(nshape[-2:], wcs)
hiflux_map[rot_pix[:, 0], rot_pix[:, 1]] = hiflux_cat
if flux_cut is not None:
tmin = flux_cut * 1e-3 * jysr2thermo(148)
loc = np.where(hiflux_map > tmin)
hiflux_map[loc] = 0
hiflux_map = hiflux_map / enmap.pixsizemap(shape, wcs)
ret[i, ...] = ret[i, ...] + hiflux_map
del hiflux_map
alms = None
if transfer is not None:
l, f = transfer
interp_func = scipy.interpolate.interp1d(l, f, bounds_error=False, fill_value=0.)
if use_sht:
l_intp = np.arange(self.lmax + 1)
f_int = interp_func(l_intp)
alms = curvedsky.map2alm(ret, lmax=self.lmax, spin=0)
for i in range(len(compts)):
alms[i] = hp.almxfl(alms[i], f_int)
ret = curvedsky.alm2map(alms, ret, spin=0)
else:
ftmap = enmap.fft(ret)
f_int = interp_func(enmap.modlmap(shape, wcs).ravel())
ftmap = ftmap * np.reshape(f_int, (shape[-2:]))
ret = enmap.ifft(ftmap).real;
del ftmap
if save_processed:
raise NotImplemented()
if ret_alm and alms is None:
alms = curvedsky.map2alm(ret, lmax=self.lmax, spin=0)
return ret if not ret_alm else (ret, alms)
if not (args.skip_aberration):
if rank == 0:
logging.info('doing aberration')
#Note - we are aberrating and not modulating! The
#modulation is a frequency-dependent, so is done
#later.
with bench.mark("boost"):
l_tqu_map_aberrated = ab.aberrate(l_tqu_map)
if rank == 0:
logging.info(f'BENCH:\n{bench.stats}')
map_list += [l_tqu_map_aberrated]
map_name_list += ['fullskyLensedAberratedCMB']
for mi, mmm in enumerate(map_list):
if rank == 0:
logging.info(f'curvedsky.map2alm for {map_name_list[mi]}')
alm = curvedsky.map2alm(mmm, lmax=args.lmax_write)
filename = cmb_dir + f"/{map_name_list[mi]}_alm_cmb_set_{cmb_set:02d}_phi_set_{phi_set:02d}_{iii:05d}.fits"
if rank == 0:
logging.info(f'writing to disk, filename is {filename}')
healpy.fitsfunc.write_alm(filename, np.complex64(alm), overwrite=True)
if rank == 0:
logging.info(f'rank 0: {(j+1.)*100./len(my_tasks)} % done...')
