def to_real_field(self):
"""
Return the RealField stored on disk.
.. note::
The mesh stored on disk must be stored with ``mode=real``
Returns
-------
real : pmesh.pm.RealField
an array-like object holding the mesh loaded from disk in
configuration space
"""
if self.isfourier:
return NotImplemented
# the real field to paint to
pmread = self.pm
with FileMPI(comm=self.comm, filename=self.path)[self.dataset] as ds:
if self.comm.rank == 0:
self.logger.info("reading real field from %s" % self.path)
real2 = RealField(pmread)
start = numpy.sum(self.comm.allgather(real2.size)[:self.comm.rank], dtype='intp')
end = start + real2.size
real2.unravel(ds[start:end])
return real2
def save(self, filename, attrs={}):
from bigfile import FileMPI
a = self.a['S']
with FileMPI(self.pm.comm, filename, create=True) as ff:
with ff.create('Header') as bb:
keylist = ['Om0', 'Tcmb0', 'Neff', 'Ob0', 'Ode0']
if getattr(self.cosmology, 'm_nu', None) is not None:
keylist.insert(3,'m_nu')
for key in keylist:
bb.attrs[key] = getattr(self.cosmology, key)
bb.attrs['Time'] = a
bb.attrs['h'] = self.cosmology.H0 / self.H0 # relative h
bb.attrs['RSDFactor'] = 1.0 / (self.H0 * a * self.cosmology.efunc(1.0 / a - 1))
for key in attrs:
try:
#best effort
bb.attrs[key] = attrs[key]
except:
pass
ff.create_from_array('1/Position', self.X)
# Peculiar velocity in km/s
ff.create_from_array('1/Velocity', self.V)
# dimensionless potential (check this)
ff.create_from_array('1/Density', self.RHO)
def to_complex_field(self):
"""
Return the ComplexField stored on disk.
.. note::
The mesh stored on disk must be stored with ``mode=complex``
Returns
-------
real : pmesh.pm.ComplexField
an array-like object holding the mesh loaded from disk in Fourier
space
"""
if not self.isfourier:
return NotImplemented
pmread = self.pm
if self.comm.rank == 0:
self.logger.info("reading complex field from %s" % self.path)
with FileMPI(comm=self.comm, filename=self.path)[self.dataset] as ds:
complex2 = ComplexField(pmread)
assert self.comm.allreduce(complex2.size) == ds.size
start = numpy.sum(self.comm.allgather(complex2.size)[:self.comm.rank], dtype='intp')
end = start + complex2.size
complex2.unravel(ds[start:end])
return complex2
def save(self, filename, dataset):
from bigfile import FileMPI
with FileMPI(self.comm, filename, create=True) as ff:
ff.create_from_array(dataset + '/Position', self.X)
# Peculiar velocity in km/s
ff.create_from_array(dataset + '/Velocity', self.V)
# dimensionless potential (check this)
ff.create_from_array(dataset + '/Density', self.RHO)
def saveIC_dmo(IC_path, dx_field, Lbox, Ng, cosmology, redshift=99):
"""
Use Zel-dovich approximation to back-scale the linear density field to initial redshift,
paint only dm particles from the grid,
and save initial condition in Mpc/h and use Constrained package to transfer to bt2 format
Parameters
---------
: dx_field : linear density field at z=0
: Lbox : BoxSize, in Mpc/h, will be converted to kpc/h in the IC output
: Ng : Number of grid on which to paint the particle
: cosmology : nbodykit.cosmology, or astropy.cosmology.FLRW
: redshift : redshift of the initial condition
"""
mesh = ArrayMesh(dx_field,
BoxSize=Lbox) # density contrast field centered at zero
dk_field = mesh.compute(mode='complex')
scale_a = 1. / (1 + redshift)
state = solver.lpt(dk_field, Q, a=scale_a, order=1)
with FileMPI(state.pm.comm, IC_path, create=True) as ff:
m1 = state.cosmology.rho_crit(0) * state.cosmology.Om0 * (
Lbox**3) / state.csize
with ff.create('Header') as bb:
bb.attrs['Time'] = state.a['S']
bb.attrs['HubbleParam'] = state.cosmology.h
bb.attrs['Omega0'] = state.cosmology.Om0
bb.attrs['OmegaM'] = state.cosmology.Omega0_m
bb.attrs['OmegaLambda'] = state.cosmology.Omega0_lambda
bb.attrs['BoxSize'] = Lbox
bb.attrs['NC'] = int(rint((state.csize)**(1 / 3)))
bb.attrs['TotNumPart'] = [0, state.csize, 0, 0, 0, 0]
bb.attrs['MassTable'] = [0, m1, 0, 0, 0, 0]
ff.create_from_array('1/Position', state.X)
# Peculiar velocity in km/s
ff.create_from_array('1/Velocity', state.V)
stateID = np.arange(state.csize)
ff.create_from_array('1/ID', stateID)
print("IC generated!")
print("*********************************************")
return state
def __init__(self, path, dataset, comm=None, **kwargs):
self.path = path
self.dataset = dataset
self.comm = comm
# update the meta-data
self.attrs.update(kwargs)
with FileMPI(comm=self.comm, filename=path)[dataset] as ff:
for key in ff.attrs:
v = ff.attrs[key]
if isinstance(v, string_types) and v.startswith('json://'):
self.attrs[key] = json.loads(v[7:], cls=JSONDecoder)
else:
self.attrs[key] = numpy.squeeze(v)
# fourier space or config space
if ff.dtype.kind == 'c':
self.isfourier = True
if ff.dtype.itemsize == 16:
dtype = 'f8'
else:
dtype = 'f4'
else:
self.isfourier = False
if ff.dtype.itemsize == 8:
dtype = 'f8'
else:
dtype = 'f4'
# determine Nmesh
if 'ndarray.shape' not in self.attrs:
raise ValueError("`ndarray.shape` should be stored in the Bigfile `attrs` to determine `Nmesh`")
if 'Nmesh' not in self.attrs:
raise ValueError("`ndarray.shape` should be stored in the Bigfile `attrs` to determine `Nmesh`")
Nmesh = self.attrs['Nmesh']
BoxSize = self.attrs['BoxSize']
MeshSource.__init__(self, BoxSize=BoxSize, Nmesh=Nmesh, dtype=dtype, comm=comm)
def save(self, filename, attrs={}):
from bigfile import FileMPI
a = self.a['S']
with FileMPI(self.comm, filename, create=True) as ff:
with ff.create('Header') as bb:
keylist = ['Om0', 'Tcmb0', 'Neff', 'Ob0', 'Ode0']
if getattr(self.cosmology, 'm_nu', None) is not None:
keylist.insert(3, 'm_nu')
for key in keylist:
bb.attrs[key] = getattr(self.cosmology, key)
bb.attrs['Time'] = a
bb.attrs['h'] = self.cosmology.H0 / H0 # relative h
bb.attrs['RSDFactor'] = 1.0 / (
H0 * a * self.cosmology.efunc(1.0 / a - 1))
for key in attrs:
try:
#best effort
bb.attrs[key] = attrs[key]
except:
pass
for k, v in sorted(self.species.items()):
v.save(filename, k)
def saveIC_bt2(IC_path, dx_field, Lbox, Ng, cosmology, redshift=99):
"""
Use Zel-dovich approximation to back-scale the linear density field to initial redshift,
paint baryon and dm particles from the grid,
and save initial condition in MP-Gadget/bluetides-ii format
Parameters
---------
: dx_field : linear density field at z=0
: Lbox : BoxSize, in Mpc/h, will be converted to kpc/h in the IC output
: Ng : Number of grid on which to paint the particle
: cosmology : nbodykit.cosmology, or astropy.cosmology.FLRW
: redshift : redshift of the initial condition
"""
mesh = ArrayMesh(dx_field,
BoxSize=Lbox) # density contrast field centered at zero
dk_field = mesh.compute(mode='complex')
shift_gas = -0.5 * (cosmology.Om0 - cosmology.Ob0) / cosmology.Om0
shift_dm = 0.5 * cosmology.Ob0 / cosmology.Om0
pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng, Ng, Ng])
solver = Solver(pm, cosmology, B=1)
Q_gas = pm.generate_uniform_particle_grid(shift=shift_gas)
Q_dm = pm.generate_uniform_particle_grid(shift=shift_dm)
scale_a = 1. / (1 + redshift)
state_gas = solver.lpt(dk_field, Q_gas, a=scale_a,
order=1) # order = 1 : Zel-dovich
state_dm = solver.lpt(dk_field, Q_dm, a=scale_a, order=1)
state = state_dm
with FileMPI(state.pm.comm, IC_path, create=True) as ff:
m0 = state.cosmology.rho_crit(0) * state.cosmology.Omega0_b * (
Lbox**3) / state.csize
m1 = state.cosmology.rho_crit(0) * (
state.cosmology.Om0 - state.cosmology.Omega0_b) * (Lbox**
3) / state.csize
with ff.create('Header') as bb:
bb.attrs['BoxSize'] = Lbox * 1000
bb.attrs['HubbleParam'] = state.cosmology.h
bb.attrs['MassTable'] = [m0, m1, 0, 0, 0, 0]
bb.attrs['OmegaM'] = state.cosmology.Om0
bb.attrs['OmegaB'] = state.cosmology.Omega0_b
bb.attrs['OmegaL'] = state.cosmology.Omega0_lambda
bb.attrs['Time'] = state.a['S']
bb.attrs['TotNumPart'] = [state.csize, state.csize, 0, 0, 0, 0]
ff.create_from_array('1/Position',
1000 * periodic_wrap(state.X, Lbox)) # in kpc/h
ff.create_from_array(
'1/Velocity',
state.V / np.sqrt(state.a['S'])) # old gadget convention for IC
dmID = np.arange(state.csize)
ff.create_from_array('1/ID', dmID)
#######################################################
ff.create_from_array('0/Position',
1000 * periodic_wrap(state_gas.X, Lbox))
ff.create_from_array('0/Velocity', state_gas.V / np.sqrt(state.a['S']))
gasID = np.arange(state.csize, 2 * state.csize)
ff.create_from_array('0/ID', gasID)
print("IC generated!")
print("*********************************************")
return state
state_gas = solver.lpt(dk_field, Q_gas, a=scale_a, order=1)
state_dm = solver.lpt(dk_field, Q_dm, a=scale_a, order=1)
# save state
# -----------------------------------------------------
def periodic_wrap(pos,Lbox):
pos[pos>Lbox] -= Lbox
pos[pos<0] += Lbox
return pos
IC_path = args.IC_path
state = state_dm
with FileMPI(state.pm.comm, IC_path, create=True) as ff:
m0 = state.cosmology.rho_crit(0)*state.cosmology.Omega0_b*(Lbox**3)/state.csize
m1 = state.cosmology.rho_crit(0)*(state.cosmology.Om0-state.cosmology.Omega0_b)*(Lbox**3)/state.csize
with ff.create('Header') as bb:
bb.attrs['BoxSize'] = Lbox*1000
bb.attrs['HubbleParam'] = state.cosmology.h
bb.attrs['MassTable'] = [m0, m1, 0, 0, 0, 0]
bb.attrs['Omega0'] = state.cosmology.Om0
bb.attrs['OmegaBaryon'] = state.cosmology.Omega0_b
bb.attrs['OmegaLambda'] = state.cosmology.Omega0_lambda
bb.attrs['Time'] = state.a['S']
bb.attrs['TotNumPart'] = [state.csize, state.csize, 0, 0, 0, 0]