def createmesh(self, bs, nc, positions, weights):
'''use this to create mesh of HI
'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
mesh = pm.create(mode='real', value=0)
comm = pm.comm
rankweight = sum([wt.sum() for wt in weights])
totweight = comm.allreduce(rankweight)
for wt in weights:
wt /= totweight / float(nc)**3
lay = pm.decompose(positions[0])
mesh.paint(positions[0], mass=weights[0], layout=lay, hold=True)
if len(positions) > 1:
for i in range(self.nsat):
shift = np.random.normal(0, positions[1])
pos = positions[0] + shift
lay = pm.decompose(pos)
mesh.paint(pos,
mass=weights[1] / self.nsat,
layout=lay,
hold=True)
return mesh
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_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 simulate(comm, ns):
pm = ParticleMesh(BoxSize=ns.BoxSize,
Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh],
dtype='f8',
comm=comm)
gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)
Q = pm.generate_uniform_particle_grid(shift=0)
print(Q.min(axis=0), Q.max(axis=0))
def convolve(k, v):
kmag = sum(ki**2 for ki in k)**0.5
ampl = (PowerSpectrum(kmag) / v.BoxSize.prod())**0.5
return v * ampl
dlinear = gaussian.apply(convolve)
DX1 = numpy.zeros_like(Q)
layout = pm.decompose(Q)
# Fill it in one dimension at a time.
for d in range(pm.ndim):
DX1[..., d] = dlinear \
.apply(dx1_transfer(d)) \
.c2r().readout(Q, layout=layout)
a0 = time_steps[0]
# 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
S = DX1 * pt.D1(a=a0)
V = S * a0**2 * pt.f1(a0) * pt.E(a0)
state = State(Q, S, V)
fpm = ParticleMesh(BoxSize=pm.BoxSize,
Nmesh=pm.Nmesh * ns.boost,
resampler='tsc',
dtype='f8')
ns.scheme(fpm, state, time_steps, ns.factors)
r = Result()
r.Q = Q
r.DX1 = DX1
r.S = S
r.V = V
r.dlinear = dlinear
return r
def createmesh(self, bs, nc, positions, weights):
'''use this to create mesh of HI
'''
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc])
mesh = pm.create(mode='real', value=0)
comm = pm.comm
rankweight = sum([wt.sum() for wt in weights])
totweight = comm.allreduce(rankweight)
for wt in weights:
wt /= totweight / float(nc)**3
for i in range(len(positions)):
lay = pm.decompose(positions[i])
mesh.paint(positions[i], mass=weights[i], layout=lay, hold=True)
return mesh
def simulate(comm, ns):
pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm)
gaussian = pm.generate_whitenoise(ns.seed, unitary=True)
time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True)
Q = pm.generate_uniform_particle_grid(shift=0)
print(Q.min(axis=0), Q.max(axis=0))
def convolve(k, v):
kmag = sum(ki**2 for ki in k) ** 0.5
ampl = (PowerSpectrum(kmag) / v.BoxSize.prod()) ** 0.5
return v * ampl
dlinear = gaussian.apply(convolve)
DX1 = numpy.zeros_like(Q)
layout = pm.decompose(Q)
# Fill it in one dimension at a time.
for d in range(pm.ndim):
DX1[..., d] = dlinear \
.apply(dx1_transfer(d)) \
.c2r().readout(Q, layout=layout)
a0 = time_steps[0]
# 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step.
S = DX1 * pt.D1(a=a0)
V = S * a0 ** 2 * pt.f1(a0) * pt.E(a0)
state = State(Q, S, V)
fpm = ParticleMesh(BoxSize=pm.BoxSize, Nmesh=pm.Nmesh * ns.boost, resampler='tsc', dtype='f8')
ns.scheme(fpm, state, time_steps, ns.factors)
r = Result()
r.Q = Q
r.DX1 = DX1
r.S = S
r.V = V
r.dlinear = dlinear
return r
def test_paint_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, type='real')
def objective(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
obj = (real[...] ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(pos, mass, layout, v_pos=None, v_mass=None):
real = pm.paint(pos, mass=mass, layout=layout)
jvp = pm.paint_jvp(pos, mass=mass, v_mass=v_mass, v_pos=v_pos, layout=layout)
return comm.allreduce((jvp * real * 2)[...].sum())
def backward_gradient(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
return pm.paint_vjp(real * 2, pos, mass=mass, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
numpy.random.seed(9999)
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += (numpy.random.uniform(size=pos.shape)) * 0.8 + 0.1
pos *= pm.BoxSize / pm.Nmesh
mass = numpy.ones(len(pos)) * 2
layout = pm.decompose(pos)
obj = objective(pos, mass, layout)
grad_pos, grad_mass = backward_gradient(pos, mass, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(real.csize):
dx1, mass1, old = perturb_mass(mass, ind1[0], dx, comm)
ng1 = (objective(pos, mass1, layout) - obj)
bag1 = get_mass(grad_mass, ind1[0], comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_mass=mass1 - mass)
# print (dx1, dx, ind1, fag1, bag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(pos1, mass, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'v_pos', (pos1 - pos)[ind1[0]], 'f*g', fag1, 'bag', bag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
class Subsample(Algorithm):
"""
Algorithm to create a subsample from a DataSource, and evaluate
the density (1 + delta), smoothed at the given scale
"""
plugin_name = "Subsample"
def __init__(self, datasource, Nmesh, seed=12345, ratio=0.01, smoothing=None, format="hdf5"):
from pmesh.pm import ParticleMesh
self.datasource = datasource
self.Nmesh = Nmesh
self.seed = seed
self.ratio = ratio
self.smoothing = smoothing
self.format = format
self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype="f4", comm=self.comm)
@classmethod
def fill_schema(cls):
s = cls.schema
s.description = "create a subsample from a DataSource, and evaluate density \n"
s.description += "(1 + delta) smoothed at the given scale"
s.add_argument(
"datasource",
type=DataSource.from_config,
help="the DataSource to read; run `nbkit.py --list-datasources` for all options",
)
s.add_argument("Nmesh", type=int, help="the size of FFT mesh for painting")
s.add_argument("seed", type=int, help="the random seed")
s.add_argument("ratio", type=float, help="the fraction of particles to keep")
s.add_argument(
"smoothing",
type=float,
help="the smoothing length in distance units. "
"It has to be greater than the mesh resolution. "
"Otherwise the code will die. Default is the mesh resolution.",
)
# this is for output..
s.add_argument("format", choices=["hdf5", "mwhite"], help="the format of the output")
def run(self):
"""
Run the Subsample algorithm
"""
import mpsort
from astropy.utils.misc import NumpyRNGContext
if self.smoothing is None:
self.smoothing = self.datasource.BoxSize[0] / self.Nmesh[0]
elif (self.datasource.BoxSize / self.Nmesh > self.smoothing).any():
raise ValueError("smoothing is too small")
def Smoothing(pm, complex):
k = pm.k
k2 = 0
for ki in k:
ki2 = ki ** 2
complex[:] *= numpy.exp(-0.5 * ki2 * self.smoothing ** 2)
def NormalizeDC(pm, complex):
""" removes the DC amplitude. This effectively
divides by the mean
"""
w = pm.w
comm = pm.comm
ind = []
value = 0.0
found = True
for wi in w:
if (wi != 0).all():
found = False
break
ind.append((wi == 0).nonzero()[0][0])
if found:
ind = tuple(ind)
value = numpy.abs(complex[ind])
value = comm.allreduce(value)
complex[:] /= value
# open the datasource and keep the cache
with self.datasource.keep_cache():
painter = Painter.create("DefaultPainter", paintbrush="cic")
real, stats = painter.paint(self.pm, self.datasource)
complex = real.r2c()
for t in [Smoothing, NormalizeDC]:
t(self.pm, complex)
complex.c2r(real)
columns = ["Position", "ID", "Velocity"]
local_seed = utils.local_random_seed(self.seed, self.comm)
dtype = numpy.dtype([("Position", ("f4", 3)), ("Velocity", ("f4", 3)), ("ID", "u8"), ("Density", "f4")])
subsample = [numpy.empty(0, dtype=dtype)]
with self.datasource.open() as stream:
for Position, ID, Velocity in stream.read(columns):
with NumpyRNGContext(local_seed):
u = numpy.random.uniform(size=len(ID))
keep = u < self.ratio
Nkeep = keep.sum()
if Nkeep == 0:
continue
data = numpy.empty(Nkeep, dtype=dtype)
data["Position"][:] = Position[keep]
data["Velocity"][:] = Velocity[keep]
data["ID"][:] = ID[keep]
layout = self.pm.decompose(data["Position"])
pos1 = layout.exchange(data["Position"])
density1 = real.readout(pos1)
density = layout.gather(density1)
# normalize the position after reading out density!
data["Position"][:] /= self.datasource.BoxSize
data["Velocity"][:] /= self.datasource.BoxSize
data["Density"][:] = density
subsample.append(data)
subsample = numpy.concatenate(subsample)
mpsort.sort(subsample, orderby="ID", comm=self.comm)
return subsample
def save(self, output, data):
if self.format == "mwhite":
self.write_mwhite_subsample(data, output)
else:
self.write_hdf5(data, output)
def write_hdf5(self, subsample, output):
import h5py
size = self.comm.allreduce(len(subsample))
offset = sum(self.comm.allgather(len(subsample))[: self.comm.rank])
if self.comm.rank == 0:
with h5py.File(output, "w") as ff:
dataset = ff.create_dataset(name="Subsample", dtype=subsample.dtype, shape=(size,))
dataset.attrs["Ratio"] = self.ratio
dataset.attrs["CommSize"] = self.comm.size
dataset.attrs["Seed"] = self.seed
dataset.attrs["Smoothing"] = self.smoothing
dataset.attrs["Nmesh"] = self.Nmesh
dataset.attrs["Original"] = self.datasource.string
dataset.attrs["BoxSize"] = self.datasource.BoxSize
for i in range(self.comm.size):
self.comm.barrier()
if i != self.comm.rank:
continue
with h5py.File(output, "r+") as ff:
dataset = ff["Subsample"]
dataset[offset : len(subsample) + offset] = subsample
def write_mwhite_subsample(self, subsample, output):
size = self.comm.allreduce(len(subsample))
offset = sum(self.comm.allgather(len(subsample))[: self.comm.rank])
if self.comm.rank == 0:
with open(output, "wb") as ff:
dtype = numpy.dtype(
[
("eflag", "int32"),
("hsize", "int32"),
("npart", "int32"),
("nsph", "int32"),
("nstar", "int32"),
("aa", "float"),
("gravsmooth", "float"),
]
)
header = numpy.zeros((), dtype=dtype)
header["eflag"] = 1
header["hsize"] = 20
header["npart"] = size
header.tofile(ff)
self.comm.barrier()
with open(output, "r+b") as ff:
ff.seek(28 + offset * 12)
numpy.float32(subsample["Position"]).tofile(ff)
ff.seek(28 + offset * 12 + size * 12)
numpy.float32(subsample["Velocity"]).tofile(ff)
ff.seek(28 + offset * 4 + size * 24)
numpy.float32(subsample["Density"]).tofile(ff)
ff.seek(28 + offset * 8 + size * 28)
numpy.uint64(subsample["ID"]).tofile(ff)
def test_readout_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, type='real')
def objective(real, pos, layout):
value = real.readout(pos, layout=layout)
obj = (value ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(real, pos, layout, v_real=None, v_pos=None):
value = real.readout(pos, layout=layout)
v_value = real.readout_jvp(pos, v_self=v_real, v_pos=v_pos, layout=layout)
return comm.allreduce((v_value * value * 2).sum())
def backward_gradient(real, pos, layout):
value = real.readout(pos, layout=layout)
return real.readout_vjp(pos, v=value * 2, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += 0.5
pos *= pm.BoxSize / pm.Nmesh
layout = pm.decompose(pos)
obj = objective(real, pos, layout)
grad_real, grad_pos = backward_gradient(real, pos, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(*grad_real.cshape):
dx1, r1 = perturb(real, ind1, dx)
ng1 = (objective(r1, pos, layout) - obj)
bag1 = grad_real.cgetitem(ind1) * dx1
fag1 = forward_gradient(real, pos, layout, v_real=r1 - real)
# print (dx1, dx, ind1, ag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(real, pos1, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(real, pos, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'a', ag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
FieldMesh(LDLmap).save(save)
elif args.target == 'kSZ':
LDLmap *= 1e-5 #The learned LDL map is actually 1e5*map_ne.
save = args.save + '/ne_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n)
FieldMesh(LDLmap).save(save)
if args.FastPMpath:
X = File(args.FastPMpath)['Position']
X = np.array(X[start:end]).astype(np.float32)
Vz = File(args.FastPMpath)['Velocity']
Vz = np.array(Vz[start:end])[:,2].astype(np.float32)
else:
from readTNG import scalefactor
a = scalefactor(args.TNGDarkpath, args.snapNum)
Vz = load_TNG_data(TNG_basepath=args.TNGDarkpath, snapNum=args.snapNum, partType='dm', field='Velocities', mdi=2) * a**0.5
layout = pm.decompose(X)
X = layout.exchange(X)
Vz = layout.exchange(Vz)
map_vz = pm.create(type="real")
map_vz.paint(X, mass=Vz, layout=None, hold=False)
map_delta = pm.create(type="real")
map_delta.paint(X, mass=1., layout=None, hold=False)
select = map_delta > 0
map_vz[select] = map_vz[select] / map_delta[select]
map_vz[~select] = 0
LDLmap = LDLmap * map_vz
save = args.save + '/' + args.target + '_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n)
FieldMesh(LDLmap).save(save)
elif args.target in ['tSZ_T', 'Xray_T']:
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)
#dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/'%(bs, nc)
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
class TidalTensor(Algorithm):
"""
Compute and save the tidal force tensor
"""
plugin_name = "TidalTensor"
def __init__(self, field, points, Nmesh, smoothing=None):
from pmesh.pm import ParticleMesh
self.field = field
self.points = points
self.Nmesh = Nmesh
self.smoothing = smoothing
self.pm = ParticleMesh(BoxSize=self.field.BoxSize,
Nmesh=[self.Nmesh] * 3,
dtype='f4',
comm=self.comm)
@classmethod
def fill_schema(cls):
s = cls.schema
s.description = "compute the tidal force tensor"
s.add_argument(
"field",
type=DataSource.from_config,
help="DataSource; run `nbkit.py --list-datasources` for all options"
)
s.add_argument(
"points",
type=DataSource.from_config,
help="A small set of points to calculate tidal force on; "
"run `nbkit.py --list-datasources` for all options")
s.add_argument("Nmesh", type=int, help='Size of FFT mesh for painting')
s.add_argument(
"smoothing",
type=float,
help='Smoothing Length in distance units. '
'It has to be greater than the mesh resolution. '
'Otherwise the code will die. Default is the mesh resolution.')
def Smoothing(self, pm, complex):
k = pm.k
k2 = 0
for ki in k:
ki2 = ki**2
complex *= numpy.exp(-0.5 * ki2 * self.smoothing**2)
def NormalizeDC(self, pm, complex):
""" removes the DC amplitude. This effectively
divides by the mean
"""
w = pm.w
comm = pm.comm
ind = []
value = 0.0
found = True
for wi in w:
if (wi != 0).all():
found = False
break
ind.append((wi == 0).nonzero()[0][0])
if found:
ind = tuple(ind)
value = numpy.abs(complex[ind])
value = comm.allreduce(value)
complex[:] /= value
def TidalTensor(self, u, v):
# k_u k_v / k **2
def TidalTensor(pm, complex):
k = pm.k
for row in range(complex.shape[0]):
k2 = k[0][row]**2
for ki in k[1:]:
k2 = k2 + ki[0]**2
k2[k2 == 0] = numpy.inf
complex[row] /= k2
complex *= k[u]
complex *= k[v]
return TidalTensor
def run(self):
"""
Run the TidalTensor Algorithm
"""
from itertools import product
if self.smoothing is None:
self.smoothing = self.field.BoxSize[0] / self.Nmesh
elif (self.field.BoxSize / self.Nmesh > self.smoothing).any():
raise ValueError("smoothing is too small")
painter = Painter.create("DefaultPainter",
weight="Mass",
paintbrush="cic")
real, stats = painter.paint(self.pm, self.field)
complex = real.r2c()
for t in [self.Smoothing, self.NormalizeDC]:
t(complex.pm, complex)
with self.points.open() as stream:
[[Position]] = stream.read(['Position'], full=True)
layout = self.pm.decompose(Position)
pos1 = layout.exchange(Position)
value = numpy.empty((3, 3, len(Position)))
for u, v in product(range(3), range(3)):
if self.comm.rank == 0:
self.logger.info("Working on tensor element (%d, %d)" % (u, v))
c2 = complex.copy()
self.TidalTensor(u, v)(c2.pm, c2)
c2.c2r(real)
v1 = real.readout(pos1)
v1 = layout.gather(v1)
value[u, v] = v1
return value.transpose((2, 0, 1))
def save(self, output, data):
self.write_hdf5(data, output)
def write_hdf5(self, data, output):
import h5py
size = self.comm.allreduce(len(data))
offset = sum(self.comm.allgather(len(data))[:self.comm.rank])
if self.comm.rank == 0:
with h5py.File(output, 'w') as ff:
dataset = ff.create_dataset(name='TidalTensor',
dtype=data.dtype,
shape=(size, 3, 3))
dataset.attrs['Smoothing'] = self.smoothing
dataset.attrs['Nmesh'] = self.Nmesh
dataset.attrs['Original'] = self.field.string
dataset.attrs['BoxSize'] = self.field.BoxSize
for i in range(self.comm.size):
self.comm.barrier()
if i != self.comm.rank: continue
with h5py.File(output, 'r+') as ff:
dataset = ff['TidalTensor']
dataset[offset:len(data) + offset] = data
def test_paint_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, mode='real')
def objective(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
obj = (real[...] ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(pos, mass, layout, v_pos=None, v_mass=None):
real = pm.paint(pos, mass=mass, layout=layout)
jvp = pm.paint_jvp(pos, mass=mass, v_mass=v_mass, v_pos=v_pos, layout=layout)
return comm.allreduce((jvp * real * 2)[...].sum())
def backward_gradient(pos, mass, layout):
real = pm.paint(pos, mass=mass, layout=layout)
return pm.paint_vjp(real * 2, pos, mass=mass, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
numpy.random.seed(9999)
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += (numpy.random.uniform(size=pos.shape)) * 0.8 + 0.1
pos *= pm.BoxSize / pm.Nmesh
mass = numpy.ones(len(pos)) * 2
layout = pm.decompose(pos)
obj = objective(pos, mass, layout)
grad_pos, grad_mass = backward_gradient(pos, mass, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(real.csize):
dx1, mass1, old = perturb_mass(mass, ind1[0], dx, comm)
ng1 = (objective(pos, mass1, layout) - obj)
bag1 = get_mass(grad_mass, ind1[0], comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_mass=mass1 - mass)
# print (dx1, dx, ind1, fag1, bag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(pos1, mass, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(pos, mass, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'v_pos', (pos1 - pos)[ind1[0]], 'f*g', fag1, 'bag', bag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
def make_component_spectra(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
try:
os.makedirs(opath)
except:
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8')
rank = pm.comm.rank
header = '1,b1,b2,bg,bk'
names = header.split(',')
dfacs = [1, dgrow, dgrow**2, dgrow**2, dgrow]
#
dyn = BigFileCatalog(spath + '/fastpm_%0.4f/1' % aa)
fpos = dyn['Position'].compute()
idd = dyn['ID']
attrs = dyn.attrs
play = pm.decompose(fpos)
#
wts = BigFileCatalog(opath + '/lagweights/')
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=dfacs[i] * wts[names[i]], layout=play))
print(rank, names[i], eul_fields[-1].cmean())
del dyn, fpos, idd, iddg, wts
print('got fields')
iz = int(100 * zz + 0.01)
k, spectra = tools.getspectra(eul_fields)
np.savetxt(opath + '/spectra-z%03d-R%d.txt' % (iz, Rsm),
np.vstack([k, spectra]).T.real,
header='k / ' + header,
fmt='%0.4e')
rank = pm.comm.rank
lin = BigFileMesh(dpath + '/mesh', 's').paint()
dyn = BigFileCatalog(dpath + '/dynamic/1')
hcat = BigFileCatalog(hpath + '/FOF/')
#
hpos = hcat['CMPosition'][1:num]
#print('Mass : ', rank, hcat['Mass'].compute())[1:num]
hmass = hcat['Mass'].compute()[1:num]
if not masswt:
hmass = hmass * 0 + 1.
suff = suff + '-pos'
else:
suff = suff + '-mass'
hlay = pm.decompose(hpos)
hmesh = pm.paint(hpos, mass=hmass, layout=hlay)
hmesh /= hmesh.cmean()
ph = FFTPower(hmesh, mode='1d').power
k, ph = ph['k'], ph['power']
#
grid = dyn['InitPosition'].compute()
fpos = dyn['Position'].compute()
print(rank, (grid - fpos).std(axis=0))
#
dgrow = cosmo.scale_independent_growth_factor(zz)
if zadisp:
fpos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow)
suff = suff + '-za'
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)
try:
os.makedirs(ofolder)
except:
pass
pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8')
rank = pm.comm.rank
header = '1,b1,b2,bg,bk'
names = header.split(',')
#
dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1' % aa)
fpos = dyn['Position'].compute()
idd = dyn['ID']
attrs = dyn.attrs
play = pm.decompose(fpos)
#
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())
dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/20stepT-B1/%d-%d-9100/' % (
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(
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')
class TidalTensor(Algorithm):
"""
Compute and save the tidal force tensor
"""
plugin_name = "TidalTensor"
def __init__(self, field, points, Nmesh, smoothing=None):
from pmesh.pm import ParticleMesh
self.field = field
self.points = points
self.Nmesh = Nmesh
self.smoothing = smoothing
self.pm = ParticleMesh(BoxSize=self.field.BoxSize, Nmesh=[self.Nmesh]*3, dtype='f4', comm=self.comm)
@classmethod
def fill_schema(cls):
s = cls.schema
s.description = "compute the tidal force tensor"
s.add_argument("field", type=DataSource.from_config,
help="DataSource; run `nbkit.py --list-datasources` for all options")
s.add_argument("points", type=DataSource.from_config,
help="A small set of points to calculate tidal force on; "
"run `nbkit.py --list-datasources` for all options")
s.add_argument("Nmesh", type=int,
help='Size of FFT mesh for painting')
s.add_argument("smoothing", type=float,
help='Smoothing Length in distance units. '
'It has to be greater than the mesh resolution. '
'Otherwise the code will die. Default is the mesh resolution.')
def Smoothing(self, pm, complex):
k = pm.k
k2 = 0
for ki in k:
ki2 = ki ** 2
complex *= numpy.exp(-0.5 * ki2 * self.smoothing ** 2)
def NormalizeDC(self, pm, complex):
""" removes the DC amplitude. This effectively
divides by the mean
"""
w = pm.w
comm = pm.comm
ind = []
value = 0.0
found = True
for wi in w:
if (wi != 0).all():
found = False
break
ind.append((wi == 0).nonzero()[0][0])
if found:
ind = tuple(ind)
value = numpy.abs(complex[ind])
value = comm.allreduce(value)
complex[:] /= value
def TidalTensor(self, u, v):
# k_u k_v / k **2
def TidalTensor(pm, complex):
k = pm.k
for row in range(complex.shape[0]):
k2 = k[0][row] ** 2
for ki in k[1:]:
k2 = k2 + ki[0] ** 2
k2[k2 == 0] = numpy.inf
complex[row] /= k2
complex *= k[u]
complex *= k[v]
return TidalTensor
def run(self):
"""
Run the TidalTensor Algorithm
"""
from itertools import product
if self.smoothing is None:
self.smoothing = self.field.BoxSize[0] / self.Nmesh
elif (self.field.BoxSize / self.Nmesh > self.smoothing).any():
raise ValueError("smoothing is too small")
painter = Painter.create("DefaultPainter", weight="Mass", paintbrush="cic")
real, stats = painter.paint(self.pm, self.field)
complex = real.r2c()
for t in [self.Smoothing, self.NormalizeDC]:
t(complex.pm, complex)
with self.points.open() as stream:
[[Position ]] = stream.read(['Position'], full=True)
layout = self.pm.decompose(Position)
pos1 = layout.exchange(Position)
value = numpy.empty((3, 3, len(Position)))
for u, v in product(range(3), range(3)):
if self.comm.rank == 0:
self.logger.info("Working on tensor element (%d, %d)" % (u, v))
c2 = complex.copy()
self.TidalTensor(u, v)(c2.pm, c2)
c2.c2r(real)
v1 = real.readout(pos1)
v1 = layout.gather(v1)
value[u, v] = v1
return value.transpose((2, 0, 1))
def save(self, output, data):
self.write_hdf5(data, output)
def write_hdf5(self, data, output):
import h5py
size = self.comm.allreduce(len(data))
offset = sum(self.comm.allgather(len(data))[:self.comm.rank])
if self.comm.rank == 0:
with h5py.File(output, 'w') as ff:
dataset = ff.create_dataset(name='TidalTensor',
dtype=data.dtype, shape=(size, 3, 3))
dataset.attrs['Smoothing'] = self.smoothing
dataset.attrs['Nmesh'] = self.Nmesh
dataset.attrs['Original'] = self.field.string
dataset.attrs['BoxSize'] = self.field.BoxSize
for i in range(self.comm.size):
self.comm.barrier()
if i != self.comm.rank: continue
with h5py.File(output, 'r+') as ff:
dataset = ff['TidalTensor']
dataset[offset:len(data) + offset] = data
wsize = pm.comm.size
comm = pm.comm
newcomm = comm.Split(color=rank//splits, key=rank)
if rank==0:
print(args)
print(outfolder)
cube_size = nc/ncube
shift = cube_size
cube_length = cube_size*bs/nc
#pmsmall = ParticleMesh(BoxSize = bs/ncube, Nmesh = [cube_size, cube_size, cube_size], dtype=np.float32, comm=MPI.COMM_SELF)
pmsmall = ParticleMesh(BoxSize = bs/ncube, Nmesh = [cube_size, cube_size, cube_size], dtype=np.float32, comm=newcomm)
gridsmall = pmsmall.generate_uniform_particle_grid(shift=0)
layoutsmall = pmsmall.decompose(gridsmall)
for zz in [3.5, 4.0, 5.0, 6.0][::-1]:
aa = 1/(1+zz)
cats, meshes = setupuvmesh(zz, suff=suff, sim=sim, pm=pm, profile=profile, stellar=stellar, stellarfluc=stellarfluc, Rg=R)
cencat, satcat = cats
h1meshfid, h1mesh, lmesh, uvmesh = meshes
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()
meshes = [dm, h1meshfid, h1mesh, lmesh, uvmesh]
names = ['dm', 'h1fid', 'h1new', 'lum', 'uv']
if rank == 0 : print('\nMeshes read\n')
satcat['blackhole'] = mbh(moster(satcat['Mass'].compute(),
z=zz,
scatter=0.3),
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()
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')
indices[bindex][0], indices[bindex][1], indices[bindex][
2] = bi, bj, bk
bi *= cube_length
bj *= cube_length
bk *= cube_length
print('For rank & index : ', rank, bindex, '-- x, y, z : ', bi,
bj, bk)
poslarge = gridsmall + np.array([bi, bj, bk])
save = True
except IndexError:
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()))
def test_readout_gradients(comm):
pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic')
real = pm.generate_whitenoise(1234, mode='real')
def objective(real, pos, layout):
value = real.readout(pos, layout=layout)
obj = (value ** 2).sum()
return comm.allreduce(obj)
def forward_gradient(real, pos, layout, v_real=None, v_pos=None):
value = real.readout(pos, layout=layout)
v_value = real.readout_jvp(pos, v_self=v_real, v_pos=v_pos, layout=layout)
return comm.allreduce((v_value * value * 2).sum())
def backward_gradient(real, pos, layout):
value = real.readout(pos, layout=layout)
return real.readout_vjp(pos, v=value * 2, layout=layout)
pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T
pos += real.start
# avoid sitting at the pmesh points
# cic gradient is zero on them, the numerical gradient fails.
pos += 0.5
pos *= pm.BoxSize / pm.Nmesh
layout = pm.decompose(pos)
obj = objective(real, pos, layout)
grad_real, grad_pos = backward_gradient(real, pos, layout)
ng = []
f*g = []
bag = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex(*grad_real.cshape):
dx1, r1 = perturb(real, ind1, dx)
ng1 = (objective(r1, pos, layout) - obj)
bag1 = grad_real.cgetitem(ind1) * dx1
fag1 = forward_gradient(real, pos, layout, v_real=r1 - real)
# print (dx1, dx, ind1, ag1, ng1)
ng.append(ng1)
f*g.append(fag1)
bag.append(bag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
# FIXME
ng = []
bag = []
f*g = []
ind = []
dx = 1e-6
for ind1 in numpy.ndindex((real.csize, real.ndim)):
dx1, pos1, old = perturb_pos(pos, ind1, dx, comm)
layout1 = pm.decompose(pos1)
ng1 = (objective(real, pos1, layout1) - obj)
bag1 = get_pos(grad_pos, ind1, comm) * dx1
fag1 = forward_gradient(real, pos, layout, v_pos=pos1 - pos)
# print ('pos', old, ind1, 'a', ag1, 'n', ng1)
ng.append(ng1)
bag.append(bag1)
f*g.append(fag1)
ind.append(ind1)
assert_allclose(bag, f*g, rtol=1e-7)
assert_allclose(ng, bag, rtol=1e-4)
class Subsample(Algorithm):
"""
Algorithm to create a subsample from a DataSource, and evaluate
the density (1 + delta), smoothed at the given scale
"""
plugin_name = "Subsample"
def __init__(self,
datasource,
Nmesh,
seed=12345,
ratio=0.01,
smoothing=None,
format='hdf5'):
from pmesh.pm import ParticleMesh
self.datasource = datasource
self.Nmesh = Nmesh
self.seed = seed
self.ratio = ratio
self.smoothing = smoothing
self.format = format
self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize,
Nmesh=[self.Nmesh] * 3,
dtype='f4',
comm=self.comm)
@classmethod
def fill_schema(cls):
s = cls.schema
s.description = "create a subsample from a DataSource, and evaluate density \n"
s.description += "(1 + delta) smoothed at the given scale"
s.add_argument(
"datasource",
type=DataSource.from_config,
help=
"the DataSource to read; run `nbkit.py --list-datasources` for all options"
)
s.add_argument("Nmesh",
type=int,
help='the size of FFT mesh for painting')
s.add_argument("seed", type=int, help='the random seed')
s.add_argument("ratio",
type=float,
help='the fraction of particles to keep')
s.add_argument(
"smoothing",
type=float,
help='the smoothing length in distance units. '
'It has to be greater than the mesh resolution. '
'Otherwise the code will die. Default is the mesh resolution.')
# this is for output..
s.add_argument("format",
choices=['hdf5', 'mwhite'],
help='the format of the output')
def run(self):
"""
Run the Subsample algorithm
"""
import mpsort
from astropy.utils.misc import NumpyRNGContext
if self.smoothing is None:
self.smoothing = self.datasource.BoxSize[0] / self.Nmesh[0]
elif (self.datasource.BoxSize / self.Nmesh > self.smoothing).any():
raise ValueError("smoothing is too small")
def Smoothing(pm, complex):
k = pm.k
k2 = 0
for ki in k:
ki2 = ki**2
complex[:] *= numpy.exp(-0.5 * ki2 * self.smoothing**2)
def NormalizeDC(pm, complex):
""" removes the DC amplitude. This effectively
divides by the mean
"""
w = pm.w
comm = pm.comm
ind = []
value = 0.0
found = True
for wi in w:
if (wi != 0).all():
found = False
break
ind.append((wi == 0).nonzero()[0][0])
if found:
ind = tuple(ind)
value = numpy.abs(complex[ind])
value = comm.allreduce(value)
complex[:] /= value
# open the datasource and keep the cache
with self.datasource.keep_cache():
painter = Painter.create("DefaultPainter", paintbrush='cic')
real, stats = painter.paint(self.pm, self.datasource)
complex = real.r2c()
for t in [Smoothing, NormalizeDC]:
t(self.pm, complex)
complex.c2r(real)
columns = ['Position', 'ID', 'Velocity']
local_seed = utils.local_random_seed(self.seed, self.comm)
dtype = numpy.dtype([
('Position', ('f4', 3)),
('Velocity', ('f4', 3)),
('ID', 'u8'),
('Density', 'f4'),
])
subsample = [numpy.empty(0, dtype=dtype)]
with self.datasource.open() as stream:
for Position, ID, Velocity in stream.read(columns):
with NumpyRNGContext(local_seed):
u = numpy.random.uniform(size=len(ID))
keep = u < self.ratio
Nkeep = keep.sum()
if Nkeep == 0: continue
data = numpy.empty(Nkeep, dtype=dtype)
data['Position'][:] = Position[keep]
data['Velocity'][:] = Velocity[keep]
data['ID'][:] = ID[keep]
layout = self.pm.decompose(data['Position'])
pos1 = layout.exchange(data['Position'])
density1 = real.readout(pos1)
density = layout.gather(density1)
# normalize the position after reading out density!
data['Position'][:] /= self.datasource.BoxSize
data['Velocity'][:] /= self.datasource.BoxSize
data['Density'][:] = density
subsample.append(data)
subsample = numpy.concatenate(subsample)
mpsort.sort(subsample, orderby='ID', comm=self.comm)
return subsample
def save(self, output, data):
if self.format == 'mwhite':
self.write_mwhite_subsample(data, output)
else:
self.write_hdf5(data, output)
def write_hdf5(self, subsample, output):
import h5py
size = self.comm.allreduce(len(subsample))
offset = sum(self.comm.allgather(len(subsample))[:self.comm.rank])
if self.comm.rank == 0:
with h5py.File(output, 'w') as ff:
dataset = ff.create_dataset(name='Subsample',
dtype=subsample.dtype,
shape=(size, ))
dataset.attrs['Ratio'] = self.ratio
dataset.attrs['CommSize'] = self.comm.size
dataset.attrs['Seed'] = self.seed
dataset.attrs['Smoothing'] = self.smoothing
dataset.attrs['Nmesh'] = self.Nmesh
dataset.attrs['Original'] = self.datasource.string
dataset.attrs['BoxSize'] = self.datasource.BoxSize
for i in range(self.comm.size):
self.comm.barrier()
if i != self.comm.rank: continue
with h5py.File(output, 'r+') as ff:
dataset = ff['Subsample']
dataset[offset:len(subsample) + offset] = subsample
def write_mwhite_subsample(self, subsample, output):
size = self.comm.allreduce(len(subsample))
offset = sum(self.comm.allgather(len(subsample))[:self.comm.rank])
if self.comm.rank == 0:
with open(output, 'wb') as ff:
dtype = numpy.dtype([('eflag', 'int32'), ('hsize', 'int32'),
('npart', 'int32'), ('nsph', 'int32'),
('nstar', 'int32'), ('aa', 'float'),
('gravsmooth', 'float')])
header = numpy.zeros((), dtype=dtype)
header['eflag'] = 1
header['hsize'] = 20
header['npart'] = size
header.tofile(ff)
self.comm.barrier()
with open(output, 'r+b') as ff:
ff.seek(28 + offset * 12)
numpy.float32(subsample['Position']).tofile(ff)
ff.seek(28 + offset * 12 + size * 12)
numpy.float32(subsample['Velocity']).tofile(ff)
ff.seek(28 + offset * 4 + size * 24)
numpy.float32(subsample['Density']).tofile(ff)
ff.seek(28 + offset * 8 + size * 28)
numpy.uint64(subsample['ID']).tofile(ff)
def main():
ns = ap.parse_args()
comm = MPI.COMM_WORLD
ff = bigfile.BigFileMPI(comm, ns.fastpm)
with ff['.'] as bb:
BoxSize = bb.attrs['BoxSize'][0]
Redshift = 1 / bb.attrs['ScalingFactor'][0] - 1
Nmesh = int(BoxSize / ns.resolution * 2)
# round it to 8.
Nmesh -= Nmesh % 8
if comm.rank == 0:
logger.info("source = %s", ns.fastpm)
logger.info("output = %s", ns.output)
logger.info("BoxSize = %g", BoxSize)
logger.info("Redshift = %g", Redshift)
logger.info("Nmesh = %g", Nmesh)
pm = ParticleMesh([Nmesh, Nmesh, Nmesh], BoxSize, comm=comm)
real = RealField(pm)
real[...] = 0
with ff['Position'] as ds:
logger.info(ds.size)
for i in range(0, ds.size, ns.chunksize):
sl = slice(i, i + ns.chunksize)
pos = ds[sl]
layout = pm.decompose(pos)
lpos = layout.exchange(pos)
real.paint(lpos, hold=True)
mean = real.cmean()
if comm.rank == 0:
logger.info("mean particle per cell = %s", mean)
real[...] /= mean
real[...] -= 1
complex = real.r2c()
for k, i, slab in zip(complex.slabs.x, complex.slabs.i, complex.slabs):
k2 = sum(kd**2 for kd in k)
# tophat
f = tophat(ns.filtersize, k2**0.5)
slab[...] *= f
# zreion
slab[...] *= Bk(k2**0.5)
slab[...] *= (1 + Redshift)
real = complex.c2r()
real[...] += Redshift
mean = real.cmean()
if comm.rank == 0:
logger.info("zreion.mean = %s", mean)
buffer = numpy.empty(real.size, real.dtype)
real.sort(out=buffer)
if comm.rank == 0:
logger.info("sorted for output")
with bigfile.BigFileMPI(comm, ns.output, create=True) as ff:
with ff.create_from_array(ns.dataset, buffer) as bb:
bb.attrs['BoxSize'] = BoxSize
bb.attrs['Redshift'] = Redshift
bb.attrs['TopHatFilterSize'] = ns.filtersize
bb.attrs['Nmesh'] = Nmesh
#
# hack: compatible with current MPGadget. This is not really needed
# we'll remove the bins later, since BoxSize and Nmesh are known.
with ff.create("XYZ_bins", dtype='f8', size=Nmesh) as bb:
if comm.rank == 0:
bins = numpy.linspace(0, BoxSize * 1000., Nmesh, dtype='f8')
bb.write(0, bins)
if comm.rank == 0:
logger.info("done. written at %s", ns.output)
