Python WLanalysis类代码示例

def OrganizeSplitFile(ifile):
	'''read in one of the split file, pick out the redshift, and sort by fields, e2 is c2 correted, with e2-=c2'''	
	
	field = genfromtxt(split_dir+ifile,usecols=0,dtype=str)
	field = array(map(field2int,field))
	print ifile

	# generate 2 random redshift and 1 peak
	Pz = genfromtxt(split_dir+ifile,usecols=arange(14,84),dtype=str)
	Pz = (np.core.defchararray.replace(Pz,',','')).astype(float)
	
	seed(99)
	z_rand1 = array(map(DrawRedshifts,Pz)).ravel()
	seed(88)
	z_rand2 = array(map(DrawRedshifts,Pz)).ravel()
	z_peak = z_arr[argmax(Pz,axis=1)]
	z_all = concatenate([[z_peak,], [z_rand1,], [z_rand2,]]).T

	sheardata = genfromtxt(split_dir+ifile,usecols=[1,2,5,6,7,8,9,10,11,12,13,84])
	ra, dec, e1, e2, w, fitclass, r, snr, mask, m, c2, mag = sheardata.T
	e2 -= c2
	
	
	i=0
	for Wx in range(1,5):
		idx=where((field==Wx)&(mask<=1.0)&(fitclass==0)&(amin(z_all,axis=-1)>=0.2)&(amax(z_all,axis=-1)<=1.3))[0]
		print ifile, Wx, len(idx)/50000.0
		if len(idx) > 0:
			#data = (np.array([ra,dec,e1,e2,w,r,snr,m,c2,mag]).T)[idx]
			data = (np.array([ra,dec,e1,e2,w,r,snr,m,c2,mag,z_peak, z_rand1, z_rand2]).T)[idx]
			radeclist = sheardata[idx][:,[0,1]]
			xylist = list2coords(radeclist, Wx)
			xy_data = concatenate([xylist,data],axis=1)
			WLanalysis.writeFits(xy_data, W_dir(Wx)+ifile+'.fit')#,fmt=['%i','%i','%s','%s','%s','%.3f'])
		i+=1

def TestCrossCorrelate (Wx, zcut, sigmaG):
	'''Input: 
	Wx - one of the W1..W4 field (= 1..4) 
	zcut - redshift cut between KS background galaxies and forground cluster probe
	sigmaG - smoothing
	Output:
	ell_arr, CCK, CCB
	'''
	galn_hi = WLanalysis.readFits(test_dir+'W%i_galn_%s_hi_sigmaG%02d.fit'%(Wx,zcut,sigmaG*10))
	galn_lo = WLanalysis.readFits(test_dir+'W%i_galn_%s_lo_sigmaG%02d.fit'%(Wx,zcut,sigmaG*10))
	galn_cut = 0.5*0.164794921875 #5gal/arcmin^2*arcmin^2/pix, arcmin/pix = 12.0*60**2/512.0**2 = 
	bmap = WLanalysis.readFits(test_dir+'W%i_Bmode_%s_hi_sigmaG%02d.fit'%(Wx,zcut,sigmaG*10))
	kmap = WLanalysis.readFits(test_dir+'W%i_KS_%s_hi_sigmaG%02d.fit'%(Wx,zcut,sigmaG*10))
	mask = where(galn_hi<galn_cut)
	bmap[mask]=0
	kmap[mask]=0
	edges=linspace(5,100,11)
	ell_arr, CCB = WLanalysis.CrossCorrelate (bmap,galn_lo,edges=edges)
	ell_arr, CCK = WLanalysis.CrossCorrelate (kmap,galn_lo,edges=edges)
	f=figure(figsize=(8,6))
	ax=f.add_subplot(111)
	ax.plot(ell_arr, CCB, 'ro',label='B-mode')
	ax.plot(ell_arr, CCK, 'bo', label='KS')
	legend()
	#ax.set_xscale('log')
	ax.set_xlabel('ell')
	ax.set_ylabel(r'$\ell(\ell+1)P_{n\kappa}(\ell)/2\pi$')
	ax.set_title('W%i_zcut%shi_sigmaG%02d'%(Wx,zcut,sigmaG*10))
	#show()
	savefig(plot_dir+'CC_edges_W%i_zcut%shi_sigmaG%02d.jpg'%(Wx,zcut,sigmaG*10))
	close()

