def _cpui_level_iterator(cpu, amr_filename, bisection_order, maxlevel, ndim) :
f = file(amr_filename, 'rb')
header = read_fortran_series(f, ramses_amr_header)
skip_fortran(f, 13)
n_per_level = read_fortran(f, _int_type, header['nlevelmax']*header['ncpu']).reshape(( header['nlevelmax'], header['ncpu']))
skip_fortran(f,1)
if header['nboundary']>0 :
skip_fortran(f,2)
n_per_level_boundary = read_fortran(f, _int_type, header['nlevelmax']*header['nboundary']).reshape(( header['nlevelmax'], header['nboundary']))
skip_fortran(f,2)
if bisection_order :
skip_fortran(f, 5)
else :
skip_fortran(f, 1)
skip_fortran(f, 3)
offset = np.array(header['ng'],dtype='f8')/2
offset-=0.5
for level in xrange(maxlevel or header['nlevelmax']) :
# loop through those CPUs with grid data (includes ghost regions)
for cpuf in 1+np.where(n_per_level[level,:]!=0)[0] :
#print "CPU=",cpu,"CPU on disk=",cpuf,"npl=",n_per_level[level,cpuf-1]
if cpuf==cpu :
# this is the data we want
skip_fortran(f,3) # grid, next, prev index
# store the coordinates in temporary arrays. We only want
# to copy it if the cell is not refined
x0,y0,z0 = [read_fortran(f, _float_type, n_per_level[level,cpu-1]) for ar in range(ndim)]
skip_fortran(f,1 # father index
+ 2*ndim # nbor index
+ 2*(2**ndim) # son index,cpumap,refinement map
)
refine = np.array([read_fortran(f,_int_type,n_per_level[level,cpu-1]) for i in xrange(2**ndim)])
if level==maxlevel :
refine[:] = 0
x0-=offset[0]; y0-=offset[1]; z0-=offset[2]
yield (x0,y0,z0),refine,cpuf,level
else :
# skip ghost regions from other CPUs
skip_fortran(f,3+ndim+1+2*ndim+3*2**ndim)
if header['nboundary']>0 :
for boundaryf in np.where(n_per_level_boundary[level, :]!=0)[0] :
skip_fortran(f,3+ndim+1+2*ndim+3*2**ndim)
def __init__(self, f, take=None) :
super(GenICSnap,self).__init__()
f_cx = file(os.path.join(f, "ic_velcx"))
self._header = read_fortran(f_cx, genic_header)
h = self._header
self._dlen = int(h['nx']*h['ny'])
self.properties['a'] = float(h['astart'])
self.properties['h'] = float(h['h0'])/100.
self.properties['omegaM0'] = float(h['omegam'])
self.properties['omegaL0'] = float(h['omegal'])
disk_family_slice = { family.dm: slice(0, self._dlen*int(h['nz'])) }
self._load_control = chunk.LoadControl(disk_family_slice, _max_buflen, take)
self._family_slice = self._load_control.mem_family_slice
self._num_particles = self._load_control.mem_num_particles
self._filename = f
def _cpui_count_particles(filename) :
distinguisher_field = int(particle_distinguisher[0])
distinguisher_type = np.dtype(particle_distinguisher[1])
f = file(filename)
header = read_fortran_series(f, ramses_particle_header)
npart_this = header['npart']
try:
skip_fortran(f,distinguisher_field)
data = read_fortran(f,distinguisher_type,header['npart'])
except TypeError:
data = np.zeros(npart_this)
my_mask = (data!=0)
nstar_this = (data!=0).sum()
return npart_this, nstar_this, my_mask
def _derive_vel(self):
self._create_array("vel", 3)
vel = self["vel"]
vel.units = "km s^-1"
h = self._header
if self.properties["a"] != float(h["astart"]):
z0 = 1.0 / h["astart"] - 1
a_bdot_original = float(h["astart"]) * analysis.cosmology.rate_linear_growth(self, z=z0)
ratio = self.properties["a"] * analysis.cosmology.rate_linear_growth(self) / a_bdot_original
warnings.warn(
"You have manually changed the redshift of these initial conditions before loading velocities; the velocities will be scaled as appropriate",
RuntimeWarning,
)
else:
ratio = 1.0
for vd in "x", "y", "z":
vel = self["v" + vd]
f = file(os.path.join(self._filename, "ic_velc" + vd))
h = read_fortran(f, genic_header)
length = self._dlen * _data_type.itemsize
alen = np.fromfile(f, util._head_type, 1)
if alen != length:
raise IOError, "Unexpected FORTRAN block length %d!=%d" % (alen, length)
readpos = 0
for readlen, buf_index, mem_index in self._load_control.iterate_with_interrupts(
family.dm,
family.dm,
np.arange(1, h["nz"]) * (h["nx"] * h["ny"]),
functools.partial(_midway_fortran_skip, f, length),
):
if buf_index is not None:
re = np.fromfile(f, _data_type, readlen)
vel[mem_index] = re[buf_index] * ratio
else:
f.seek(_data_type.itemsize * readlen, 1)
alen = np.fromfile(f, util._head_type, 1)
if alen != length:
raise IOError, "Unexpected FORTRAN block length (tail) %d!=%d" % (alen, length)
def _cpui_load_particle_block(filename, dm_ar, star_ar, offset, ind0_dm, ind0_star, _type, star_mask, nstar) :
f = file(filename)
header = read_fortran_series(f, ramses_particle_header)
skip_fortran(f, offset)
ind1_dm = ind0_dm+header['npart']-nstar
ind1_star = ind0_star+nstar
data = read_fortran(f, _type, header['npart'])
if len(star_mask)>0 :
dm_ar[ind0_dm:ind1_dm]=data[~star_mask]
star_ar[ind0_star:ind1_star]=data[star_mask]
else :
dm_ar[ind0_dm:ind1_dm]=data
f.close()
def __init__(self, f, take=None):
super(GrafICSnap, self).__init__()
f_cx = file(os.path.join(f, "ic_velcx"))
self._header = read_fortran(f_cx, genic_header)
h = self._header
self._dlen = int(h["nx"] * h["ny"])
self.properties["a"] = float(h["astart"])
self.properties["h"] = float(h["h0"]) / 100.0
self.properties["omegaM0"] = float(h["omegam"])
self.properties["omegaL0"] = float(h["omegal"])
disk_family_slice = {family.dm: slice(0, self._dlen * int(h["nz"]))}
self._load_control = chunk.LoadControl(disk_family_slice, _max_buflen, take)
self._family_slice = self._load_control.mem_family_slice
self._num_particles = self._load_control.mem_num_particles
self._filename = f
boxsize = self._header["dx"] * self._header["nx"]
self.properties["boxsize"] = boxsize * units.Unit("Mpc a")
def _cpui_load_particle_block(filename, dm_ar, star_ar, offset, ind0_dm, ind0_star, _type, star_mask, nstar) :
f = file(filename)
header = read_fortran_series(f, ramses_particle_header)
skip_fortran(f, offset)
ind1_dm = ind0_dm+header['npart']-nstar
ind1_star = ind0_star+nstar
data = read_fortran(f, _type, header['npart'])
try:
if len(star_mask)>0 :
dm_ar[ind0_dm:ind1_dm]=data[~star_mask]
star_ar[ind0_star:ind1_star]=data[star_mask]
else :
dm_ar[ind0_dm:ind1_dm]=data
except ValueError:
# this translates into the data block loaded from disk not being
# long enough
raise IOError, "Could not load particle block"
f.close()
def _load_array(self, name, fam=None) :
if fam is not family.dm and fam is not None :
raise IOError, "Only DM particles supported"
if name=="mass" :
boxsize = self._header['dx']*self._header['nx']
rho = analysis.cosmology.rho_M(self, unit='Msol Mpc^-3 a^-3')
tot_mass = rho*boxsize**3 # in Msol
part_mass = tot_mass/self._header['nx']**3
self._create_array('mass')
self['mass'][:] = part_mass
self['mass'].units="Msol"
elif name=="pos" :
self._create_array('pos', 3)
self['pos'].units="Mpc a"
pos=self['pos']
nx,ny,nz = [int(self._header[x]) for x in 'nx','ny','nz']
# the following is equivalent to
#
# self['z'],self['y'],self['x'] = np.mgrid[0.0:self._header['nx'], 0.0:self._header['ny'], 0.0:self._header['nz']]
#
# but works on partial loading without having to generate the entire mgrid
# (which might easily exceed the available memory for a big grid)
pos_cache = np.empty((_max_buflen,3))
fp0 = 0
for readlen, buf_index, mem_index in self._load_control.iterate(family.dm, family.dm) :
if mem_index is not None :
pos[mem_index] = _grid_gen(slice(fp0, fp0+readlen), nx, ny, nz, pos=pos_cache)[buf_index]
fp0+=readlen
self['pos']*=self._header['dx']
a = self.properties['a']
bdot_by_b = analysis.cosmology.rate_linear_growth(self, unit='km Mpc^-1 s^-1')/analysis.cosmology.linear_growth_factor(self)
# offset position according to zeldovich approximation
self['offset'] = self['vel']/(a*bdot_by_b)
self['offset'].units=self['vel'].units/units.Unit('km Mpc^-1 s^-1 a^-1')
self['pos']+=self['offset']
elif name=="vel" :
self._create_array('vel', 3)
vel = self['vel']
vel.units='km s^-1'
h = self._header
if self.properties['a']!=float(h['astart']) :
z0 = 1./h['astart']-1
a_bdot_original = (float(h['astart']) * analysis.cosmology.rate_linear_growth(self, z=z0))
ratio = self.properties['a'] * analysis.cosmology.rate_linear_growth(self) / a_bdot_original
warnings.warn("You have manually changed the redshift of these initial conditions before loading velocities; the velocities will be scaled as appropriate", RuntimeWarning)
else :
ratio = 1.0
for vd in 'x','y','z' :
vel = self['v'+vd]
f = file(os.path.join(self._filename, 'ic_velc'+vd))
h = read_fortran(f, genic_header)
length = self._dlen * _data_type.itemsize
alen = np.fromfile(f, util._head_type, 1)
if alen!=length :
raise IOError, "Unexpected FORTRAN block length %d!=%d"%(alen,length)
readpos = 0
for readlen, buf_index, mem_index in (self._load_control.iterate_with_interrupts(family.dm, family.dm,
np.arange(1,h['nz'])*(h['nx']*h['ny']),
functools.partial(_midway_fortran_skip, f, length))) :
if buf_index is not None :
re = np.fromfile(f, _data_type ,readlen)
vel[mem_index] = re[buf_index]*ratio
else :
f.seek(_data_type.itemsize*readlen, 1)
alen = np.fromfile(f, util._head_type, 1)
if alen!=length :
raise IOError, "Unexpected FORTRAN block length (tail) %d!=%d"%(alen,length)
else :
raise IOError, "No such array"
def _cpui_load_gas_vars(dims, maxlevel, ndim, filename, cpu, lia, i1,
mode = _gv_load_hydro ) :
if config['verbose'] :
print>>sys.stderr, cpu,
sys.stderr.flush()
nvar = len(dims)
grid_info_iter = _cpui_level_iterator(*lia)
f = file(filename)
check_nvar_file = False
if mode is _gv_load_hydro :
header = read_fortran_series(f, ramses_hydro_header)
nvar_file = header['nvarh']
else :
header = read_fortran_series(f, ramses_grav_header)
nvar_file=4
if nvar_file<nvar :
warnings.warn("Fewer hydro variables are in this RAMSES dump than are defined in config.ini (expected %d, got %d in file)"%(nvar, nvar_file), RuntimeWarning)
nvar = nvar_file
dims = dims[:nvar]
elif nvar_file>nvar :
warnings.warn("More hydro variables are in this RAMSES dump than are defined in config.ini", RuntimeWarning)
for level in xrange(maxlevel or header['nlevelmax']) :
for cpuf in xrange(1,header['ncpu']+1) :
flevel = read_fortran(f, 'i4')[0]
ncache = read_fortran(f, 'i4')[0]
assert flevel-1==level
if ncache>0 :
if cpuf==cpu :
coords, refine, gi_cpu, gi_level = grid_info_iter.next()
mark = np.where(refine==0)
assert gi_level==level
assert gi_cpu==cpu
if cpuf==cpu and len(mark[0])>0 :
for icel in xrange(2**ndim) :
i0 = i1
i1 = i0+(refine[icel]==0).sum()
for ar in dims :
ar[i0:i1] = read_fortran(f, _float_type, ncache)[(refine[icel]==0)]
skip_fortran(f, (nvar_file-nvar))
else :
skip_fortran(f, (2**ndim)*nvar_file)
for boundary in xrange(header['nboundary']) :
flevel = read_fortran(f, 'i4')[0]
ncache = read_fortran(f, 'i4')[0]
if ncache>0 :
skip_fortran(f, (2**ndim)*nvar_file)
