def getPow(data1, data2=None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den1 = pm.paint(q + data1.reshape([-1, 3]))
if (data2 is not None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
den2 = pm.paint(q + data2.reshape([-1, 3]))
temp = FFTPower(first=den1, second=den2, mode='1d', BoxSize=128, dk=dk)
k, powspec = temp.power['k'], temp.power['power']
else:
temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk)
k, powspec = temp.power['k'], temp.power['power']
return [k, powspec.real]
def test_c2c(comm):
# this test requires pfft-python 0.1.16.
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128')
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
real2 = complex.c2r()
assert numpy.iscomplexobj(real)
assert numpy.iscomplexobj(real2)
assert numpy.iscomplexobj(complex)
assert_array_equal(complex.cshape, pm.Nmesh)
assert_array_equal(real2.cshape, pm.Nmesh)
assert_array_equal(real.cshape, pm.Nmesh)
real.readout(npos)
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def fiddlebias(aa, saveb=False, bname='h1bias.txt', ofolder=None):
#Fiddling bias for halos. Deprecated.
print('Read in catalogs')
halocat = readincatalog(aa, matter=False)
hpos = halocat['Position']
hmass = halocat['Mass']
h1mass = dohod.HI_mass(hmass, aa)
dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint()
if ofolder is None: ofolder = project + '/%s/fastpm_%0.4f/'%(sim, aa)
#measure power
pm = ParticleMesh(BoxSize = bs, Nmesh = [nc, nc, nc])
pkm = FFTPower(dm/dm.cmean(), mode='1d').power
k, pkm = pkm['k'], pkm['power']
##H1
print("H1")
h1mesh = pm.paint(hpos, mass=h1mass)
pkh1 = FFTPower(h1mesh/h1mesh.cmean(), mode='1d').power['power']
pkh1m = FFTPower(h1mesh/h1mesh.cmean(), second=dm/dm.cmean(), mode='1d').power['power']
#Bias
b1h1 = pkh1m/pkm
b1h1sq = pkh1/pkm
if saveb:
np.savetxt(ofolder+bname, np.stack((k, b1h1, b1h1sq**0.5), axis=1),
header='k, pkh1xm/pkm, pkh1/pkm^0.5')
return k, b1h1, b1h1sq
def test_nody():
""" Checking end to end nbody
"""
a0 = 0.1
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32)
solver = Solver(pm, Planck15, B=1)
stages = np.linspace(0.1, 1.0, 10, endpoint=True)
# Generate initial state with fastpm
whitec = pm.generate_whitenoise(100, 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)
statelpt = solver.lpt(lineark, grid, a0, order=1)
finalstate = solver.nbody(statelpt, leapfrog(stages))
final_cube = pm.paint(finalstate.X)
# Same thing with flowpm
tlinear = tf.expand_dims(np.array(lineark.c2r()), 0)
state = tfpm.lpt_init(tlinear, a0, order=1)
state = tfpm.nbody(state, stages, nc)
tfread = pmutils.cic_paint(tf.zeros_like(tlinear), state[0]).numpy()
assert_allclose(final_cube, tfread[0], atol=1.2)
def test_inplace_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(
-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
complex2 = real.r2c(out=Ellipsis)
assert real._base in complex2._base
assert_almost_equal(numpy.asarray(complex),
numpy.asarray(complex2),
decimal=7)
real = complex2.c2r()
real2 = complex2.c2r(out=Ellipsis)
assert real2._base in complex2._base
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_c2c(comm):
# this test requires pfft-python 0.1.16.
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128')
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(
-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
real2 = complex.c2r()
assert numpy.iscomplexobj(real)
assert numpy.iscomplexobj(real2)
assert numpy.iscomplexobj(complex)
assert_array_equal(complex.cshape, pm.Nmesh)
assert_array_equal(real2.cshape, pm.Nmesh)
assert_array_equal(real.cshape, pm.Nmesh)
real.readout(npos)
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_inplace_fft(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8')
real = RealField(pm)
numpy.random.seed(1234)
if comm.rank == 0:
Npar = 100
else:
Npar = 0
pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7)
pos %= (pm.Nmesh + 1)
layout = pm.decompose(pos)
npos = layout.exchange(pos)
real = pm.paint(npos)
complex = real.r2c()
complex2 = real.r2c(out=Ellipsis)
assert real._base in complex2._base
assert_almost_equal(numpy.asarray(complex), numpy.asarray(complex2), decimal=7)
real = complex2.c2r()
real2 = complex2.c2r(out=Ellipsis)
assert real2._base in complex2._base
assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def make_galaxy_pk(scale_factor,nc,seed,bs=1536,T=40,B=2,simpath=sc_simpath,outpath=sc_outpath,Rsm=0):
aa = scale_factor # since particle IDs are ordered only need to load at one redshift
zz = 1/aa-1
dgrow = cosmo.scale_independent_growth_factor(zz)
fname = get_filename(nc,seed,T=T,B=B)
spath = simpath + fname
opath = outpath + fname
galpath = opath + '/hod/'
try: os.makedirs(opath+'spectra/')
except : pass
zz = 1/aa-1
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8')
rank = pm.comm.rank
dyn = BigFileCatalog(spath + '/fastpm_%0.4f/1'%aa)
# Load Matter Mesh
fpos = dyn['Position'].compute()
mlay = pm.decompose(fpos)
mmesh = pm.paint(fpos, layout=mlay)
mmesh /= mmesh.cmean()
# Load halo mesh: HOD parameters fixed for now...
mmin = 10**12.5; m1fac = 20
cendir ='cencat-aa-%.04f-Mmin-%.1f-M1f-%.1f-alpha-0p8-subvol'%(aa,np.log10(mmin), m1fac)
satdir ='satcat-aa-%.04f-Mmin-%.1f-M1f-%.1f-alpha-0p8-subvol'%(aa,np.log10(mmin), m1fac)
cencat = BigFileCatalog(galpath + cendir)
satcat = BigFileCatalog(galpath + satdir)
hpos = np.concatenate((cencat['Position'],satcat['Position']))
hlay = pm.decompose(hpos)
hmesh = pm.paint(hpos, layout=hlay)
hmesh /= hmesh.cmean()
phm = FFTPower(mmesh, second=hmesh, mode='1d').power
k, phm = phm['k'], phm['power'].real
phh = FFTPower(hmesh, mode='1d').power
k, phh = phh['k'], phh['power'].real
np.savetxt(opath+'spectra/pks-%04d-%04d-%04d-gal.txt'%(aa*10000, bs, nc), np.vstack([k, phh, phm]).T.real, header='k, phh, phm/ ', fmt='%0.4e')
def test_grid_shifted(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
grid = grid + 4.0
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
g2 = layout.exchange(grid)
real = pm.paint(grid, layout=layout)
#print(real, g2 % 8, real.slices)
assert_allclose(real, 1.0)
grid = grid - 6.1
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
real = pm.paint(grid, layout=layout)
assert_allclose(real, 1.0)
def test_grid_shifted(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
grid = grid + 4.0
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
g2 = layout.exchange(grid)
real = pm.paint(grid, layout=layout)
#print(real, g2 % 8, real.slices)
assert_allclose(real, 1.0)
grid = grid - 6.1
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
layout = pm.decompose(grid)
real = pm.paint(grid, layout=layout)
assert_allclose(real, 1.0)
def savecatalogmesh(bs, nc, aa):
dmcat = BigFileCatalog(scratchyf + sim + '/fastpm_%0.4f/' % aa,
dataset='1')
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
pos = dmcat['Position'].compute()
layout = pm.decompose(pos)
mesh = pm.paint(pos, layout=layout)
mesh = FieldMesh(mesh)
path = project + sim + '/fastpm_%0.4f/' % aa + '/dmesh_N%04d' % nc
mesh.save(path, dataset='1', mode='real')
def getPow_dis(data1, data2=None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den1 = pm.paint(q)
power = 0
if (data2 is not None):
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
den2 = pm.paint(q)
for ii in range(3):
den1[:] = data1[:, :, :, ii]
den2[:] = data2[:, :, :, ii]
temp = FFTPower(first=den1,
second=den2,
mode='1d',
BoxSize=128,
dk=dk)
k, power = temp.power['k'], power + temp.power['power']
else:
for ii in range(3):
den1[:] = data1[:, :, :, ii]
temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk)
k, power = temp.power['k'], power + temp.power['power']
return [k, power.real]
def test_cic_paint():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
wts = np.random.random(nparticle).astype(np.float32)
# Painting with pmesg
pmmesh = pm.paint(pos, mass=wts)
mesh = cic_paint(tf.zeros((1, nc, nc, nc), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)),
weight=wts.reshape(1, nparticle))
tfmesh = mesh.numpy()
assert_allclose(pmmesh, tfmesh[0], atol=1e-06)
def test_cic_paint():
bs = 50
nc = 16
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
nparticle = 100
pos = bs * np.random.random(3 * nparticle).reshape(-1, 3).astype(
np.float32)
wts = np.random.random(nparticle).astype(np.float32)
# Painting with pmesg
pmmesh = pm.paint(pos, mass=wts)
with tf.Session() as sess:
mesh = cic_paint(tf.zeros((1, nc, nc, nc), dtype=tf.float32),
(pos * nc / bs).reshape((1, nparticle, 3)),
weight=wts.reshape(1, nparticle))
sess.run(tf.global_variables_initializer())
tfmesh = sess.run(mesh)
assert_allclose(pmmesh, tfmesh[0], atol=1e-06)
def fiddlebiasgal(aa, suff, nc=nc, mcfv=[1.], saveb=False, bname='h1bias', ofolder=None):
'''Fiddle bias for galaxies'''
if ofolder is None: ofolder = project + '/%s/fastpm_%0.4f/'%(sim, aa)
pm = ParticleMesh(BoxSize = bs, Nmesh = [nc, nc, nc])
print('Read in catalogs')
cencat = BigFileCatalog(project + sim + '/fastpm_%0.4f/cencat'%aa)
satcat = BigFileCatalog(project + sim + '/fastpm_%0.4f/satcat'%aa+suff)
cpos, spos = cencat['Position'], satcat['Position']
cmass, smass = cencat['Mass'], satcat['Mass']
pos = np.concatenate((cpos, spos), axis=0)
dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint()
pkm = FFTPower(dm/dm.cmean(), mode='1d').power
k, pkm = pkm['k'], pkm['power']
b1, b1sq = np.zeros((k.size, len(mcfv))), np.zeros((k.size, len(mcfv)))
for imc, mcf in enumerate(mcfv):
print(mcf)
ch1mass = HI_masscutfiddle(cmass, aa, mcutf=mcf)
sh1mass = HI_masscutfiddle(smass, aa, mcutf=mcf)
h1mass = np.concatenate((ch1mass, sh1mass), axis=0)
#
h1mesh = pm.paint(pos, mass=h1mass)
pkh1 = FFTPower(h1mesh/h1mesh.cmean(), mode='1d').power['power']
pkh1m = FFTPower(h1mesh/h1mesh.cmean(), second=dm/dm.cmean(), mode='1d').power['power']
#Bias
b1[:, imc] = pkh1m/pkm
b1sq[:, imc] = pkh1/pkm
np.savetxt(ofolder+bname+'auto'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1sq**0.5), axis=1),
header='mcut factors = %s\nk, pkh1xm/pkm, pkh1/pkm^0.5'%mcfv)
np.savetxt(ofolder+bname+'cross'+suff+'.txt', np.concatenate((k.reshape(-1, 1), b1), axis=1),
header='mcut factors = %s\nk, pkh1xm/pkm, pkh1mx/pkm'%mcfv)
return k, b1, b1sq
cencat['HImass'] = HI_hod(cencat['Mass'], aa)
satcat['HImass'] = HI_hod(satcat['Mass'], aa)
### totHImass = cencat['HImass'].sum().compute() +\
### satcat['HImass'].sum().compute()
### cencat['HImass']/= totHImass/float(nc)**3
### satcat['HImass']/= totHImass/float(nc)**3
### allcat = MultipleSpeciesCatalog(['cen','sat'],cencat,satcat)
###
### h1mesh = allcat.to_mesh(BoxSize=bs,Nmesh=[nc,nc,nc],\
### position='RSDpos',weight='HImass')
###
## cenmesh = pm.paint(cencat['RSDpos'], mass=cencat['HImass'])
## satmesh = pm.paint(satcat['RSDpos'], mass=satcat['HImass'])
## h1mesh = cenmesh + satmesh
## h1mesh /= h1mesh.cmean()
## end = time()
## if rank ==0: print('Time to create mesh : ', end-start)
## calc_pk1d(aa,h1mesh)
## calc_pkmu(aa,h1mesh)
## calc_pkll(aa,h1mesh)
## #
cenmesh = pm.paint(cencat['Position'], mass=cencat['HImass'])
satmesh = pm.paint(satcat['Position'], mass=satcat['HImass'])
h1mesh = cenmesh + satmesh
h1mesh /= h1mesh.cmean()
end = time()
if rank == 0: print('Time to create mesh : ', end - start)
calc_pkreal(aa, h1mesh)
#
print(rank, 'lag field created')
for Rsm in [0]:
for zadisp in [False, True]:
#
#fpos = dyn['Position']
#if zadisp : fpos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow)
glay, play = pm.decompose(grid), pm.decompose(fpos)
for i, ff in enumerate(lag_fields):
print(rank, i)
x = FieldMesh(pm.paint(fpos, mass=ff.readout(grid, layout = glay, resampler='nearest'), layout=play))
if zadisp : x.save(lpath + 'za-z%03d-R%d'%(zz*100, Rsm), dataset=names[i], mode='real')
else : x.save(lpath + 'eul-z%03d-R%d'%(zz*100, Rsm), dataset=names[i], mode='real')
del x
## k, spectra = tools.getspectra(eul_fields)
## #k, spectra = tools.getspectra(lag_fields)
##
## if zadisp: np.savetxt(ofolder + '/spectraza-z%03d-R%d.txt'%(zz*100, Rsm), np.vstack([k, spectra]).T.real, header='k / '+header, fmt='%0.4e')
## else: np.savetxt(ofolder + '/spectra-z%03d-R%d.txt'%(zz*100, Rsm), np.vstack([k, spectra]).T.real, header='k / '+header, fmt='%0.4e')
##
##
##
##
##
#measure power
dmesh[zz] = BigFileMesh(
dpath + sim + '/fastpm_%0.4f/' % aafiles[iz] + '/dmesh_N%04d' % nc,
'1').paint()
pk = FFTPower(dmesh[zz] / dmesh[zz].cmean(), mode='1d').power
k, pkm = pk['k'], pk['power']
hbins[zz] = np.logspace(
np.log10(hmass[zz][-1]) - 0.01,
np.log10(hmass[zz][0]) - 0.01, 100)
bias_table = []
massval = []
inb = 6
posmesh = pm.paint(hpos[zz])
massmesh = pm.paint(hpos[zz], mass=hmass[zz])
pkhm = FFTPower(massmesh / massmesh.cmean(), mode='1d').power['power']
pkhmx = FFTPower(massmesh / massmesh.cmean(),
second=dmesh[zz] / dmesh[zz].cmean(),
mode='1d').power['power']
pkhp = FFTPower(posmesh / posmesh.cmean(), mode='1d').power['power']
pkhpx = FFTPower(posmesh / posmesh.cmean(),
second=dmesh[zz] / dmesh[zz].cmean(),
mode='1d').power['power']
biases = [(pkhm[1:inb] / pkm[1:inb]).mean()**0.5,
(pkhp[1:inb] / pkm[1:inb]).mean()**0.5,
(pkhmx[1:inb] / pkm[1:inb]).mean(),
(pkhpx[1:inb] / pkm[1:inb]).mean()]
print(biases)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
nnmass = ntools.applynet(mftt, mtup).reshape(nc, nc, nc)
predict = pm.create(mode='real')
predict[...] = nnpred * nnmass
predictr = ntools.relu(predict[...])
#predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
print('Generating data')
dpath = proj + '/data/z%02d/L%04d_N%04d_S%04d_40step/' % (zz * 10, bs,
sfine * nc, seed)
galcat = BigFileCatalog(dpath + 'galaxies_n%02d/galcat' % (numd * 1e4),
header='Header',
comm=pm.comm)
datap = pm.paint(galcat['Position'], mass=galcat['Mass'])
datapR = ft.smooth(datap, Rsm, 'fingauss')
with open(dpath + 'galaxies_n%02d/hodparams.json' % (numd * 1e4)) as fp:
hodparams = json.load(fp)
#hodparams = {'alpha':0.775668, 'logMmin':13.053998, 'logM1':14.3, 'logM0':13.6176, 'sigma_logM':0.397894}
with open(mtpath + 'hodparams.json', 'w') as fp:
json.dump(hodparams, fp, sort_keys=True, indent=4)
print('Data generated')
#########
#Get the prediction for galaxy mass
def mhtomstar2(mh, satfunc=False):
ftt = ntools.createdata(pm, meshdict, pdict['pftname'], plocal)
mftt = ntools.createdata(pm, meshdict, mftname, mlocal)
nnpred = ntools.applynet(ftt, ptup).reshape(nc, nc, nc)
if regression:
nnmass = ntools.applymassreg(mftt, mtup).reshape(nc, nc, nc)
else:
nnmass = ntools.applynet(mftt, mtup).reshape(nc, nc, nc)
if doexp:
nnmass = mf.fmexp(ntools.relu(nnmass), mexp, cc)
predict = pm.create(mode='real')
predict[...] = nnpred * nnmass
if rsd:
cat = BigFileCatalog(scratch + '/data/L%04d_N%04d_S%04d_05step/dynamic/1' %
(bs, nc, seed),
header='Header')
mass = pm.paint(cat['Position']) + 1e-5
mom = pm.paint(cat['Position'], mass=cat['Velocity'][:, 2])
vel = mom / mass
gridpt = pm.generate_uniform_particle_grid(shift=0)
vgrid = vel.readout(gridpt)
predictgrid = predict.readout(gridpt)
predict = pm.paint(
gridpt + vgrid.reshape(-1, 1) * np.array([0, 0, 1]).reshape(1, -1),
mass=predictgrid)
predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
print('Generating data')
#hdictf = ntools.gridhalos(pm, scratch +'/data/L%04d_N%04d_S%04d_40step/'%(bs, 4*nc, seed), rank=num, R1=R1, R2=R2, pmesh=True)
for i, zz in enumerate(zzfiles):
plt.hist(np.log10(smass[zz]), color='C%d'%i, bins=20, histtype='step', label=zz, log=True, lw=1.5, normed=False)
plt.hist(np.log10(cmass[zz]), color='C%d'%i, bins=20, histtype='step', log=True, lw=2, normed=False, ls="--")
plt.axvline(np.log10(dohod.HI_mass(1, aafiles[i], 'mcut')), color='C%d'%i, ls=":")
plt.legend(ncol=2, loc=1)
plt.xlim(8.5, 12)
plt.ylim(10, 2e7)
plt.savefig('./figs/%s/massfunc%s'%(subf, suff))
## Plot bias
start = time()
fig, ax = plt.subplots(1, 2, figsize = (9, 4))
for iz, zz in enumerate(zzfiles):
sh1mesh = pm.paint(spos[zz], mass=sh1mass[zz])
ch1mesh = pm.paint(cpos[zz], mass=ch1mass[zz])
h1mesh = sh1mesh + ch1mesh
pkh1 = FFTPower(h1mesh/h1mesh.cmean(), mode='1d').power['power']
pkh1m = FFTPower(h1mesh/h1mesh.cmean(), second=dmesh[zz]/dmesh[zz].cmean(), mode='1d').power['power']
b1h1 = pkh1m/pkm[zz]
b1h1sq = pkh1/pkm[zz]
ax[0].plot(k, b1h1, 'C%d'%iz, lw=1.5)
ax[0].plot(k, b1h1sq**0.5, 'C%d--'%iz, lw=2)
ax[1].plot(zz, b1h1[1:5].mean(), 'C%do'%iz, label='%0.2f'%(zz))
ax[1].plot(zz, (b1h1sq**0.5)[1:5].mean(), 'C%d*'%iz)
for axis in ax:
axis.legend()
ax[0].set_xscale('log')
def test_grid(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
real = pm.paint(grid)
assert_array_equal(real, 1.0)
#
wts = BigFileCatalog(ldpath + '/lagweights/')
wts['1'] = wts['b1'] * 0 + 1.
grid = wts['InitPosition']
iddg = wts['ID']
print('asserting')
np.allclose(idd.compute(), iddg.compute())
print('read fields')
disp = grid.compute() - 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))
eul_fields = []
for i in range(len(names)):
eul_fields.append(pm.paint(fpos, mass=wts[names[i]], layout=play))
print(rank, names[i], eul_fields[-1].cmean())
#if abs(eul_fields[-1].cmean()) > 0.1: eul_fields[-1] /= eul_fields[-1].cmean()
del dyn, fpos, idd, iddg, wts
print('got fields')
k, spectra = tools.getspectra(eul_fields)
np.savetxt(ofolder + '/spectra2-z%03d-R%d.txt' % (zz * 100, Rsm),
np.vstack([k, spectra]).T.real,
header='k / ' + header,
fmt='%0.4e')
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8')
rank = pm.comm.rank
hcat = BigFileCatalog(dpath + '/fastpm_%0.4f/LL-0.200/' % aa)
print(rank, 'files read')
#print('Mass : ', rank, hcat['Length'][-1].compute()*hcat.attrs['M0']*1e10)
numd = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4, 5e-5]
num = [int(bs**3 * i) for i in numd]
for i in range(len(num)):
if rank == 0: print(numd[i])
cat = hcat.gslice(start=0, stop=num[i])
hlay = pm.decompose(cat['Position'])
hmesh = pm.paint(cat['Position'], layout=hlay)
hmesh /= hmesh.cmean()
ph = FFTPower(hmesh, mode='1d').power
k, ph = ph['k'], ph['power']
np.savetxt(ofolder + '/ph-%05d.txt' % (numd[i] * 1e5),
np.vstack([k, ph]).T.real,
header='k ph',
fmt='%0.4e')
bs, nc)
#dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/'%(bs, nc)
aa = 1.0000
zz = 1 / aa - 1
Rsm = 0
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
rank = pm.comm.rank
#grid = pm.mesh_coordinates()*bs/nc
hcat = BigFileCatalog(dpath + '/fastpm_%0.4f/LL-0.200/' % aa)
hpos = hcat['Position'].compute()
#hmass = print('Mass : ', rank, hcat['Mass'][-1].compute())
hlay = pm.decompose(hpos)
hmesh = pm.paint(hpos, layout=hlay)
hmesh /= hmesh.cmean()
ph = FFTPower(hmesh, mode='1d').power
k, ph = ph['k'], ph['power'].real
sn = (hpos.shape[0] / bs**3)**-1
print(sn)
if rank == 0:
for Rsm in [0, 2]:
for zadisp in [True, False]:
header = '1, b1, b2, bg, bk'
if zadisp:
spectra = np.loadtxt(
'./output/%s/spectraza-%04d-%04d-%04d-R%d.txt' %
elif bs == 256:
censuff = ''
satsuff='-m1_5p0min-alpha_0p8'
pks = []
for aa in alist[:2]:
if rank ==0 : print('For z = %0.2f'%(1/aa-1))
#cencat = BigFileCatalog(scratch2+sim+'/fastpm_%0.4f/cencat'%aa+censuff)
cencat = BigFileCatalog(scratch1+sim+'/fastpm_%0.4f/LL-0.200'%aa)
mp = cencat.attrs['M0'][0]*1e10
cencat['Mass'] = cencat['Length'] * mp
#cencat['HImass'] = HI_hod(cencat['Mass'],aa)
h1mesh = pm.paint(cencat['Position'], mass=cencat['Mass'])
#h1mesh = pm.paint(cencat['Position'])
print('Rank, mesh.cmean() : ', rank, h1mesh.cmean())
h1mesh /= h1mesh.cmean()
#
pkh1h1 = FFTPower(h1mesh,mode='1d').power
# Extract the quantities we want and write the file.
kk = pkh1h1['k']
sn = pkh1h1.attrs['shotnoise']
pk = np.abs(pkh1h1['power'])
pks.append(pk)
#if rank ==0 :
aas = [
0.5,
]
for aa in aas:
zz = 1 / aa - 1
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8')
rank = pm.comm.rank
dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1' % aa)
# Load Matter Mesh
fpos = dyn['Position'].compute()
mlay = pm.decompose(fpos)
mmesh = pm.paint(fpos, layout=mlay)
mmesh /= mmesh.cmean()
# Load halo mesh
cencat = BigFileCatalog(galpath + '/fastpm_%0.4f/' % (aa) + cendir)
satcat = BigFileCatalog(galpath + '/fastpm_%0.4f/' % (aa) + satdir)
hpos = np.concatenate((cencat['Position'], satcat['Position']))
hlay = pm.decompose(hpos)
hmesh = pm.paint(hpos, layout=hlay)
hmesh /= hmesh.cmean()
phm = FFTPower(mmesh, second=hmesh, mode='1d').power
k, phm = phm['k'], phm['power'].real
phh = FFTPower(hmesh, mode='1d').power
predict = pm.create(mode='real', value=nnpred * nnmass)
predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
print('Generating data')
hdictf = ntools.gridhalos(pm,
dpath=scratch +
'/data/z%02d/L%04d_N%04d_S%04d_40step/' %
(zz * 10, bs, sfine * nc, seed),
R1=R1,
R2=R2,
pmesh=False,
abund=abund)[1]
datap = pm.paint(hdictf['position'][:num], mass=hdictf['mass'][:num])
datapR = ft.smooth(datap, Rsm, 'fingauss')
print('Data generated')
func = dg.normal
colors = ['r', 'b', 'g', 'y', 'm', 'orange', 'brown', 'k']
###########################################
mbins = np.logspace(10, 13, 16)[::-1]
msave = [mbins[0] * 100] + list(mbins)
print('mbins -- ', mbins)
####
def test_grid(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8')
grid = pm.generate_uniform_particle_grid(shift=0.5)
assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod())
real = pm.paint(grid)
assert_array_equal(real, 1.0)
def look_den_slice(net1, net2, net3, s, title):
matplotlib.rc('xtick', labelsize=12)
matplotlib.rc('ytick', labelsize=12)
matplotlib.rc('font', size=12)
cmap = 'coolwarm'
assert net1.shape[-1] == 3
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
q = pm.generate_uniform_particle_grid()
net1 = pm.paint(q + net1.reshape([-1, 3]))
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
net2 = pm.paint(q + net2.reshape([-1, 3]))
pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32])
net3 = pm.paint(q + net3.reshape([-1, 3]))
fig = plt.figure(figsize=(6, 4))
amp_low = min(np.percentile(net1[:, :, s], 5),
np.percentile(net2[:, :, s], 5),
np.percentile(net3[:, :, s], 5))
amp_high = max(np.percentile(net1[:, :, s], 95),
np.percentile(net2[:, :, s], 95),
np.percentile(net3[:, :, s], 95))
amp_low1 = min(np.percentile((net1[:, :, s] - net2[:, :, s]), 5),
np.percentile((net1[:, :, s] - net3[:, :, s]), 5))
amp_high1 = max(np.percentile((net1[:, :, s] - net2[:, :, s]), 95),
np.percentile((net1[:, :, s] - net3[:, :, s]), 95))
amp = max(np.abs(amp_low1), amp_high1)
gs = GridSpec(2, 4, height_ratios=[1, 1], width_ratios=[1, 1, 1, 0.1])
plt.subplot(gs[0, 0])
im = plt.imshow(net1[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('fastPM')
plt.subplot(gs[0, 1])
plt.imshow(net2[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('2LPT')
plt.subplot(gs[0, 2])
plt.imshow(net3[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title('U-Net')
cbax = plt.subplot(gs[0, 3])
cbar = fig.colorbar(mappable=im,
cax=cbax,
orientation='vertical',
ticklocation='right')
cbar.ax.tick_params(labelsize=12)
plt.subplot(gs[1, 1])
im = plt.imshow((net1[:, :, s] - net2[:, :, s]),
cmap=cmap,
vmin=-amp,
vmax=amp)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title(r'fastPM $-$ 2LPT')
plt.subplot(gs[1, 2])
plt.imshow((net1[:, :, s] - net3[:, :, s]), cmap=cmap, vmin=-amp, vmax=amp)
#plt.xticks([0,16,32],[0,64,128])
#plt.yticks([16,0],[64,128])
plt.axis('off')
plt.title(r'fastPM $-$ U-Net')
cbax = plt.subplot(gs[1, 3])
cbar = fig.colorbar(mappable=im,
cax=cbax,
orientation='vertical',
ticklocation='right')
cbar.ax.tick_params(labelsize=12)
fig.subplots_adjust(hspace=.2)
poslarge = np.empty((1, 3))
save = False
#print(rank, 'Position created')
layout = pm.decompose(poslarge)
#if rank == 0 : print(rank, 'Layout decomposed')
for i in range(len(meshes)):
vals = meshes[i].readout(poslarge,
layout=layout,
resampler='nearest').astype(
np.float32)
name = names[i]
if save:
savemesh = pmsmall.paint(gridsmall,
mass=vals,
resampler='nearest')
totals[bindex, i] = savemesh.csum()
#print('In rank {:3d}, sum for mesh {:8} is equal to {:.3e}'.format(rank, name, savemesh.csum()))
indices = comm.gather(indices, root=0)
totals = comm.gather(totals, root=0)
if rank == 0:
indices = np.concatenate(
[indices[ii][bindexsplit[ii]] for ii in range(wsize)], axis=0)
totals = np.concatenate(
[totals[ii][bindexsplit[ii]] for ii in range(wsize)], axis=0)
tosave = np.concatenate((indices, totals), axis=1)
#print(tosave)
header = 'ix, iy, iz, dm, h1fid, h1new, lum, uv'
alpha,
beta,
scatter=0.3)
cencat['luminosity'] = lq(cencat['blackhole'],
fon=switchon,
eta=0.1,
scatter=0.3)
satcat['luminosity'] = lq(satcat['blackhole'],
fon=switchon,
eta=0.1,
scatter=0.3)
clayout = pm.decompose(cencat['Position'])
slayout = pm.decompose(satcat['Position'])
cmesh = pm.paint(cencat['Position'],
mass=cencat['luminosity'],
layout=clayout)
smesh = pm.paint(satcat['Position'],
mass=satcat['luminosity'],
layout=slayout)
lmesh = cmesh + smesh
if lumspectra: calc_bias(aa, lmesh / lmesh.cmean(), suff, fname='Lum')
mfpath = mfpathz(zz)
if rank == 0:
print('At redshift {:.2f}, mean free path is {:.2f}'.format(
zz, mfpath))
meshc = lmesh.r2c()
kmesh = sum(i**2 for i in meshc.x)**0.5
wt = np.arctan(kmesh * mfpath) / kmesh / mfpath
wt[kmesh == 0] = 1
predictR = ft.smooth(predict, Rsm, 'fingauss')
#data
hdictf = ntools.gridhalos(pm,
scratch + '/data/L%04d_N%04d_S%04d_40step/' %
(bs, fine * nc, seed),
R1=R1,
R2=R2,
pmesh=True,
abund=abund,
doexp=doexp,
mexp=mexp,
cc=cc,
stellar=stellar)[1]
datap = pm.paint(hdictf['position'][:num], hdictf['mass'][:num])
datapR = ft.smooth(datap, Rsm, 'fingauss')
predicts.append(predictR[...].copy())
datas.append(datapR[...].copy())
for i, sg in enumerate(sgs):
hmass, hpos = dg.scatter_catalog(hdictf['mass'],
hdictf['position'],
sg,
seed=i)
datasg = pm.paint(hpos[:num], hmass[:num])
datasgR = ft.smooth(datasg, Rsm, 'fingauss')
datasgs[i].append(datasgR[...].copy())
#####Setup
def measurepk(nc=nc, dpath=scratchyf):
'''plot the power spectrum of halos on 'nc' grid'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
for i, aa in enumerate(aafiles):
zz = zzfiles[i]
if rank == 0: print('redshift = ', zz)
if ncsim == 10240:
dm = BigFileMesh(scratchyf+sim+'/fastpm_%0.4f/'%aa+\
'/1-mesh/N%04d'%nc,'').paint()
else: dm = BigFileMesh(project+sim+'/fastpm_%0.4f/'%aa+\
'/dmesh_N%04d/1/'%nc,'').paint()
#dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint()
halos = BigFileCatalog(scratchyf + sim + '/fastpm_%0.4f/' % aa,
dataset='LL-0.200')
mp = halos.attrs['MassTable'][1] * 1e10
if rank == 0: print('Mass of particle is = %0.2e' % mp)
hmass = halos["Length"].compute() * mp
hpos = halos['Position'].compute()
layout = pm.decompose(hpos)
if rank == 0: print("paint")
hpmesh = pm.paint(hpos, layout=layout)
hmesh = pm.paint(hpos, mass=hmass, layout=layout)
print(rank, dm.cmean(), hmesh.cmean(), hpmesh.cmean())
pkm = FFTPower(dm / dm.cmean(), mode='1d').power
pkh = FFTPower(hmesh / hmesh.cmean(), mode='1d').power
pkhp = FFTPower(hpmesh / hpmesh.cmean(), mode='1d').power
pkhm = FFTPower(hmesh / hmesh.cmean(),
second=dm / dm.cmean(),
mode='1d').power
pkhpm = FFTPower(hpmesh / hpmesh.cmean(),
second=dm / dm.cmean(),
mode='1d').power
def savebinned(path, binstat, header):
if halos.comm.rank == 0:
k, p, modes = binstat['k'].real, binstat[
'power'].real, binstat['modes'].real
np.savetxt(path,
np.stack((k, p, modes), axis=1),
header=header)
ofolder = "../data/outputs/halos/{}/".format(sim)
if rank == 0:
print(ofolder)
try:
os.makedirs(ofolder)
except:
pass
savebinned(ofolder + 'pkhp_%0.4f.txt' % aa,
pkhp,
header='k, P(k), Nmodes')
savebinned(ofolder + 'pkhm_%0.4f.txt' % aa,
pkh,
header='k, P(k), Nmodes')
savebinned(ofolder + 'pkd_%0.4f.txt' % aa,
pkm,
header='k, P(k), Nmodes')
savebinned(ofolder + 'pkhpxd_%0.4f.txt' % aa,
pkhpm,
header='k, P(k), Nmodes')
savebinned(ofolder + 'pkhmxd_%0.4f.txt' % aa,
pkhm,
header='k, P(k), Nmodes')
