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)