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()