Page MenuHomec4science
Files F86508372

gadget.py
No OneTemporary
Actions

File Metadata

Created
Sun, Oct 6, 22:17

gadget.py
View Options

####################################################################################################################################
#
# GADGET CLASS
#
####################################################################################################################################
class Nbody_gadget(NbodyDefault):
def get_read_fcts(self):
return [self.read_particles]
def get_write_fcts(self):
return [self.write_particles]
def get_spec_vars(self):
'''
return specific variables default values for the class
'''
return {'massarr' :array([0,0,self.nbody,0,0,0]),
'atime' :0.,
'redshift' :0.,
'flag_sfr' :0,
'flag_feedback' :0,
'nall' :array([0,0,self.nbody,0,0,0]),
'flag_cooling' :0,
'num_files' :0,
'boxsize' :0.,
'omega0' :0.,
'omegalambda' :0.,
'hubbleparam' :0.,
'flag_age' :0.,
'hubbleparam' :0.,
'flag_metals' :0.,
'nallhw' :array([0,0,0,0,0,0]),
'flag_entr_ic' :0,
'critical_energy_spec' :0.,
'empty' :52*' '
}
def read_particles(self,f):
'''
read gadget file
'''
##########################################
# read the header and send it to each proc
##########################################
tpl = (24,48,Float,Float,Int,Int,24,Int,Int,Float,Float,Float,Float,Int,Int,24,Int,Float,52)
header = io.ReadBlock(f,tpl,byteorder=self.byteorder,pio=self.pio)
npart,massarr,atime,redshift,flag_sfr,flag_feedback,nall,flag_cooling,num_files,boxsize,omega0,omegalambda,hubbleparam,flag_age,flag_metals,nallhw,flag_entr_ic,critical_energy_spec,empty = header
npart = fromstring(npart,Int)
massarr = fromstring(massarr,Float)
nall = fromstring(nall,Int)
nallhw = fromstring(nallhw,Int)
if sys.byteorder != self.byteorder:
npart.byteswap()
massarr.byteswap()
nall.byteswap()
nallhw.byteswap()
##########################################
# computes variables depending on pio
##########################################
if self.pio == 'yes':
# count number of particles that have non constant masses defined further
# compute masses
nzero=0
mass = array([])
for i in range(6):
if massarr[i] == 0:
nzero = nzero+npart[i]
else:
mass = concatenate((mass,ones(npart[i])*massarr[i]))
mass = mass.astype(Float32)
npart_all = None # not used
ngas_all = None # not used
nzero_all = None # not used
nbody = sum(npart)
else:
##########################################
# compute the particle repartition
##########################################
# here, we should have : sum(npart_all) == npart
npart_all = zeros((mpi.NTask,6))
ngas_all = zeros((mpi.NTask,6))
for i in range(6):
for Task in range(mpi.NTask-1,-1,-1):
npart_all[Task,i] = npart[i]/mpi.NTask + npart[i]%mpi.NTask*(Task==0)
if i==0:
ngas_all[Task,i] = npart[i]/mpi.NTask + npart[i]%mpi.NTask*(Task==0)
# compute local number of particle
nall = npart
npart = npart_all[mpi.ThisTask]
nbody = sum(npart)
##########################################
# compute mass
##########################################
# count number of particles that have non constant masses defined further
# and create mass from massarr
nzero=0
nzero_all = zeros((mpi.NTask,len(npart)))
mass = array([],Float32)
for i in range(6):
if massarr[i] == 0:
nzero = nzero+npart[i]
nzero_all[mpi.ThisTask,i] = nzero_all[mpi.ThisTask,i] + npart[i]
else:
mass = concatenate((mass,ones(npart[i],Float32)*massarr[i]))
nzero_all = mpi.mpi_reduce(nzero_all)
##########################################
# read and send particles attribute
##########################################
pos = io.ReadArray(f,Float32,shape=(nbody,3),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
vel = io.ReadArray(f,Float32,shape=(nbody,3),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
num = io.ReadArray(f,Int32,byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
##########################################
# read mass if needed
##########################################
if nzero != 0:
massnzero = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=nzero_all)
mass = concatenate((mass,massnzero))
# extentions
u = None
rho = None
rsp = None
erd = None
dte = None
if not io.end_of_file(f,MPI=MPI):
u = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
u = concatenate((u,zeros(nbody-npart[0]).astype(Float32)))
if not io.end_of_file(f,MPI=MPI):
rho = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
rho = concatenate((rho,zeros(nbody-npart[0]).astype(Float32)))
if not io.end_of_file(f,MPI=MPI):
rsp = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
rsp = concatenate((rsp,zeros(nbody-npart[0]).astype(Float32)))
if not io.end_of_file(f,MPI=MPI):
erd = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
if (len(erd)==npart[0]): # to be compatible with old files
erd = concatenate((erd,zeros(nbody-npart[0]).astype(Float32)))
if not io.end_of_file(f,MPI=MPI):
dte = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
dte = concatenate((dte,zeros(nbody-npart[0]).astype(Float32)))
if not io.end_of_file(f,MPI=MPI): # if sn feedback is enabeled
qq = io.ReadArray(f,Float32,byteorder=self.byteorder,pio=self.pio,nlocal=ngas_all)
#if u!=None and rho==None:
# rho = 0*u
# rsp
#rsp = None
#if rho != None:
# c = (rho==0.)
# rho = where(c,1,rho)
# rsp = rho**(-1./3.)
# rsp = where(c,0,rsp)
# make global
self.npart = npart
self.massarr = massarr
self.atime = atime
self.redshift = redshift
self.flag_sfr = flag_sfr
self.flag_feedback = flag_feedback
self.nall = nall
self.flag_cooling = flag_cooling
self.num_files = num_files
self.boxsize = boxsize
self.omega0 = omega0
self.omegalambda = omegalambda
self.hubbleparam = hubbleparam
self.flag_age = flag_age
self.flag_metals = flag_metals
self.nallhw = nallhw
self.flag_entr_ic = flag_entr_ic
self.critical_energy_spec = critical_energy_spec
self.empty = empty
self.nbody = nbody
self.pos = pos
self.vel = vel
self.mass = mass
self.num = num
self.nzero = nzero
self.u = u
self.rho = rho
self.erd = erd
self.dte = dte
self.rsp = rsp
def write_particles(self,f):
'''
specific format for particle file
'''
massarr,nzero = self.get_massarr_and_nzero()
if self.pio == 'yes':
'''
here, we have to compute also
mass for each proc
'''
npart = self.npart
nall = self.npart_tot
num_files = mpi.NTask
npart_all = None
else:
npart = self.npart
nall = self.npart_tot
num_files = 1
npart_all = array(mpi.mpi_Allgather(self.npart))
# compute the global massarr and global nzero
nzero_tot = 0
massarr_all = array(mpi.mpi_Allgather(massarr))
massarr_tot = zeros(len(npart),Float)
for i in range(len(npart)):
c = (massarr_all[:,i] == massarr_all[0,i]).astype(Int)
if sum(c) == len(c): # all are equal
massarr_tot[i] = massarr_all[0,i]
else: # not equal
massarr_tot[i] = 0.0
nzero_tot = nzero_tot + sum(npart_all[:,i])
# now, re-compute nzero for the current node
massarr = massarr_tot
nzero = 0
for i in range(len(npart)):
if massarr[i] == 0:
nzero = nzero+npart[i]
# now that we have the right massarr and nzero,
# we can compute massnzero for each node
nzero_all = zeros((mpi.NTask,len(self.npart)))
if nzero != 0:
ni = 0
massnzero=array([],Float32)
for i in range(len(self.npart)):
if npart[i]!=0 and massarr[i]==0.:
massnzero = concatenate((massnzero,self.mass[ni:ni+self.npart[i]]))
nzero_all[mpi.ThisTask,i] = nzero_all[mpi.ThisTask,i] + npart[i]
ni = ni + self.npart[i]
nzero_all = mpi.mpi_reduce(nzero_all)
if self.pio == 'yes':
pass
else:
npart = self.npart_tot
nzero = mpi.mpi_reduce(nzero) # to ensure that all nodes
# will do -> write mass if needed
# header
if sys.byteorder == self.byteorder:
npart = array(npart,Int).tostring()
massarr = array(massarr,Float).tostring()
else:
npart = array(npart,Int).byteswapped().tostring()
massarr = array(massarr,Float).byteswapped().tostring()
atime = self.atime
redshift = self.redshift
flag_sfr = self.flag_sfr
flag_feedback = self.flag_feedback
if sys.byteorder == self.byteorder:
nall = array(nall,Int).tostring()
else:
nall = array(nall,Int).byteswapped().tostring()
flag_cooling = self.flag_cooling
num_files = num_files
boxsize = self.boxsize
omega0 = self.omega0
omegalambda = self.omegalambda
hubbleparam = self.hubbleparam
flag_age = self.flag_age
flag_metals = self.flag_metals
if sys.byteorder == self.byteorder:
nallhw = array(self.nallhw,Float).tostring()
else:
nallhw = array(self.nallhw,Float).byteswapped().tostring()
flag_entr_ic = self.flag_entr_ic
critical_energy_spec = self.critical_energy_spec
empty = self.empty
tpl = ((npart,24),(massarr,48),(atime,Float),(redshift,Float),(flag_sfr,Int),(flag_feedback,Int),(nall,24),(flag_cooling,Int),(num_files,Int),(boxsize,Float),(omega0,Float),(omegalambda,Float),(hubbleparam,Float),(flag_age,Int),(flag_metals,Int),(nallhw,24),(flag_entr_ic,Int),(critical_energy_spec,Float),(empty,52))
io.WriteBlock(f,tpl,byteorder=self.byteorder)
# positions
io.WriteArray(f,self.pos.astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
# velocities
io.WriteArray(f,self.vel.astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
# id
io.WriteArray(f,self.num.astype(Int),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
# write mass if needed
if nzero != 0:
io.WriteArray(f,massnzero.astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=nzero_all)
# write extension
if self.has_array('u'):
if self.u != None:
io.WriteArray(f,self.u[:self.npart[0]].astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
if self.has_array('rho'):
if self.rho != None:
io.WriteArray(f,self.rho[:self.npart[0]].astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
if self.has_array('rsp'):
if self.rsp != None:
io.WriteArray(f,self.rsp[:self.npart[0]].astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
if self.has_array('erd'):
if self.erd != None:
io.WriteArray(f,self.erd[:self.npart[0]].astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
if self.has_array('dte'):
if self.dte != None:
io.WriteArray(f,self.dte[:self.npart[0]].astype(Float32),byteorder=self.byteorder,pio=self.pio,nlocal=npart_all)
def spec_info(self):
"""
Write spec info
"""
infolist = []
infolist.append("")
#infolist.append("nzero : %s"%self.nzero)
#infolist.append("npart : %s"%self.npart)
#infolist.append("massarr : %s"%self.massarr)
infolist.append("atime : %s"%self.atime)
infolist.append("redshift : %s"%self.redshift)
infolist.append("flag_sfr : %s"%self.flag_sfr)
infolist.append("flag_feedback : %s"%self.flag_feedback)
infolist.append("nall : %s"%self.nall)
infolist.append("flag_cooling : %s"%self.flag_cooling)
infolist.append("num_files : %s"%self.num_files)
infolist.append("boxsize : %s"%self.boxsize)
infolist.append("omega0 : %s"%self.omega0)
infolist.append("omegalambda : %s"%self.omegalambda)
infolist.append("hubbleparam : %s"%self.hubbleparam)
infolist.append("flag_age : %s"%self.flag_age)
infolist.append("flag_metals : %s"%self.flag_metals)
infolist.append("nallhw : %s"%self.nallhw)
infolist.append("flag_entr_ic : %s"%self.flag_entr_ic)
infolist.append("critical_energy_spec: %s"%self.critical_energy_spec)
infolist.append("")
if self.has_array('u'):
infolist.append("len u : %s"%len(self.u))
infolist.append("u[0] : %s"%self.u[0])
infolist.append("u[-1] : %s"%self.u[-1])
if self.has_array('rho'):
infolist.append("len rho : %s"%len(self.rho))
infolist.append("rho[0] : %s"%self.rho[0])
infolist.append("rho[-1] : %s"%self.rho[-1])
if self.has_array('rsp'):
infolist.append("len rsp : %s"%len(self.rsp))
infolist.append("rsp[0] : %s"%self.rsp[0])
infolist.append("rsp[-1] : %s"%self.rsp[-1])
if self.has_array('erd'):
infolist.append("len erd : %s"%len(self.erd))
infolist.append("erd[0] : %s"%self.erd[0])
infolist.append("erd[-1] : %s"%self.erd[-1])
if self.has_array('dte'):
infolist.append("len dte : %s"%len(self.dte))
infolist.append("dte[0] : %s"%self.dte[0])
infolist.append("dte[-1] : %s"%self.dte[-1])
return infolist
def select(self,tpe='gas',val=None):
"""
Return an N-body object that contain only particles of a
certain type, defined by in gadget:
gas : gas particles
halo : halo particles
disk : disk particles
bulge : bulge particles
stars : stars particles
bndry : bndry particles
sph : gas with u > u_c
sticky : gas with u < u_c
"""
index = {'gas':0,'halo':1,'disk':2,'bulge':3,'stars':4,'bndry':5}
if (tpe=='sph'):
nb = self.select('gas')
return nb.selectc((self.u>self.critical_energy_spec))
if (tpe=='sticky'):
nb = self.select('gas')
return nb.selectc((nb.u<=nb.critical_energy_spec))
if not index.has_key(tpe):
print "unknown type %s"%(tpe)
return self
i = index[tpe]
if self.npart[i]==0:
#print "no particle of type %s"%(tpe)
return self.selectc(zeros(self.nbody))
n1 = sum(self.npart[:i])
n2 = n1 + self.npart[i]-1
return self.sub(n1,n2)
def subdis(self,mode='dd',val=None):
"""
Equivalent of select
"""
return self.select(mode)

Event Timeline

c4science · Help