def PeakPos (Wx, z_lo=0.6, z_hi='0.6_lo',noise=False, Bmode=False):
	'''For a map(kappa or bmode), find peaks, and its(RA, DEC)
	return 3 columns: [kappa, RA, DEC]
	'''
	#print 'noise', noise, Wx
	if Bmode:
		kmap = bmodeGen(Wx, z=z_hi)
	else:
		kmap = kmapGen(Wx, z=z_hi)
	ipeak_mat = WLanalysis.peaks_mat(kmap)
	imask = maskGen (Wx, z=z_lo)
	ipeak_mat[where(imask==0)]=nan #get ipeak_mat, masked region = nan
	if noise: #find the index for peaks in noise map
		idx_all = where((imask==1)&isnan(ipeak_mat))
		sample = randint(0,len(idx_all[0])-1,sum(~isnan(ipeak_mat)))
		idx = array([idx_all[0][sample],idx_all[1][sample]])
	else:#find the index for peaks in kappa map
		idx = where(~isnan(ipeak_mat)==True)
	kappaPos_arr = zeros(shape=(len(idx[0]),3))#prepare array for output
	for i in range(len(idx[0])):
		x, y = idx[0][i], idx[1][i]#x, y
		kappaPos_arr[i,0] = kmap[x, y]
		x = int(x-sizes[Wx-1]/2)+1
		y = int(y-sizes[Wx-1]/2)+1
		x /= PPR512# convert from pixel to radians
		y /= PPR512
		kappaPos_arr[i,1:] = WLanalysis.gnom_inv((y, x), centers[Wx-1])
	return kappaPos_arr.T

def compute_GRF_PDF_ps_pk (r):
	'''for a convergence map with filename fn, compute the PDF and the power spectrum. sizedeg = 3.5**2, or 1.7**2'''
	print cosmo, r
	kmap = kmapGen(r)
	#kmap = load(CMBlensing_dir+'GRF_fidu/'+'GRF_fidu_%04dr.npy'%(r))
	
	i_arr = arange(len(sigmaP_arr))
	
	if not doGRF:
            kmap_smoothed = [WLanalysis.smooth(kmap, sigmaP) for sigmaP in sigmaP_arr]
            ps = WLanalysis.PowerSpectrum(kmap_smoothed[0])[1]

            PDF = [PDFGen(kmap_smoothed[i], PDFbin_arr[i]) for i in i_arr]
            peaks = [peaksGen(kmap_smoothed[i], peak_bins_arr[i]) for i in i_arr]

	###### generate GRF
	else:
            ps=0
            random.seed(r)
            GRF = (WLanalysis.GRF_Gen(kmap)).newGRF()
            #save(CMBlensing_dir+'GRF_fidu/'+'GRF_fidu_%04dr.npy'%(r), GRF)		
            #GRF = load(CMBlensing_dir+'GRF_fidu/'+'GRF_fidu_%04dr.npy'%(r))
            GRF_smoothed = [WLanalysis.smooth(GRF, sigmaP) for sigmaP in sigmaP_arr]
            PDF = [PDFGen(GRF_smoothed[i], PDFbin_arr[i]) for i in i_arr]
            peaks = [peaksGen(GRF_smoothed[i], peak_bins_arr[i]) for i in i_arr]
	#############

	return [ps,], PDF, peaks#, PDF_GRF, peaks_GRF

