def make_data(bs, nc, seed, nsteps, path='', z=0):
#initiate
pm = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8')
ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep/'%(z*10, bs, nc, seed, nsteps)
#
if pm.comm.rank == 0:
print('Make data for seed = %d'%seed)
print('Data to be saved at path - %s'%ofolder)
klin, plin = numpy.loadtxt('ics_matterpow_0.dat', unpack = True)
pk = interpolate(klin, plin)
cosmo = Planck15.clone(Omega_cdm = 0.2685, h = 0.6711, Omega_b = 0.049)
#pk = EHPower(Planck15, redshift=0)
#cosmo = Planck15
s_truth = pm.generate_whitenoise(seed, type='complex')\
.apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
.c2r()
#dynamics
aa = 1.0/(1+z)
if pm.comm.rank == 0:
print('Evolve to redshift = %0.1f, scale factor = %0.2f'%(z, aa))
stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)
start = time()
dynamic_model = NBodyModel(cosmo, pm, B=2, steps=stages)
#dynamic_model = LPTModel(cosmo, pm, B=2, steps=stages)
#Save data
X, V, final = dynamic_model.get_code().compute(['X', 'V', 'final'], init={'parameters':s_truth})
end = time()
print('Time taken = ', end-start)
save_map(s_truth, ofolder + 'mesh', 's')
save_map(final, ofolder + 'mesh', 'd')
if pm.comm.rank == 0:
print('X, V computed')
cat = ArrayCatalog({'Position': X, 'Velocity' : V}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
kdd = KDDensity(cat).density
cat['KDDensity'] = kdd
cat['InitPosition'] = dynamic_model.get_code().engine.q
cat.save(ofolder + 'dynamic/1', ('InitPosition', 'Position', 'Velocity', 'KDDensity'))
if pm.comm.rank == 0:
print('dynamic model created')
#FOF
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.find_features(peakcolumn='KDDensity')
fofcat['Mass'] = fofcat['Length'] * cosmo.Om0 * 27.7455 * pm.BoxSize.prod() / pm.Nmesh.prod()
fofcat.save(ofolder+'FOF', ('CMPosition', 'CMVelocity', 'PeakPosition', 'PeakVelocity', 'Length', 'Mass'))
def randoms(ns):
zmin = 0
zmax = ns.zmax
dNdZ = read_Nz(ns.nz, ns.ncol, zmin, zmax)
zedges = numpy.linspace(zmin, zmax, 12)
zcenters = 0.5 * (zedges[1:] + zedges[:-1])
dNdZ_ran = lambda z: dNdZ(z) * ns.boost
icdf, Ntotal = geticdf(dNdZ_ran, numpy.linspace(zmin, zmax, 1024))
# drop about 20% of samples to give it enough freedom.
cat = RandomCatalog(csize=int(Ntotal), seed=SEED * 20 + 19)
if cat.comm.rank == 0:
cat.logger.info("Total = %d" % Ntotal)
# generate points uniformly on a sphere.
# FIXME: add uniform sphere to nbodykit's RandomCatalog
z = cat.rng.uniform(-1., 1.)
r = abs(1 - z**2)**0.5
phi = cat.rng.uniform(0, 2 * numpy.pi)
x = da.cos(phi) * r
y = da.sin(phi) * r
z = z
ra, dec = transform.CartesianToEquatorial(da.stack((x, y, z), axis=1),
frame='galactic')
z = icdf(cat.rng.uniform(0, 1))
cat['RA'] = ra.compute()
cat['DEC'] = dec.compute()
cat['Z'] = z
cat['Aemit'] = (1 / (z + 1))
ntarget = numpy.ones_like(z, dtype='i4')
randoms = ArrayCatalog(
{
'RA': cat['RA'].compute().repeat(ntarget, axis=0),
'DEC': cat['DEC'].compute().repeat(ntarget, axis=0),
'Z': cat['Z'].compute().repeat(ntarget, axis=0),
'Aemit': cat['Aemit'].compute().repeat(ntarget, axis=0),
},
comm=cat.comm)
randoms = randoms.sort(('Aemit', ), usecols=['RA', 'DEC', 'Z', 'Aemit'])
randoms.save(ns.output, dataset=ns.odataset)
def monitor(action, ai, ac, af, state, event):
if not state.synchronized: return
global a0, X0, V0
for a in a_output:
if a > a0 and a <= af:
#interpolate particle positions and velocities
pt = MatterDominated(state.cosmology.Om0,
a=[a0, a, af],
a_normalize=state.solver.a_linear)
fac1 = (pt.Gp(af) - pt.Gp(a)) / (pt.Gp(af) - pt.Gp(a0))
fac2 = (pt.Gp(a) - pt.Gp(a0)) / (pt.Gp(af) - pt.Gp(a0))
X = fac1 * X0 + fac2 * state.X
fac1 = (pt.Gf(af) - pt.Gf(a)) / (pt.Gf(af) - pt.Gf(a0))
fac2 = (pt.Gf(a) - pt.Gf(a0)) / (pt.Gf(af) - pt.Gf(a0))
V = fac1 * V0 + fac2 * state.V
#save the snapshot
cat = ArrayCatalog({
'Position': X,
'Velocity': V
},
BoxSize=state.pm.BoxSize)
cat.save(
args.save + '/FastPM_Nmesh%d_Nstep%d_z%.2f' %
(args.Nmesh, args.Nstep, 1. / a - 1.),
('Position', 'Velocity'))
del X, V
if state.pm.comm.rank == 0:
print(
'Finish writing snapshot at redshift %.2f' % (1. / a - 1.),
'Time:',
time.time() - t)
a0 = np.array(af)
X0 = state.X
V0 = state.V
if state.pm.comm.rank == 0:
print('Finish redshift %.2f' % (1. / af - 1.), 'Time:',
time.time() - t)
def monitor(action, ai, ac, af, state, event):
if not state.synchronized: return
global a0, X0, V0
for a in a_output:
if a > a0 and a <= af:
#interpolate particle positions and velocities
pt = MatterDominated(state.cosmology.Om0, a=[a0, a, af], a_normalize=state.solver.a_linear)
fac1 = (pt.Gp(af) - pt.Gp(a)) / (pt.Gp(af) - pt.Gp(a0))
fac2 = (pt.Gp(a) - pt.Gp(a0)) / (pt.Gp(af) - pt.Gp(a0))
X = fac1 * X0 + fac2 * state.X
fac1 = (pt.Gf(af) - pt.Gf(a)) / (pt.Gf(af) - pt.Gf(a0))
fac2 = (pt.Gf(a) - pt.Gf(a0)) / (pt.Gf(af) - pt.Gf(a0))
V = fac1 * V0 + fac2 * state.V
#save the snapshot
cat = ArrayCatalog({'Position' : X, 'Velocity' : V}, BoxSize=state.pm.BoxSize)
cat.save(args.save + '/FastPM_Nmesh%d_Nstep%d_z%.2f' % (args.Nmesh, args.Nstep, 1./a-1.), ('Position', 'Velocity'))
del X, V
if state.pm.comm.rank == 0:
print('Finish writing snapshot at redshift %.2f' % (1./a-1.), 'Time:', time.time()-t)
#mass = 1.0 * state.pm.Nmesh.prod() / state.pm.comm.allreduce(len(X), op=MPI.SUM)
#layout = pm.decompose(X)
#pos1 = layout.exchange(X)
#del X
#TNGmap = state.pm.create(type="real")
#TNGmap.paint(pos1, mass=mass, layout=None, hold=False)
#FieldMesh(TNGmap).save('dmmap_Nmesh64_z%.2f' % (1/af-1))
a0 = np.array(af)
X0 = state.X
V0 = state.V
if state.pm.comm.rank == 0:
print('Finish redshift %.2f' % (1./af-1.), 'Time:', time.time()-t)
def make_galcat(aa, mmin, m1, alpha=0.9, censuff=None, satsuff=None, ofolder=None, seed=3333):
'''Assign 0s to
'''
zz = atoz(aa)
#halocat = readincatalog(aa)
halocat = BigFileCatalog(dpath + '/fastpm_%0.4f/'%aa, dataset='LL-0.200')
rank = halocat.comm.rank
halocat.attrs['BoxSize'] = np.broadcast_to(halocat.attrs['BoxSize'], 3)
ghid = halocat.Index.compute()
halocat['GlobalIndex'] = ghid
mp = halocat.attrs['MassTable'][1]*1e10
halocat['Mass'] = halocat['Length'] * mp
halocat['Position'] = halocat['Position']%bs # Wrapping positions assuming periodic boundary conditions
rank = halocat.comm.rank
halocat = halocat.to_subvolumes()
if rank == 0: print('\n ############## Redshift = %0.2f ############## \n'%zz)
hmass = halocat['Mass'].compute()
hpos = halocat['Position'].compute()
hvel = halocat['Velocity'].compute()
rvir = HaloRadius(hmass, cosmo, 1/aa-1).compute()/aa
vdisp = HaloVelocityDispersion(hmass, cosmo, 1/aa-1).compute()
ghid = halocat['GlobalIndex'].compute()
# Select halos that have galaxies
rands = np.random.uniform(size=len(hmass))
ncen = hod.ncen(mh=hmass,mcutc=mmin,sigma=0.2)
ws = (rands < ncen)
hmass = hmass[ws]
hpos = hpos[ws]
hvel = hvel[ws]
rvir = rvir[ws]
vdisp = vdisp[ws]
ghid = ghid[ws]
print('In rank = %d, Catalog size = '%rank, hmass.size)
#Do hod
start = time()
ncen = np.ones_like(hmass)
#nsat = hod.nsat_martin(msat = mmin, mh=hmass, m1f=m1f, alpha=alpha).astype(int)
nsat = hod.nsat_zheng(mh=hmass, m0=mmin, m1=20*mmin, alpha=alpha).astype(int)
#Centrals
cpos, cvel, gchid, chid = hpos, hvel, ghid, np.arange(ncen.size)
spos, svel, shid = hod.mksat(nsat, pos=hpos, vel=hvel,
vdisp=vdisp, conc=7, rvir=rvir, vsat=0.5, seed=seed)
gshid = ghid[shid]
svelh1 = svel*2/3 + cvel[shid]/3.
smmax = hmass[shid]/10.
smmin = np.ones_like(smmax)*mmin
mask = smmin > smmax/3. #Some fudge that should be discussed
smmin[mask] = smmax[mask]/3.
smass = hod.get_msat(hmass[shid], smmax, smmin, alpha)
sathmass = np.zeros_like(hmass)
tot = np.bincount(shid, smass)
sathmass[np.unique(shid)] = tot
cmass = hmass - sathmass # assign remaining mass in centrals
print('In rank = %d, Time taken = '%rank, time()-start)
print('In rank = %d, Number of centrals & satellites = '%rank, ncen.sum(), nsat.sum())
print('In rank = %d, Satellite occupancy: Max and mean = '%rank, nsat.max(), nsat.mean())
#
#Save
cencat = ArrayCatalog({'Position':cpos, 'Velocity':cvel, 'Mass':cmass, 'GlobalID':gchid,
'Nsat':nsat, 'HaloMass':hmass},
BoxSize=halocat.attrs['BoxSize'], Nmesh=halocat.attrs['NC'])
minid, maxid = cencat['GlobalID'].compute().min(), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
cencat = cencat.sort('GlobalID')
minid, maxid = cencat['GlobalID'].compute().min(), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if censuff is not None:
colsave = [cols for cols in cencat.columns]
cencat.save(ofolder+'cencat'+censuff, colsave)
satcat = ArrayCatalog({'Position':spos, 'Velocity':svel, 'Velocity_HI':svelh1, 'Mass':smass,
'GlobalID':gshid, 'HaloMass':hmass[shid]},
BoxSize=halocat.attrs['BoxSize'], Nmesh=halocat.attrs['NC'])
minid, maxid = satcat['GlobalID'].compute().min(), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
satcat = satcat.sort('GlobalID')
minid, maxid = satcat['GlobalID'].compute().min(), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if satsuff is not None:
colsave = [cols for cols in satcat.columns]
satcat.save(ofolder+'satcat'+satsuff, colsave)
def make_galcat(aa,
mmin,
m1f,
alpha=-1,
censuff=None,
satsuff=None,
ofolder=None,
seed=3333):
'''Assign 0s to
'''
zz = 1 / aa - 1
#halocat = readincatalog(aa)
halocat = BigFileCatalog(args['halofilez'] % aa,
dataset=args['halodataset'])
rank = halocat.comm.rank
mpart, Lbox, rsdfac, acheck = read_conversions(args['headerfilez'] % aa)
halocat.attrs['BoxSize'] = [bs, bs, bs]
halocat.attrs['NC'] = nc
ghid = halocat.Index.compute()
halocat['GlobalIndex'] = ghid
halocat['Mass'] = halocat['Length'] * mpart
halocat['Position'] = halocat[
'Position'] % bs # Wrapping positions assuming periodic boundary conditions
rank = halocat.comm.rank
halocat = halocat.to_subvolumes()
if rank == 0:
print('\n ############## Redshift = %0.2f ############## \n' % zz)
hmass = halocat['Mass'].compute()
hpos = halocat['Position'].compute()
hvel = halocat['Velocity'].compute()
rvir = HaloRadius(hmass, cosmo, 1 / aa - 1).compute() / aa
vdisp = HaloVelocityDispersion(hmass, cosmo, 1 / aa - 1).compute()
ghid = halocat['GlobalIndex'].compute()
print('In rank = %d, Catalog size = ' % rank, hmass.size)
#Do hod
start = time()
ncen = np.ones_like(hmass)
nsat = hod.nsat_martin(msat=mmin, mh=hmass, m1f=m1f,
alpha=alpha).astype(int)
#Centrals
cpos, cvel, gchid, chid = hpos, hvel, ghid, np.arange(ncen.size)
spos, svel, shid = hod.mksat(nsat,
pos=hpos,
vel=hvel,
vdisp=vdisp,
conc=7,
rvir=rvir,
vsat=0.5,
seed=seed)
gshid = ghid[shid]
svelh1 = svel * 2 / 3 + cvel[shid] / 3.
smmax = hmass[shid] / 10.
smmin = np.ones_like(smmax) * mmin
mask = smmin > smmax / 3. #If Mmin and Mmax are too close for satellites, adjust Mmin
smmin[mask] = smmax[mask] / 3.
smass = hod.get_msat(hmass[shid], smmax, smmin, alpha)
sathmass = np.zeros_like(hmass)
tot = np.bincount(shid, smass)
sathmass[np.unique(shid)] = tot
cmass = hmass - sathmass # assign remaining mass in centrals
print('In rank = %d, Time taken = ' % rank, time() - start)
print('In rank = %d, Number of centrals & satellites = ' % rank,
ncen.sum(), nsat.sum())
print('In rank = %d, Satellite occupancy: Max and mean = ' % rank,
nsat.max(), nsat.mean())
#
#Save
cencat = ArrayCatalog(
{
'Position': cpos,
'Velocity': cvel,
'Mass': cmass,
'GlobalID': gchid,
'Nsat': nsat,
'HaloMass': hmass
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
minid, maxid = cencat['GlobalID'].compute().min(
), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
cencat = cencat.sort('GlobalID')
minid, maxid = cencat['GlobalID'].compute().min(
), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if censuff is not None:
colsave = [cols for cols in cencat.columns]
cencat.save(ofolder + 'cencat' + censuff, colsave)
satcat = ArrayCatalog(
{
'Position': spos,
'Velocity': svel,
'Velocity_HI': svelh1,
'Mass': smass,
'GlobalID': gshid,
'HaloMass': hmass[shid]
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
minid, maxid = satcat['GlobalID'].compute().min(
), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
satcat = satcat.sort('GlobalID')
minid, maxid = satcat['GlobalID'].compute().min(
), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if satsuff is not None:
colsave = [cols for cols in satcat.columns]
satcat.save(ofolder + 'satcat' + satsuff, colsave)
#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
path = '/project/projectdirs/astro250/chmodi/cosmo4d/data/'
ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep_fpm/' % (
z * 10, bs, nc, seed, nsteps)
FieldMesh(linear).save(ofolder + 'mesh', dataset='s', mode='real')
FieldMesh(final).save(ofolder + 'mesh', dataset='d', mode='real')
cat.save(ofolder + 'dynamic/1', ('Position', 'Velocity', 'KDDensity'))
cat = ArrayCatalog({'ID': idd, 'InitPosition': grid}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
#cat.save(lpath + 'initpos', tosave)
#cat.attrs = dyn.attrs
#header = '1,b1,b2,bg,bk'
header = 'b1,b2,bg,bk'
names = header.split(',')
#cat.save(lpath + 'dynamic/1', ('ID', 'InitPosition', 'Position'))
tosave = ['ID', 'InitPosition'] + names
if rank == 0: print(tosave)
for i in range(len(names)):
ff = BigFileMesh(lpath+ '/lag', names[i]).paint()
pm = ff.pm
rank = pm.comm.rank
glay, play = pm.decompose(grid), pm.decompose(fpos)
wts = ff.readout(grid, layout = glay, resampler='nearest')
cat[names[i]] = wts
#x = FieldMesh(pm.paint(fpos, mass=wts, layout=play))
#x.save(lpath + 'eul-z%03d-R%d'%(zz*100, Rsm), dataset=names[i], mode='real')
del pm, ff, wts
if rank == 0: print(rank, ' got weights')
cat.save(lpath + 'lagweights/', tosave)
def make_galcat(aa,
nc,
seed,
bs=1536,
T=40,
B=2,
mmin=10**12.5,
m1=20 * 10**12.5,
alpha=0.9,
censuff=None,
satsuff=None,
simpath=sc_simpath,
outpath=sc_outpath):
'''Assign 0s to
'''
fname = get_filename(nc, seed, T=T, B=B)
spath = simpath + fname
opath = outpath + fname + 'hod/'
try:
os.makedirs(opath)
except:
pass
zz = atoz(aa)
#halocat = readincatalog(aa)
halocat = BigFileCatalog(spath + '/fastpm_%0.4f/' % aa, dataset='LL-0.200')
rank = halocat.comm.rank
halocat.attrs['BoxSize'] = np.broadcast_to(halocat.attrs['BoxSize'], 3)
ghid = halocat.Index.compute()
halocat['GlobalIndex'] = ghid
mp = halocat.attrs['MassTable'][1] * 1e10
halocat['Mass'] = halocat['Length'] * mp
halocat['Position'] = halocat[
'Position'] % bs # Wrapping positions assuming periodic boundary conditions
rank = halocat.comm.rank
halocat = halocat.to_subvolumes()
if rank == 0:
print('\n ############## Redshift = %0.2f ############## \n' % zz)
hmass = halocat['Mass'].compute()
hpos = halocat['Position'].compute()
hvel = halocat['Velocity'].compute()
rvir = HaloRadius(hmass, cosmo, 1 / aa - 1).compute() / aa
vdisp = HaloVelocityDispersion(hmass, cosmo, 1 / aa - 1).compute()
ghid = halocat['GlobalIndex'].compute()
# Select halos that have galaxies
rands = np.random.uniform(size=len(hmass))
ncen = hod.ncen(mh=hmass, mcutc=mmin, sigma=0.2)
ws = (rands < ncen)
hmass = hmass[ws]
hpos = hpos[ws]
hvel = hvel[ws]
rvir = rvir[ws]
vdisp = vdisp[ws]
ghid = ghid[ws]
print('In rank = %d, Catalog size = ' % rank, hmass.size)
#Do hod
start = time()
ncen = np.ones_like(hmass)
#nsat = hod.nsat_martin(msat = mmin, mh=hmass, m1f=m1f, alpha=alpha).astype(int)
nsat = hod.nsat_zheng(mh=hmass, m0=mmin, m1=20 * mmin,
alpha=alpha).astype(int)
#Centrals
cpos, cvel, gchid, chid = hpos, hvel, ghid, np.arange(ncen.size)
spos, svel, shid = hod.mksat(nsat,
pos=hpos,
vel=hvel,
vdisp=vdisp,
conc=7,
rvir=rvir,
vsat=0.5,
seed=seed)
gshid = ghid[shid]
print('In rank = %d, Time taken = ' % rank, time() - start)
print('In rank = %d, Number of centrals & satellites = ' % rank,
ncen.sum(), nsat.sum())
print('In rank = %d, Satellite occupancy: Max and mean = ' % rank,
nsat.max(), nsat.mean())
#
#Save
cencat = ArrayCatalog({
'Position': cpos,
'GlobalID': gchid,
'Nsat': nsat
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
minid, maxid = cencat['GlobalID'].compute().min(
), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
cencat = cencat.sort('GlobalID')
minid, maxid = cencat['GlobalID'].compute().min(
), cencat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if censuff is not None:
colsave = [cols for cols in cencat.columns]
print("Saving centrals.")
cencat.save(opath + 'cencat' + censuff, colsave)
#satcat = ArrayCatalog({'Position':spos, 'Velocity':svel, 'Velocity_HI':svelh1, 'Mass':smass,
# 'GlobalID':gshid, 'HaloMass':hmass[shid]},
# BoxSize=halocat.attrs['BoxSize'], Nmesh=halocat.attrs['NC'])
satcat = ArrayCatalog({
'Position': spos,
'GlobalID': gshid
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
minid, maxid = satcat['GlobalID'].compute().min(
), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('before ', rank, minid, maxid)
satcat = satcat.sort('GlobalID')
minid, maxid = satcat['GlobalID'].compute().min(
), satcat['GlobalID'].compute().max()
if minid < 0 or maxid < 0:
print('after ', rank, minid, maxid)
if satsuff is not None:
colsave = [cols for cols in satcat.columns]
print("Saving sats.")
satcat.save(opath + 'satcat' + satsuff, colsave)
def make_finer(bs, nc, seed, nsteps, cosmo, pk, ofolder, zz=0):
#initiate
pmf = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8')
try:
if pmf.comm.rank == 0:
print('Trying to read existing catalog from folder \n%s' % ofolder)
fofcat = BigFileCatalog(ofolder + 'FOF', header='Header')
except:
print('File not found in rank = ', pmf.comm.rank)
try:
os.makedirs(ofolder)
except:
pass
if pmf.comm.rank == 0:
print(
'Creating a new simulation with box, mesh, seed = %d, %d, %d'
% (bs, nc, seed))
print('pk works at k = 0.01, p = %0.2f' % pk(0.01))
s_truth = pmf.generate_whitenoise(seed, mode='complex')\
.apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
.c2r()
## s_truth = pmf.generate_whitenoise(seed, mode='complex')\
## .apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\
## .c2r()
##
if pmf.comm.rank == 0:
print('truth generated')
#dynamics
aa = 1.0 / (1 + zz)
stages = numpy.linspace(0.1, aa, nsteps, endpoint=True)
dynamic_model = NBodyModel(cosmo, pmf, B=2, steps=stages)
if pmf.comm.rank == 0:
print('dynamic model created')
print('Starting sim')
X, V, final = dynamic_model.get_code().compute(
['X', 'V', 'final'], init={'parameters': s_truth})
if pmf.comm.rank == 0:
print('X, Y, final computed')
save_map(s_truth, ofolder + 'mesh', 's')
save_map(final, ofolder + 'mesh', 'd')
cat = ArrayCatalog({
'Position': X,
'Velocity': V
},
BoxSize=pmf.BoxSize,
Nmesh=pmf.Nmesh)
kdd = KDDensity(cat).density
cat['KDDensity'] = kdd
cat.save(ofolder + 'dynamic/1', ('Position', 'Velocity', 'KDDensity'))
if pmf.comm.rank == 0:
print('High-res dynamic model created')
#FOF
fof = FOF(cat, linking_length=0.2, nmin=12)
fofcat = fof.find_features(peakcolumn='KDDensity')
fofcat['Mass'] = fofcat[
'Length'] * cosmo.Om0 * 27.7455 * pmf.BoxSize.prod(
) / pmf.Nmesh.prod()
fofile = ofolder + 'FOF'
fofcat.save(fofile, ('CMPosition', 'CMVelocity', 'PeakPosition',
'PeakVelocity', 'Length', 'Mass'))
if pmf.comm.rank == 0:
print('Halos found in high-res simulation')
print('Number of halos found in rank 0 = ', fofcat['Mass'].size)
return fofcat
def mock(ns):
if ns.idataset is None:
ns.idataset = ns.odataset
cat = BigFileCatalog(ns.input, dataset=ns.idataset)
if ns.simcov == 'NGP':
fsky = 0.5
elif ns.simcov == 'FULL':
fsky = 1.0
else:
raise
cat['ZREAL'] = (1 / cat['Aemit'] - 1)
def compute_va(vel, pos):
u = pos / (pos**2).sum(axis=-1)[:, None]**0.5
return numpy.einsum('ij,ij->i', vel, u)
VZ = da.apply_gufunc(compute_va, '(3),(3)->()', cat['Velocity'],
cat['Position'])
C = 299792458. / 1000
cat['Z'] = (1 + cat['ZREAL']) * (1 + VZ / C) - 1
zmin, zmax = da.compute(cat['Z'].min(), cat['Z'].max())
zmax = max(cat.comm.allgather(zmax))
zmin = min(cat.comm.allgather(zmin))
dNdZ = read_Nz(ns.nz, ns.ncol, zmin, zmax)
zedges = numpy.linspace(zmin, zmax, 128)
zcenters = 0.5 * (zedges[:-1] + zedges[1:])
dNdZ1 = fit_dNdZ(cat, zedges, fsky)
Z = cat['Z'].compute()
ntarget = dNdZ(Z) / dNdZ1(Z)
ntarget[numpy.isnan(ntarget)] = 0
#ntarget = ntarget.clip(0, 10)
rng = numpy.random.RandomState((SEED * 20 + 11) * cat.comm.size +
cat.comm.rank)
if all(cat.comm.allgather((ntarget < 1.0).all())):
ntarget = rng.binomial(1, ntarget)
else:
ntarget = rng.poisson(ntarget)
if cat.comm.rank == 0:
cat.logger.info(
"Up-sampling with poisson because number density is too low")
pos = cat['Position'].compute().repeat(ntarget, axis=0)
redshift = cat['Z'].compute().repeat(ntarget, axis=0)
aemit = cat['Aemit'].compute().repeat(ntarget, axis=0)
ra, dec = transform.CartesianToEquatorial(pos, frame='galactic')
if ns.simcov == 'NGP':
if cat.comm.rank == 0:
cat.logger.info(
"Patching the half sky simulation into full sky by flipping z axis"
)
ra2, dec2 = transform.CartesianToEquatorial(pos * [1, 1, -1],
frame='galactic')
cat1 = ArrayCatalog(
{
'RA': numpy.concatenate([ra, ra2], axis=0),
'DEC': numpy.concatenate([dec, dec2], axis=0),
'Aemit': numpy.concatenate([aemit, aemit], axis=0),
'Z': numpy.concatenate([redshift, redshift], axis=0),
},
comm=cat.comm)
elif ns.simcov == 'FULL':
cat1 = ArrayCatalog(
{
'RA': ra,
'DEC': dec,
'Aemit': aemit,
'Z': redshift,
},
comm=cat.comm)
cat1.save(ns.output, dataset=ns.odataset)
def get_lagweights(nc,
seed,
bs=1536,
T=40,
B=2,
simpath=sc_simpath,
outpath=sc_outpath,
dyn=None):
# Note that dyn is the particle catalog, which we reload if not given
aa = 1.0000 # since particle IDs are ordered only need to load at one redshift
zz = 1 / aa - 1
fname = get_filename(nc, seed, T=T, B=B)
spath = simpath + fname
opath = outpath + fname
try:
os.makedirs(opath)
except:
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4')
rank = pm.comm.rank
if dyn is None:
particle_file = spath + '/fastpm_%0.4f/1' % aa
print("Loading particle data from: " + particle_file)
dyn = BigFileCatalog(particle_file)
#
fpos = dyn['Position'].compute()
idd = dyn['ID'].compute()
attrs = dyn.attrs
grid = tools.getqfromid(idd, attrs, nc)
if rank == 0: print('grid computed')
cat = ArrayCatalog({
'ID': idd,
'InitPosition': grid
},
BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh)
header = '1,b1,b2,bg,bk'
names = header.split(',')
tosave = ['ID', 'InitPosition'] + names
if rank == 0: print(tosave)
for i in range(len(names)):
ff = BigFileMesh(opath + '/lag', names[i]).paint()
pm = ff.pm
rank = pm.comm.rank
glay, play = pm.decompose(grid), pm.decompose(fpos)
wts = ff.readout(grid, layout=glay, resampler='nearest')
cat[names[i]] = wts
del pm, ff, wts
if rank == 0: print(rank, ' got weights')
cat.save(opath + 'lagweights/', tosave)
assert action == 'F'
yield action, af, ac, ai
stack = []
elif action == 'F':
# need to do F after D to ensure the time tag is right.
assert ac == af
stack.append((action, af, ai, ai))
else:
yield action, af, ac, ai
print(list(leapfrog(stages)))
print('----')
print(list(gorfpael(stages)))
print('++++')
print(list(leapfrog(stages[::-1])))
state = solver.nbody(lpt.copy(), leapfrog(stages), monitor=monitor)
if pm.comm.rank == 0:
print('----------------')
reverse = solver.nbody(state, leapfrog(stages[::-1]), monitor=monitor)
#print((lpt.X - reverse.X).max())
assert_allclose(lpt.X, reverse.X)
cat1 = ArrayCatalog({'Position' : lpt.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat2 = ArrayCatalog({'Position' : reverse.X}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh)
cat1.save('Janus/Truth', ('Position',))
cat2.save('Janus/Reverse', ('Position',))
def make_galcat(aa,
mmin,
mcutc,
m1,
sigma=0.25,
kappa=1,
alpha=1,
censuff=None,
satsuff=None,
seed=3333):
'''do HOD with Zheng model using nbodykit'''
zz = tools.atoz(aa)
halocat = readincatalog(aa)
rank = halocat.comm.rank
if rank == 0:
print('\n ############## Redshift = %0.2f ############## \n' % zz)
hmass = halocat['Mass'].compute()
hpos = halocat['Position'].compute()
hvel = halocat['Velocity'].compute()
rvir = HaloRadius(hmass, cosmo, 1 / aa - 1).compute() / aa
vdisp = HaloVelocityDispersion(hmass, cosmo, 1 / aa - 1).compute()
print('In rank = %d, Catalog size = ' % rank, hmass.size)
#Do hod
ofolder = myscratch + '/%s/fastpm_%0.4f/' % (sim, aa)
try:
os.makedirs(os.path.dirname(ofolder))
except IOError:
pass
start = time()
#(ncen, cpos, cvel), (nsat, spos, svel) = hod(seed*rank, hmass, halocat['Position'].compute(), halocat['Velocity'].compute(),\
# conc=7, rvir=3, vdisp=1100, mcut=mcutc, m1=m1, sigma=0.25, \
# kappa=kappa, alpha=alpha,vcen=0, vsat=0.5)
(ncen, cpos, cvel), (nsat, spos, svel) = hod(seed*rank, hmass, hpos, hvel,
conc=7, rvir=rvir, vdisp=vdisp, mcut=mcutc, m1=m1, sigma=0.25, \
kappa=kappa, alpha=alpha, vcen=0, vsat=0.5)
print('In rank = %d, Time taken = ' % rank, time() - start)
print('In rank = %d, Number of centrals & satellites = ' % rank,
ncen.sum(), nsat.sum())
print('In rank = %d, Satellite occupancy: Max and mean = ' % rank,
nsat.max(), nsat.mean())
#
#Assign mass to centrals
hid = np.repeat(range(len(hmass)), ncen).astype(int)
cmass = hmass[hid]
cencat = ArrayCatalog(
{
'Position': cpos,
'Velocity': cvel,
'Mass': cmass,
'HaloID': hid
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
if censuff is not None:
colsave = [cols for cols in cencat.columns]
cencat.save(ofolder + 'cencat' + censuff, colsave)
#
#Assign mass to satellites
hid = np.repeat(range(len(hmass)), nsat).astype(int)
np.random.seed(seed * rank)
smass = np.random.uniform(size=hid.size)
mmax = hmass[hid] / 3.
mmin = np.ones_like(mmax) * mmin
mask = mmin > hmass[hid] / 10. #Some fudge that should be discussed
mmin[mask] = hmass[hid][mask] / 10.
#smass = mmin * mmax / ((1-smass)*mmax + smass*mmin)
smass = hmass[hid] * assign_msat(smass, mmin / hmass[hid],
mmax / hmass[hid], alpha)
satcat = ArrayCatalog(
{
'Position': spos,
'Velocity': svel,
'Mass': smass,
'HaloID': hid
},
BoxSize=halocat.attrs['BoxSize'],
Nmesh=halocat.attrs['NC'])
if satsuff is not None:
colsave = [cols for cols in satcat.columns]
satcat.save(ofolder + 'satcat' + satsuff, colsave)