def __signal_postprocessing__(self,
patch,
signal_idx,
alm_patch,
save_map,
oshape,
owcs,
apply_window=True):
signal = self.get_template(patch, shape=oshape, wcs=owcs)
signal = signal if len(alm_patch.shape) > 1 else signal[0, ...]
curvedsky.alm2map(alm_patch, signal, spin=[0, 2], verbose=True)
if apply_window:
print('apply window')
axes = [-2, -1]
for idx in range(signal.shape[0]):
kmap = pfft.fft(signal[idx], axes=axes)
wy, wx = enmap.calc_window(kmap.shape)
wind = wy[:, None]**1 * wx[None, :]**1
kmap *= wind
signal[idx] = (pfft.ifft(kmap, axes=axes, normalize=True)).real
del kmap
if save_map:
self.signals[signal_idx] = signal.copy()
self.manage_cache(self.signals, self.max_cached)
return signal
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 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 compute_output_map(self, alm):
if self.nside is None:
assert (self.shape is not None) and (self.wcs is not None)
n_comp = 3 if self.has_polarization else 1
output_map = enmap.empty((n_comp,) + self.shape[-2:], self.wcs)
curvedsky.alm2map(alm, self.output_map, spin=[0, 2], verbose=True)
elif self.nside is not None:
output_map = hp.alm2map(alm, self.nside)
else:
raise ValueError("You must specify either nside or both of shape and wcs")
return output_map
def __init__(self,camb_file='./CMB_fiducial_totalCls.dat',seed=1,pixRes_arcmin=4,lmax_sim=500):
cls_camb = np.loadtxt(camb_file,unpack=True)
cl_tt = cls_camb[1]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_tt = np.append([0,0],cl_tt)
cl_ee = cls_camb[2]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_ee = np.append([0,0],cl_ee)
cl_bb = cls_camb[3]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)*0.05
cl_bb = np.append([0,0],cl_bb)
cl_te = cls_camb[4]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_te = np.append([0,0],cl_te)
ells = np.append([0,1],cls_camb[0])
self.pixRes_arcmin=pixRes_arcmin/180./60*np.pi
self.lmax_sim = lmax_sim
shape,wcs = enmap.fullsky_geometry(self.pixRes_arcmin,dims=(2,))
np.random.seed(seed)
self.shape = shape
self.wcs = wcs
self.alm_cmb_E = curvedsky.rand_alm(cl_ee,lmax=lmax_sim)#hp.synalm(cl_ee,lmax=3*nside-1,new=True)
self.alm_cmb_B = curvedsky.rand_alm(cl_bb,lmax=lmax_sim)#hp.synalm(cl_bb,lmax=3*nside-1,new=True)
tmp_empty = enmap.empty(shape,wcs)
self.Q_cmb,self.U_cmb = curvedsky.alm2map(np.array([self.alm_cmb_E,self.alm_cmb_B]),tmp_empty,spin=[2])
ps = np.zeros([2,2,lmax_sim])
ps[0,0] = synchrotron_Cl('E','E')[:lmax_sim]
ps[1,1] = galaticDust_Cl('E','E')[:lmax_sim]
ps[1,0] = syncxdust_Cls('E','E')[:lmax_sim]
self.alm_sync_E,self.alm_dust_E = curvedsky.rand_alm(ps,lmax=lmax_sim)
ps = np.zeros([2,2,lmax_sim])
ps[0,0] = synchrotron_Cl('B','B')[:lmax_sim]
ps[1,1] = galaticDust_Cl('B','B')[:lmax_sim]
ps[1,0] = syncxdust_Cls('B','B')[:lmax_sim]
self.alm_sync_B,self.alm_dust_B = curvedsky.rand_alm(ps,lmax=lmax_sim)
tmp_empty = enmap.empty(shape,wcs)
self.Q_dust,self.U_dust = curvedsky.alm2map(np.array([self.alm_dust_E,self.alm_dust_B]),tmp_empty,spin=[2])
tmp_empty = enmap.empty(shape,wcs)
self.Q_sync,self.U_sync = curvedsky.alm2map(np.array([self.alm_sync_E,self.alm_sync_B]),tmp_empty,spin=[2])
def __init__(self,camb_file='./CMB_fiducial_totalCls.dat',seed=1,pixRes_arcmin=2.,lmax_sim=500,nside_pysm=512):
cls_camb = np.loadtxt(camb_file,unpack=True)
cl_tt = cls_camb[1]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_tt = np.append([0,0],cl_tt)
cl_ee = cls_camb[2]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_ee = np.append([0,0],cl_ee)
cl_bb = cls_camb[3]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_bb = np.append([0,0],cl_bb)
cl_te = cls_camb[4]/(cls_camb[0]*(cls_camb[0]+1)/2/np.pi)
cl_te = np.append([0,0],cl_te)
ells = np.append([0,1],cls_camb[0])
self.pixRes_arcmin=pixRes_arcmin/180./60*np.pi
self.lmax_sim = lmax_sim
shape,wcs = enmap.fullsky_geometry(self.pixRes_arcmin,dims=(2,))
self.shape = shape
self.wcs = wcs
np.random.seed(seed)
self.nside_pysm =nside_pysm
self.alm_cmb_E = curvedsky.rand_alm(cl_ee,lmax=lmax_sim)#hp.synalm(cl_ee,lmax=3*nside-1,new=True)
self.alm_cmb_B = curvedsky.rand_alm(cl_bb,lmax=lmax_sim)#hp.synalm(cl_bb,lmax=3*nside-1,new=True)
tmp_empty = enmap.empty(shape,wcs)
self.Q_cmb,self.U_cmb = curvedsky.alm2map(np.array([self.alm_cmb_E,self.alm_cmb_B]),tmp_empty,spin=[2])
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_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 alm2stripe(alm, width, resol, proj='car'):
"""Return the stripe centered at dec=0 of the map corresponding to alm.
Return a slice of the map corresponding to alm from declination -width/2
to +width/2, and right ascension -180deg to +180deg.
Parameters
----------
alm: array
The set of alms to convert to map-space.
width: value
The width of the map (centered at dec=0) in degrees.
resol: value
The resolution of the map in arcmin.
proj: str, default='car'
The projection of the map. Must be compatible with pixell.
Returns
-------
stmap: array
The output map.
"""
box = np.array([[-width/2,180], [width/2,-180]]) * utils.degree
shape, wcs = enmap.geometry(pos=box, res=resol*utils.arcmin, proj=proj)
cmap = enmap.zeros(shape, wcs)
return curvedsky.alm2map(alm, cmap, method='cyl')
def get_map(self,seed=None):
"""Get correlated galaxy and kappa map """
from pixell import curvedsky
alms = curvedsky.rand_alm_healpy(self.ps, lmax=int(self.lmax)+1, seed=seed)
ncomp = 1 if len(self.shape)==2 else self.shape[0]
omap = enmap.empty((ncomp,)+self.shape[-2:], self.wcs, dtype=np.float64)
omap = curvedsky.alm2map(alms, omap, spin=0)
return alms,omap
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 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 alm2map(self,alm,spin,ncomp,mlmax):
if self.hpix:
if spin!=0:
res = hp.alm2map_spin(alm,nside=self.nside,spin=spin,lmax=mlmax)
return res
else:
# complex64 not supported here
return hp.alm2map(alm.astype(np.complex128),nside=self.nside,pol=False)[None]
else:
omap = enmap.empty((ncomp,)+self.shape,self.wcs,dtype=self.dtype)
return cs.alm2map(alm,omap,spin=spin)
def compute_output_map(self, alm):
lmax = hp.Alm.getlmax(alm.shape[-1]) # we assume mmax = lmax
if self.nside is None:
assert (self.shape is not None) and (self.wcs is not None)
n_comp = 3 if self.has_polarization else 1
output_map = enmap.empty((n_comp, ) + self.shape[-2:], self.wcs)
curvedsky.alm2map(alm, output_map, spin=[0, 2], verbose=True)
elif self.nside is not None:
if lmax > 3 * self.nside - 1:
clip = np.ones(3 * self.nside)
if alm.ndim == 1:
alm_clipped = hp.almxfl(alm, clip)
else:
alm_clipped = [hp.almxfl(each, clip) for each in alm]
else:
alm_clipped = alm
output_map = hp.alm2map(alm_clipped, self.nside)
else:
raise ValueError(
"You must specify either nside or both of shape and wcs")
return (output_map << self.input_units).to(
u.uK_CMB, equivalencies=self.equivalencies)
def alm2map_spin(self,alm,spin_alm,spin_transform,ncomp,mlmax):
"""
Returns
X(n) = sum_lm alm_s1 s2_Y_lm(n)
where s1=spin_alm and s2=spin_transform are different spins.
alm should be (alm_+|s|, alm_-|s|)
"""
# Transform (alm_+|s|, alm_-|s|) to (alm_+, alm_-)
ap_am = irot2d(alm,spin=spin_alm)
if self.hpix:
res = hp.alm2map_spin(ap_am,nside=self.nside,spin=abs(spin_transform),lmax=mlmax)
else:
omap = enmap.empty((ncomp,)+self.shape[-2:],self.wcs,dtype=self.dtype)
res = cs.alm2map(ap_am,omap,spin=abs(spin_transform))
# Rotate maps M+ and M- to M_+|s| and M_-|s|
# M_+|s| is the first component
# M_-|s| is the second component
return rot2d(res)
def get_input(input_name,set_idx,sim_idx,shape,wcs):
if input_name=='CMB':
cmb_type = 'LensedUnabberatedCMB'
signal_path = sints.dconfig['actsims']['signal_path']
cmb_file = os.path.join(signal_path, 'fullsky%s_alm_set%02d_%05d.fits' %(cmb_type, set_idx, sim_idx))
alms = hp.fitsfunc.read_alm(cmb_file, hdu = 1)
elif input_name=='tSZ':
ellmax = 8101
ells = np.arange(ellmax)
cyy = fgs.power_y(ells)[None,None]
cyy[0,0][ells<2] = 0
comptony_seed = seed_tracker.get_fg_seed(set_idx, sim_idx, 'comptony')
alms = curvedsky.rand_alm(cyy, seed = comptony_seed,lmax=ellmax)
elif input_name=='kappa':
signal_path = sints.dconfig['actsims']['signal_path']
k_file = os.path.join(signal_path, 'fullskyPhi_alm_%05d.fits' %(sim_idx))
palms = hp.fitsfunc.read_alm(k_file)
from pixell import lensing as plensing
alms = plensing.phi_to_kappa(palms)
omap = enmap.zeros((1,)+shape[-2:],wcs)
omap = curvedsky.alm2map(np.complex128(alms),omap,spin=0)[0]
return omap
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 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)
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 compute_map(self,oshape,owcs,qid,pixwin_taper_deg=0.3,pixwin_pad_deg=0.3,
include_cmb=True,include_tsz=True,include_fgres=True,sht_beam=True):
"""
1. get total alm
2. apply beam, and pixel window if Planck
3. ISHT
4. if ACT, apply a small taper and apply pixel window in Fourier space
"""
# pad to a slightly larger geometry
tot_pad_deg = pixwin_taper_deg + pixwin_pad_deg
res = maps.resolution(oshape,owcs)
pix = np.deg2rad(tot_pad_deg)/res
omap = enmap.pad(enmap.zeros((3,)+oshape,owcs),pix)
ishape,iwcs = omap.shape[-2:],omap.wcs
# get data model
dm = sints.models[sints.arrays(qid,'data_model')](region_shape=ishape,region_wcs=iwcs,calibrated=True)
# 1. get total alm
array_index = self.qids.index(qid)
tot_alm = int(include_cmb)*self.alms['cmb']
if include_tsz:
try:
assert self.tsz_fnu.ndim==2
tot_alm[0] = tot_alm[0] + hp.almxfl(self.alms['comptony'][0] ,self.tsz_fnu[array_index])
except:
tot_alm[0] = tot_alm[0] + self.alms['comptony'][0] * self.tsz_fnu[array_index]
if self.cfgres is not None: tot_alm[0] = tot_alm[0] + int(include_fgres)*self.alms['fgres'][array_index]
assert tot_alm.ndim==2
assert tot_alm.shape[0]==3
ells = np.arange(self.lmax+1)
# 2. get beam, and pixel window for Planck
if sht_beam:
beam = tutils.get_kbeam(qid,ells,sanitize=False,planck_pixwin=False) # NEVER SANITIZE THE BEAM IN A SIMULATION!!!
for i in range(3): tot_alm[i] = hp.almxfl(tot_alm[i],beam)
if dm.name=='planck_hybrid':
pixwint,pixwinp = hp.pixwin(nside=tutils.get_nside(qid),lmax=self.lmax,pol=True)
tot_alm[0] = hp.almxfl(tot_alm[0],pixwint)
tot_alm[1] = hp.almxfl(tot_alm[1],pixwinp)
tot_alm[2] = hp.almxfl(tot_alm[2],pixwinp)
# 3. ISHT
omap = curvedsky.alm2map(np.complex128(tot_alm),omap,spin=[0,2])
assert omap.ndim==3
assert omap.shape[0]==3
if not(sht_beam):
taper,_ = maps.get_taper_deg(ishape,iwcs,taper_width_degrees=pixwin_taper_deg,pad_width_degrees=pixwin_pad_deg)
modlmap = omap.modlmap()
beam = tutils.get_kbeam(qid,modlmap,sanitize=False,planck_pixwin=True)
kmap = enmap.fft(omap*taper,normalize='phys')
kmap = kmap * beam
# 4. if ACT, apply a small taper and apply pixel window in Fourier space
if dm.name=='act_mr3':
if sht_beam: taper,_ = maps.get_taper_deg(ishape,iwcs,taper_width_degrees=pixwin_taper_deg,pad_width_degrees=pixwin_pad_deg)
pwin = tutils.get_pixwin(ishape[-2:])
if sht_beam:
omap = maps.filter_map(omap*taper,pwin)
else:
kmap = kmap * pwin
if not(sht_beam): omap = enmap.ifft(kmap,normalize='phys').real
return enmap.extract(omap,(3,)+oshape[-2:],owcs)
def hdowngrade(imap, shape, wcs, lmax):
return (cs.alm2map(cs.map2alm(imap, lmax=lmax),
enmap.empty(shape, wcs, dtype=imap.dtype)))
pl.add(ls,nls,ls="--",label='theory noise per mode') pl._ax.set_xlim(20,4000) pl.done(config['data_path']+"xcls_%s.png" % polcomb) # Filtered input fikalm = hp.almxfl(ikalm,wfilt) fimap = hp.alm2map(fikalm,nside=256) # Resampled mask dmask = hp.ud_grade(mask,nside_out=256) dmask[dmask<0] = 0 # Mollview plots io.mollview(frmap*dmask,config['data_path']+"wrmap1.png",xsize=1600,lim=8e-6) io.mollview(fimap*dmask,config['data_path']+"wimap1.png",xsize=1600,lim=8e-6) # CAR plots shape,wcs = enmap.band_geometry(np.deg2rad((-70,30)),res=np.deg2rad(0.5*8192/512/60.)) omask = reproject.enmap_from_healpix(dmask, shape, wcs,rot=None) omap = cs.alm2map(fkalm, enmap.empty(shape,wcs)) io.hplot(omap*omask,config['data_path']+"cwrmap1",grid=True,ticks=20,color='gray') omap = cs.alm2map(fikalm, enmap.empty(shape,wcs)) io.hplot(omap*omask,config['data_path']+"cwimap1",grid=True,ticks=20,color='gray')
def _generate_gaussian_kappa(self, seed=None):
if seed is not None:
np.random.seed(seed_tracker.get_kappa_seed(seed))
clkk = self.clkk_spec[:, 1]
alm = curvedsky.rand_alm(clkk)
return curvedsky.alm2map(alm, self.template.copy())[np.newaxis, ...]
def getActpolCmbFgSim(beamfileDict,
shape, wcs,
iterationNum, cmbDir, freqs, psa,
cmbSet = 0, \
doBeam = True, applyWindow = True, verbose = True, cmbMaptype = 'LensedCMB',
foregroundSeed = 0, simType = 'cmb', foregroundPowerFile = None, applyModulation = True):
nTQUs = len('TQU')
firstTime = True
output = enmap.empty((
len(freqs),
nTQUs,
) + shape[-2:], wcs)
if simType == 'cmb':
filename = cmbDir + "/fullsky%s_alm_set%02d_%05d.fits" % (
cmbMaptype, cmbSet, iterationNum)
if verbose:
print('getActpolCmbFgSim(): loading CMB a_lms from %s' % filename)
import healpy
almTebFullskyOnecopy = np.complex128(
healpy.fitsfunc.read_alm(filename, hdu=(1, 2, 3)))
#Now tile the same for all the frequencies requested (i.e. CMB is same at all frequencies).
#The beam convolution happens below.
almTebFullsky = np.tile(almTebFullskyOnecopy, (len(freqs), 1, 1))
if verbose:
print('getActpolCmbFgSim(): done')
elif simType == 'foregrounds':
outputFreqs = ['f090', 'f150']
foregroundPowers \
= powspec.read_spectrum(os.path.join(os.path.dirname(os.path.abspath(__file__)),
'../data/',
foregroundPowerFile),
ncol = 3,
expand = 'row')
if verbose:
print('getActpolCmbFgSim(): getting foreground a_lms')
almTFullsky90and150 = curvedsky.rand_alm(foregroundPowers,
seed=foregroundSeed)
almTebFullsky = np.zeros((
len(freqs),
nTQUs,
) + (len(almTFullsky90and150[-1]), ),
dtype=np.complex128)
for fi, freq in enumerate(freqs):
if freq in outputFreqs:
almTebFullsky[fi, 'TQU'.index('T'), :] \
= almTFullsky90and150[outputFreqs.index(freq), :]
#Convolve with beam on full sky
for fi, freq in enumerate(freqs):
if doBeam:
beamFile = beamfileDict[psa + '_' + freq]
if verbose:
print('getActpolCmbFgSim(): applying beam from %s' % beamFile)
beamData = (np.loadtxt(beamFile))[:, 1]
else:
if verbose:
print('getActpolCmbFgSim(): not convolving with beam')
beamData = np.repeat(1., almTebFullsky.shape[-1])
import healpy
#couldn't quickly figure out how to vectorize this so loop from 0 to 2.
for tqui in range(nTQUs):
almTebFullsky[fi, tqui] = healpy.sphtfunc.almxfl(
almTebFullsky[fi, tqui].copy(), beamData)
#These lines stolen from curvedsky.rand_map
#doing all freqs at once gives error:
#sharp.pyx in sharp.execute_dp (cython/sharp.c:12118)()
#ValueError: ndarray is not C-contiguous
#so loop over all freqs once for now.
curvedsky.alm2map(almTebFullsky, output, spin=[0, 2], verbose=True)
# outputThisfreq = enmap.empty(( nTQUs,)+shape[-2:], wcs)
# curvedsky.alm2map(almTebFullsky[fi,:,:], outputThisfreq, spin = [0,2], verbose = True)
# output[fi,...] = outputThisfreq
# curvedsky.alm2map(almTebFullsky[fi,:,:], output[fi,:,:,:], spin = [0,2], verbose = True)
if applyModulation:
from pixell import aberration
for fi, freq in enumerate(freqs):
print('applying modulation for frequency %s' % freq)
output, A = aberration.boost_map(output,
aberration.dir_equ,
aberration.beta,
return_modulation=True,
aberrate=False,
freq=freqStrToValGhz[freq] * 1e9)
if applyWindow:
from pixell import fft
#The axes along which to FFT
axes = [-2, -1]
if verbose:
print('getActpolCmbFgSim(): applying pixel window function')
nfreq = len(freqs)
for idx in range(nfreq):
fd = fft.fft(output[idx], axes=axes)
wy, wx = enmap.calc_window(fd.shape)
twoDWindow = wy[:, None]**1 * wx[None, :]**1
#Careful, this is quietly multiplying an array with shape [N_freq, N_TQU, N_y, N_x] with one of shape [N_y, N_x]
fd *= twoDWindow
output[idx] = (fft.ifft(fd, axes=axes, normalize=True)).real
del fd
if verbose:
print('getActpolCmbFgSim(): done')
return enmap.ndmap(output, wcs)
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)
kappa_map_hp = healpy.read_map(p['websky_dir'] + p['kappa_map_name'])
kappa_alms = healpy.sphtfunc.map2alm(kappa_map_hp)
shape, wcs = enmap.fullsky_geometry(p['PIX_SIZE_ARCMIN']*utils.arcmin)
# phi_alms = healpy.sphtfunc.almxfl(kappa_alms, 1. / (ells * (ells + 1) / 2.))
# phi_map =
kappa_map_car = enmap.enmap(np.zeros(shape), wcs)
curvedsky.alm2map(kappa_alms, kappa_map_car)
lmax = healpy.Alm.getlmax(len(kappa_alms))
ainfo = sharp.alm_info(lmax)
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)
imap = f'{my_sim_path}/kap_lt4.5.fits'
kmap = reproject.enmap_from_healpix(imap,
fshape,
fwcs,
ncomp=1,
unit=1,
lmax=6000,
rot=None)
print(kmap.shape)
filename = f'{my_sim_path}/lensed_alm.fits'
alm = np.complex128(hp.read_alm(filename, hdu=(1, 2, 3)))
#print(alm.shape) (3, 34043626)
ncomp = 1
omap = enmap.empty((ncomp, ) + fshape[-2:], fwcs, dtype=np.float64)
lmap = curvedsky.alm2map(alm, omap, spin=[0, 2], oversample=2.0, method="auto")
print(lmap.shape)
tap_per = 12.0
pad_per = 3.0
fwhm = 1.4
nlevel = 20
xlmin = 200
xlmax = 2000
ylmin = 200
ylmax = 6000
ylcut = 2
klmin = 200
klmax = 5000
binner = stats.bin2D(modlmap, bin_edges)
cents = binner.centers
power = lambda x, y: binner.bin(np.real(x * y.conj()))[1]
norm = maps.interp(ls, Als['TE'])(modlmap)
fc = maps.FourierCalc(shape, wcs)
for task in my_tasks:
# Load sims
sindex = str(task).zfill(5)
alm = maps.change_alm_lmax(
hp.read_alm(sim_location +
"fullskyLensedUnabberatedCMB_alm_set00_%s.fits" % sindex,
hdu=(1, 2, 3)), mlmax).astype(dtype)
imap = cs.alm2map(alm, enmap.zeros((3, ) + shape[-2:], wcs), spin=0)
taper, w2 = maps.get_taper_deg(shape, wcs, args.tapwidth)
w3 = np.mean(taper**3.)
imap = imap * taper
_, kmap, _ = fc.power2d(imap, rot=False)
# Inputs
ikalm = lensing.phi_to_kappa(
maps.change_alm_lmax(
hp.read_alm(sim_location + "fullskyPhi_alm_%s.fits" % (sindex)),
mlmax))
ikmap = cs.alm2map(ikalm, enmap.zeros((1, ) + shape[-2:], wcs), spin=0)
ikmap = ikmap * taper
_, kikmap, _ = fc.power2d(ikmap)