def partialdata2grid (icount):
    '''for a small portion of the data, put into a grid'''
    print 'icount',icount
    idata = load(mask_dir+'smaller/weight0_cat_W%i_step%i_start%i.npy'%(Wx,step, icount))
    print 'loaded',icount
    ix, iy=np.indices(idata.shape)
    iy+=step*icount
    #radeclist = (array(w.wcs_pix2world(ix, iy, 0)).reshape(2,-1)).T ////jia changed on 12/9, since the coordinates seems to be off..
    #radeclist = (array(w.wcs_pix2world(iy, ix, 0)).reshape(2,-1)).T 
    #y, x = f_Wx (radeclist)
    y, x = f_Wx ((array(w.wcs_pix2world(iy, ix, 0)).reshape(2,-1)).T )
    print icount,'done f_wx %s'%(icount)#,time.strftime("%Y-%m-%d %H:%M")

    if Wx!=1:
        ipix_mask, ipix = WLanalysis.coords2grid(x,y, idata.flatten().reshape(1,-1), size=sizes[Wx-1])
        
    else:
        istep=ceil(len(y)/10.0)
        ipix_mask = zeros(shape=(sizes[Wx-1], sizes[Wx-1]))
        ipix = ipix_mask.copy()
        jj=0    
        while jj< len(y):
            print 'icount, jj',icount,jj
            iipix_mask,iipix = WLanalysis.coords2grid(x[jj:jj+istep], y[jj:jj+istep], idata.flatten().reshape(1,-1)[:,jj:jj+istep], size=sizes[Wx-1])
            ipix_mask += iipix_mask
            ipix += iipix
            jj+=istep
    print icount,'W%i done coords2grid %s'%(Wx,icount)#,time.strftime("%Y-%m-%d %H:%M")    
    
    save(mask_dir+'smaller/weight0_W%i_%i_numpix'%(Wx,icount), ipix)
    save(mask_dir+'smaller/weight0_W%i_%i_nummask'%(Wx,icount), ipix_mask)
    #ipix is the num. of pixels fall in that big pix, ipix_mask is the mask
    return ipix, ipix_mask

	def randmap (iseed, Wx=Wx):	
		Me1rnd, Me2rnd = WLanalysis.rndrot(Me1, Me2, iseed=iseed)
		Me1smooth = WLanalysis.weighted_smooth(Me1rnd, Mwm)
		Me2smooth = WLanalysis.weighted_smooth(Me2rnd, Mwm)
		kmap_rand = WLanalysis.KSvw(Me1smooth, Me2smooth)
		print Wx, iseed, kmap_rand.shape
		np.save(bmap_fn(Wx, iseed), kmap_rand)

def iskew (i):
    print i
    ikmap_NL = kmapNL(i)
    ikmap_NOISY = kmapNOISY(i)
    skewness_NL = [skew(WLanalysis.smooth(ikmap_NL, ismooth).flatten() ) for ismooth in sigmaG_arr*PPA_NL] 
    skewness_NOISY = [skew(WLanalysis.smooth(ikmap_NOISY, ismooth).flatten() ) for ismooth in sigmaG_arr*PPA_NOISY]
    return [skewness_NL, skewness_NOISY]

def KSmap_massproduce(iiRcosmo):
	'''Input:
	iiRcosmo = [i, R, cosmo]
	i: subfield range from (1, 2..13)
	R: realization range from (1..1000)
	cosmo: one of the 1000 cosmos
	Return:
	KS inverted map
	'''
	i, R, cosmo = iiRcosmo
	create_kmap = 0
	for sigmaG in sigmaG_arr:
		KS_fn = KSfn(i, cosmo, R, sigmaG)
		if not os.path.isfile(KS_fn):
			create_kmap = 1
			break
	if create_kmap:
		print 'Mass Produce, creating KS: ', i, R, cosmo
		Me1, Me2 = fileGen(i, R, cosmo)
		for sigmaG in sigmaG_arr:
			KS_fn = KSfn(i, cosmo, R, sigmaG)	
			Me1_smooth = WLanalysis.weighted_smooth(Me1, Mw, PPA=PPA512, sigmaG=sigmaG)
			Me2_smooth = WLanalysis.weighted_smooth(Me2, Mw, PPA=PPA512, sigmaG=sigmaG)
			kmap = WLanalysis.KSvw(Me1_smooth, Me2_smooth)
			np.save(KS_fn, kmap)

def create_prob_plane(psPDFpk='pk', sigmaG_idx=0):
    if psPDFpk =='ps':
        idx2000=where(ell_gadget<10000)[0]
        obs_arr = load(CMBNG_dir+'mat/mat_ps_avg.npy')[:,idx2000]
        fidu_mat = load(CMBNG_dir+'mat/mat_ps_fidu.npy')[:,idx2000]
    else:
        obs_arr = load(CMBNG_dir+'mat/mat_%s_sigmaG%i_avg.npy'%(psPDFpk, sigmaG_idx))
        
        fidu_mat = load(CMBNG_dir+'mat/mat_%s_sigmaG%i_fidu.npy'%(psPDFpk, sigmaG_idx))
    
    idx = where(~isnan(mean(fidu_mat,axis=0))&(mean(fidu_mat,axis=0)!=0))[0]
    fidu_mat=fidu_mat[:,idx]
    interp_cosmo=WLanalysis.buildInterpolator2D(obs_arr[:,idx], cosmo_params)
    
    #cov_mat = 
    cov_mat = cov(fidu_mat,rowvar=0)/(2e4/12.5)
    cov_inv = mat(cov_mat).I
    
    def chisq_fcn(param1, param2):
        model = interp_cosmo((param1,param2))
        del_N = np.mat(model - mean(fidu_mat,axis=0))
        chisq = float(del_N*cov_inv*del_N.T)
        return chisq
        
    prob_plane = WLanalysis.prob_plane(chisq_fcn, om_arr, si8_arr)
    return prob_plane[1]

