def savenoisep2d(Nmu=4): for ia, aa in enumerate([0.3333, 0.2000, 0.1429]): zz = 1/aa-1 for iw, wopt in enumerate(['opt', 'pess']): #for iw, wopt in enumerate(['opt']): for it, thopt in enumerate(['opt', 'pess', 'reas']): #for it, thopt in enumerate([ 'reas']): if rank == 0: print(aa, wopt, thopt) dpath = '/global/cscratch1/sd/chmodi/m3127/21cm_cleaning/recon/fastpm_%0.4f/wedge_kmin%0.2f_%s/'%(aa, 0.03, wopt) dpath += 'L%04d-N%04d-R//thermal-%s-hex/ZA/opt_s999_h1massA_fourier_rsdpos/'%(bs, nc, thopt) ofolder = '../../data/ZArecon-rsd/kmin-003_wedge-%s/thermal-%s-hex/Nmu%d/'%(wopt, thopt, Nmu) try: os.makedirs(ofolder) except: pass fname = ofolder + 'pnoise-L%04d_%0.4f.txt'%(bs, aa) if args.pp : fname = fname[:-4] + '-up.txt' try: rep = np.loadtxt(fname).T rpfit = [{'k':rep[0], 'power':rep[i+1]} for i in range(3)] except: pass if args.pp: ivar = BigFileMesh(dpath+'/ivarmesh_up', 'ivar').paint() else: ivar = BigFileMesh(dpath+'/ivarmesh', 'ivar').paint() svar = (ivar.r2c()**-0.5).c2r() p0 = FFTPower(svar, mode='2d', Nmu=Nmu).power if rank == 0: np.savetxt(fname, p0['power'].real)
for aa in [0.1429, 0.2000, 0.3333]: #for aa in [0.1429]: zz = 1 / aa - 1 for wed in ['opt', 'pess']: for th in ['opt', 'pess', 'reas']: optfolder = recfolder + '/fastpm_%0.4f/wedge_kmin0.03_%s/L1024-N0256-R/thermal-%s-hex/ZA/opt_s999_h1massA_fourier_rsdpos/' % ( aa, wed, th) if nc == 256: dd = BigFileMesh(optfolder + 'datap/', 'mapp').paint() dn = BigFileMesh(optfolder + 'datan/', 'mapp').paint() else: dd = BigFileMesh(optfolder + 'datap_up/', 'mapp').paint() dn = BigFileMesh(optfolder + 'datan_up/', 'mapp').paint() dn -= dd if kk is None: kk = dd.r2c().x pd2 = FFTPower(dd, mode='2d', Nmu=Nmu).power k2, modes2, pd2 = pd2['k'], pd2['modes'], pd2['power'] pn2 = FFTPower(dn, mode='2d', Nmu=Nmu).power['power'] #Lost due to wedge wedang = mapnoise.wedge(1 / aa - 1, att=wed, angle=True) wedmu = mapnoise.wedge(1 / aa - 1, att=wed, angle=False) mask = getmask(kk, wedang, kmin=0.03) lostw = 1 - mask.sum() / mask.size #lost due to noise: noisedom = 0 k0 = 10 for ii in range(Nmu - 1):
assert conf['stages'].size == nsteps grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) kvec = tools.fftk((nc, nc, nc), bs, dtype=np.float32, symmetric=False) solver = Solver(pm, conf['cosmology']) conf['kvec'] = kvec conf['grid'] = grid #PM #whitec = pm.generate_whitenoise(seed, mode='complex', unitary=False) #lineark = whitec.apply(lambda k, v:Planck15.get_pklin(sum(ki ** 2 for ki in k)**0.5, 0) ** 0.5 * v / v.BoxSize.prod() ** 0.5) #linear = lineark.c2r() linear = BigFileMesh( '/project/projectdirs/astro250/chmodi/cosmo4d/data/z00/L0400_N0128_S0100_05step/mesh/', 's').paint() lineark = linear.r2c() state = solver.lpt(lineark, grid, conf['stages'][0], order=2) solver.nbody(state, leapfrog(conf['stages'])) final = pm.paint(state.X) if pm.comm.rank == 0: print('X, V computed') cat = ArrayCatalog({ 'Position': state.X, 'Velocity': state.V }, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) kdd = KDDensity(cat).density cat['KDDensity'] = kdd #Save
% (bs, nc), 'LinearDensityK').paint() tosave = [] for aa in [0.1429, 0.2000, 0.3333]: zz = 1 / aa - 1 print(aa) dyn = BigFileCatalog( '/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100/fastpm_%0.4f/1' % (bs, nc, aa)) hmesh = BigFileMesh( '/global/cscratch1/sd/chmodi/m3127/H1mass/highres/%d-9100/fastpm_%0.4f/HImesh-N%04d/' % (bs * 10, aa, nc), model).paint() fpos = dyn['Position'] dgrow = cosmo.scale_independent_growth_factor(zz) zapos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) ph = FFTPower(hmesh, mode='1d').power k, ph = ph['k'], ph['power'] ik = numpy.where(k > 0.3)[0][0] paramsza, zamod = getbias(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid, ik=ik) pmod = FFTPower(zamod, mode='1d').power['power'] tf = (pmod / ph)**0.5 if rank == 0: print(tf) if rank == 0: print(paramsza)
def make_bsum_plot(): """bias rcc, tf, error, bk""" ik = args.ik grid = pm.mesh_coordinates()*bs/nc noises = np.loadtxt('/global/u1/c/chmodi/Programs/21cm/21cm_cleaning/data/summaryHI.txt').T for i in range(noises[0].size): if noises[0][i] == np.round(1/aa-1, 2): noise = noises[3][i] print(noise) if args.ray: hmesh = BigFileMesh('/global/cscratch1/sd/chmodi/m3127/H1mass/highres/%d-9100/fastpm_%0.4f/HImesh-N%04d/'%(bs*10, aa, nc), 'ModelA').paint() lin = BigFileMesh('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100/linear'%(bs, nc), 'LinearDensityK').paint() dyn = BigFileCatalog('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100/fastpm_%0.4f/1'%(bs, nc, aa), comm=pm.comm, header='Header') else: hmesh = BigFileMesh('/global/cscratch1/sd/chmodi/m3127/H1mass/highres/%d-9100-fixed/fastpm_%0.4f/HImesh-N%04d/'%(bs*10, aa, nc), 'ModelA').paint() lin = BigFileMesh('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/linear'%(bs, nc), 'LinearDensityK').paint() dyn = BigFileCatalog('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/fastpm_%0.4f/1'%(bs, nc, aa), comm=pm.comm, header='Header') print(pm.comm.rank, dyn.Index.compute()) fpos = dyn['Position'] flay = pm.decompose(fpos) # dgrow = cosmo.scale_independent_growth_factor(zz) zapos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) print(zapos.shape, fpos.shape) paramsza, zamod = getbias(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid, ik=ik) _, pmmod = getbias(pm, basemesh=lin, hmesh=hmesh, pos=fpos, grid=grid, ik=ik) kk, paramszak, _ = getbiask(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid) # fin = pm.paint(fpos, layout=flay) fin /= fin.cmean() fin -= 1 finsm = ft.smooth(fin, 3, 'gauss') _, finmod = getbias(pm, basemesh=finsm, hmesh=hmesh, pos=grid, grid=grid, ik=ik) models = [zamod, pmmod, finmod, finsm.copy()] lbls = ['ZA shift', 'PM shift', 'Eulerian (R=3)', 'Eulerian (R=3), $b_1^E$'] lss = ['-', '-', '-', '--'] print('Setup done') ##### fig, ax = plt.subplots(2, 2, figsize=(8, 6), sharex=True) for ii, mod in enumerate(models): pmod = FFTPower(mod, mode='1d').power k, pmod = pmod['k'], pmod['power'] ph = FFTPower(hmesh, mode='1d').power['power'] if ii == 3: mod *= (ph[1]/pmod[1]).real**0.5 pmod *= (ph[1]/pmod[1]).real pxmodh = FFTPower(hmesh, second=mod, mode='1d').power['power'] perr = FFTPower(hmesh -mod, mode='1d').power['power'] ax[0, 0].plot(k, pxmodh/(pmod*ph)**0.5, label=lbls[ii], lw=2, ls=lss[ii]) ax[0, 0].set_ylabel('$r_{cc}$', fontdict=font) #if ii == 3: pmod *= ph[1]/pmod[1] ax[0, 1].plot(k,(pmod/ph)**0.5, lw=2, ls=lss[ii]) ax[0, 1].set_ylabel(r'$T_f$', fontdict=font) #ax[1].set_ylabel(r'$\sqrt{P_{\rm mod}/P_{hh}}$', fontdict=font) ax[1, 0].plot(k, perr, lw=2, ls=lss[ii]) ax[1, 0].set_ylabel(r'$P_{\rm err}$', fontdict=font) ax[1, 0].set_yscale('log') ax[1, 0].axhline(noise, color='k', ls="--") ax[1, 0].set_ylim(2, 5e2) lbls = ['$b_1$', '$b_2$', '$b_s$'] for ii in range(3): ax[1, 1].plot(kk, paramszak[ii], lw=2, color='C%d'%ii, label=lbls[ii]) ax[1, 1].axhline(paramsza[ii], lw=1, color='C%d'%ii, ls="--") ax[1, 1].set_ylabel(r'$b(k)$', fontdict=font) for axis in ax.flatten(): axis.set_xscale('log') axis.grid(which='both', alpha=0.2, color='gray', lw=0.2) axis.legend(prop=fontmanage) axis.set_xlabel(r'$k\quad [h\,{\rm Mpc}^{-1}]$', fontdict=font) axis.set_xscale('log') axis.axvline(kk[ik], color='k', ls="--") # Put on some more labels. for axis in ax.flatten(): axis.set_xscale('log') for tick in axis.xaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) for tick in axis.yaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) ##and finish plt.tight_layout(rect=[0, 0, 1, 0.95]) if args.ray: plt.savefig(figpath + '/bsum_L%04d_N%04d_%04d_ik%d.pdf'%(bs, nc, aa*10000, ik)) else: plt.savefig(figpath + '/bsumfix_L%04d_N%04d_%04d_ik%d.pdf'%(bs, nc, aa*10000, ik))
def make_biask_plot(): """Does the work of making the real-space xi(r) and b(r) figure.""" hmesh = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier/datap', 'mapp').paint() noises = np.loadtxt('/global/u1/c/chmodi/Programs/21cm/21cm_cleaning/data/summaryHI.txt').T for i in range(noises[0].size): if noises[0][i] == np.round(1/aa-1, 2): noise = noises[3][i] print(noise) zamod = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier/fitp/', 'mapp').paint() pmmod = BigFileMesh(dpath+'T05-B1/opt_s999_h1massA_fourier/fitp/', 'mapp').paint() fin = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier/datap/', 'd').paint() fin /= fin.cmean() fin -= 1 finsm = ft.smooth(fin, 3, 'gauss') grid = pm.mesh_coordinates()*bs/nc params, finmod = getbias(pm, basemesh=finsm, hmesh=hmesh, pos=grid, grid=grid) models = [zamod, pmmod, finmod, finsm.copy()] lbls = ['ZA shift', 'PM shift', 'Eulerian (R=3)', 'Eulerian (R=3), $b_1$'] lss = ['-', '-', '-', '--'] lin = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier/datap/', 's').paint() dyn = BigFileCatalog('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/fastpm_%0.4f/1'%(bs, nc, aa)) dgrow = cosmo.scale_independent_growth_factor(zz) zapos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) fpos = dyn['Position'] params, _ = getbias(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid) kk, paramsk, _ = getbiask(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid) fig, ax = plt.subplots(1, 2, figsize=(9, 4)) for ii, mod in enumerate(models): pmod = FFTPower(mod, mode='1d').power k, pmod = pmod['k'], pmod['power'] ph = FFTPower(hmesh, mode='1d').power['power'] if ii == 3: mod *= (ph[1]/pmod[1]).real **0.5 pmod *= (ph[1]/pmod[1]).real pxmodh = FFTPower(hmesh, second=mod, mode='1d').power['power'] perr = FFTPower(hmesh -mod, mode='1d').power['power'] ax[0].plot(k, perr, lw=2, ls=lss[ii], label=lbls[ii]) ax[0].set_ylabel(r'$P_{\rm err}$', fontdict=font) ax[0].set_yscale('log') ax[0].axhline(noise, color='k', ls="--") ax[0].set_ylim(2, 5e2) lbls = ['$b_1$', '$b_2$', '$b_s$'] for ii in range(3): ax[1].plot(kk, paramsk[ii], lw=2, color='C%d'%ii, label=lbls[ii]) ax[1].axhline(params[ii], lw=1, color='C%d'%ii, ls="--") ax[1].set_ylabel(r'$b(k)$', fontdict=font) for axis in ax: axis.set_xscale('log') axis.grid(which='both') axis.legend(prop=fontmanage, ncol=2) axis.set_xlabel(r'$k\quad [h\,{\rm Mpc}^{-1}]$', fontdict=font) axis.set_xscale('log') # Put on some more labels. for axis in ax: axis.set_xscale('log') for tick in axis.xaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) for tick in axis.yaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) ##and finish plt.tight_layout(rect=[0, 0, 1, 0.95]) plt.savefig(figpath + '/bfitk_L%04d_%04d.pdf'%(bs, aa*10000))
def make_rsdbias_plot(): """Does the work of making the real-space xi(r) and b(r) figure.""" hmesh = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier_rsdpos/datap', 'mapp').paint() noises = np.loadtxt('/global/u1/c/chmodi/Programs/21cm/21cm_cleaning/data/summaryHI.txt').T for i in range(noises[0].size): if noises[0][i] == np.round(1/aa-1, 2): noise = noises[3][i] print(noise) dgrow = cosmo.scale_independent_growth_factor(zz) rsdfac = 0 with open('/global/cscratch1/sd/chmodi/m3127/H1mass/highres/2560-9100-fixed/fastpm_%0.4f/Header/attr-v2'%aa) as ff: for line in ff.readlines(): if 'RSDFactor' in line: rsdfac = float(line.split()[-2]) print(rsdfac) los = np.array([0, 0, 1]).reshape(1, -1) grid = pm.mesh_coordinates()*bs/nc lin = BigFileMesh(dpath+'ZA/opt_s999_h1massA_fourier/datap/', 's').paint() dyn = BigFileCatalog('/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/fastpm_%0.4f/1'%(bs, nc, aa)) zapos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) zdisp = za.doza(lin.r2c(), grid, z=zz, displacement=True, dgrow=dgrow) zarsd = cosmo.scale_independent_growth_rate(zz)*zdisp*los pmrsd = dyn['Velocity']*rsdfac*los fpos = dyn['Position'] params, fitrsd = getbias(pm, basemesh=lin, hmesh=hmesh, pos=zapos+zarsd, grid=grid, fpos=zapos+zarsd) params, fitreal = getbias(pm, basemesh=lin, hmesh=hmesh, pos=zapos, grid=grid, fpos=zapos+zarsd) models = [fitrsd, fitreal] lbls = ['Fit Redshift', 'Fit Real'] lss = ['-', '-'] fig, ax = plt.subplots(1, 3, figsize=(15, 4)) for ii, mod in enumerate(models): pmod = FFTPower(mod, mode='1d').power k, pmod = pmod['k'], pmod['power'] ph = FFTPower(hmesh, mode='1d').power['power'] pxmodh = FFTPower(hmesh, second=mod, mode='1d').power['power'] perr = FFTPower(hmesh -mod, mode='1d').power['power'] ax[0].plot(k, pxmodh/(pmod*ph)**0.5, label=lbls[ii], lw=2, ls=lss[ii]) ax[0].set_ylabel('$r_{cc}$', fontdict=font) ax[1].plot(k,(pmod/ph)**0.5, lw=2, ls=lss[ii]) ax[1].set_ylabel(r'$\sqrt{P_{\rm mod}/P_{hh}}$', fontdict=font) ax[2].plot(k, perr, lw=2, ls=lss[ii]) ax[2].set_ylabel(r'$P_{\delta_{\rm mod}-\delta_h}$', fontdict=font) ax[2].set_yscale('log') ax[2].axhline(noise, color='k', ls="--") for axis in ax: axis.set_xscale('log') axis.grid(which='both') axis.legend(prop=fontmanage) axis.set_xlabel(r'$k\quad [h\,{\rm Mpc}^{-1}]$', fontdict=font) axis.set_xscale('log') # Put on some more labels. for axis in ax: axis.set_xscale('log') for tick in axis.xaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) for tick in axis.yaxis.get_major_ticks(): tick.label.set_fontproperties(fontmanage) ##and finish plt.tight_layout(rect=[0, 0, 1, 0.95]) plt.savefig(figpath + '/bfitrsd_L%04d_%04d.pdf'%(bs, aa*10000))
Rsms = [ 0, ] zadisps = [ True, ] for Rsm in Rsms: for zadisp in zadisps: # fpos = dyn['Position'].compute() dgrow = cosmo.scale_independent_growth_factor(zz) if zadisp: fpos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) dlay = pm.decompose(fpos) disp = grid - fpos mask = abs(disp) > bs / 2. disp[mask] = (bs - abs(disp[mask])) * -np.sign(disp[mask]) print(rank, ' Max disp: ', disp.max()) print(rank, ' Std disp: ', disp.std(axis=0)) # hpos = hcat['Position'] print('Mass : ', rank, hcat['Length'][-1].compute() * hcat.attrs['M0'] * 1e10) hlay = pm.decompose(hpos) hmesh = pm.paint(hpos, layout=hlay) hmesh /= hmesh.cmean()