def compute_GRF_PDF_ps_pk (cosmo, r, Gaus=0,sigmaG=8.0):
    kmap = FT2real(cosmo, r, Gaus=Gaus)
    ps = WLanalysis.PowerSpectrum(WLanalysis.smooth(kmap, 0.18), bins=bins)[1]#*2.0*pi/ell_arr**2
    if not filtered:
        kmap = WLanalysis.smooth(kmap, 2.93*sigmaG/8.0)
    PDF = PDFGen(kmap, PDFbins)
    peaks = peaksGen(kmap, peak_bins)
    return concatenate([ps, PDF, peaks])

	def ips_pk_single (R):#, sigmaG, zg, bins):
		kmap = WLanalysis.readFits(KSsim_fn(i, cosmo, R, sigmaG, zg))
		if pk:#peaks
			mask = WLanalysis.readFits(Mask_fn(i, sigmaG))
			peaks_hist = WLanalysis.peaks_mask_hist(kmap, mask, bins, kmin=kmin, kmax=kmax)
			return peaks_hist
		else:#powspec
			ell_arr, powspec = WLanalysis.PowerSpectrum(kmap, sizedeg=12.0)
			return powspec

	def randmap (iseedWx):
		iseed, Wx = iseedWx
		Me1, Me2 = Me1_arr[Wx-1], Me2_arr[Wx-1]
		Mwm = Mwm_arr[Wx-1]
		Me1rnd, Me2rnd = WLanalysis.rndrot(Me1, Me2, iseed=iseed)
		Me1smooth = WLanalysis.weighted_smooth(Me1rnd, Mwm)
		Me2smooth = WLanalysis.weighted_smooth(Me2rnd, Mwm)
		kmap_rand = WLanalysis.KSvw(Me1smooth, Me2smooth)
		print Wx, iseed, kmap_rand.shape
		np.save(bmap_fn(Wx, iseed), kmap_rand)

def kmapPk_1sim (r):
	print r,'1sim'
	k, s1, s2 = kappaGen_1sim(r)[:3]
	A, galn = WLanalysis.coords2grid(x, y, array([s1,s2]))
	Ms1,Ms2 = A
	s1_smooth = WLanalysis.weighted_smooth(Ms1, galn, PPA=PPA512, sigmaG=sigmaG)
	s2_smooth = WLanalysis.weighted_smooth(Ms1, galn, PPA=PPA512, sigmaG=sigmaG)
	kmap = WLanalysis.KSvw(s1_smooth, s2_smooth)
	pk = WLanalysis.peaks_mask_hist(kmap, mask, bins=25, kmin = -0.04, kmax = 0.12)
	return pk

def kmapPs (r):
	print r
	k, s1, s2 = kappaGen(r)[:3]
	A, galn = WLanalysis.coords2grid(x, y, array([s1,s2 ]))
	Ms1,Ms2 = A
	s1_smooth = WLanalysis.weighted_smooth(Ms1, galn, PPA=PPA512, sigmaG=sigmaG)
	s2_smooth = WLanalysis.weighted_smooth(Ms1, galn, PPA=PPA512, sigmaG=sigmaG)
	kmap = WLanalysis.KSvw(s1_smooth, s2_smooth)
	kmap*=mask
	ps = WLanalysis.PowerSpectrum(kmap,sizedeg=12.0)[-1]
	return ps

def FT_PowerSpectrum (cosmo, r, bins=10, return_ell_arr=0, Gaus=0):
    if Gaus:
        a = FTmapGen_Gaus(r)
    else:
        a = FTmapGen(cosmo, r)
    PS2D=np.abs(fftpack.fftshift(a))**2
    ell_arr,psd1D=WLanalysis.azimuthalAverage(PS2D, bins=bins)    
    if return_ell_arr:
        ell_arr = WLanalysis.edge2center(ell_arr)* 360./3.5
        return ell_arr
    else:
        return psd1D