diff --git a/MANIFEST b/MANIFEST index bbfc87c..e64237c 100644 --- a/MANIFEST +++ b/MANIFEST @@ -1,312 +1,308 @@ README setup.py config/defaultparameters config/unitsparameters -config/formats/binary.py -config/formats/bnbf.py config/formats/default.py config/formats/gadget.py -config/formats/simpleformat.py -config/formats/tipsy.py -config/formats/tipsybig.py +config/formats/ascii.py config/plugins/cg.py config/plugins/cmcenter.py config/plugins/cut.py config/plugins/dd.py config/plugins/dg.py config/plugins/disipatives.py config/plugins/dn.py config/plugins/gadget_bndry.py config/plugins/gadget_bulge.py config/plugins/gadget_disk.py config/plugins/gadget_gas.py config/plugins/gadget_halo.py config/plugins/gadget_stars.py config/plugins/gas.py config/plugins/gasdeg.py config/plugins/hg.py config/plugins/ncg.py config/plugins/nd.py config/plugins/newstar.py config/plugins/nhg.py config/plugins/ns.py config/plugins/reduc.py config/plugins/selectbynum.py config/plugins/st.py config/plugins/stars.py config/plugins/wg.py config/rgb_tables/aips0 config/rgb_tables/backgr config/rgb_tables/bgyrw config/rgb_tables/blackwhite config/rgb_tables/blue config/rgb_tables/blulut config/rgb_tables/color config/rgb_tables/cyan config/rgb_tables/green config/rgb_tables/greenlut config/rgb_tables/grey config/rgb_tables/half config/rgb_tables/heat config/rgb_tables/idl11 config/rgb_tables/idl12 config/rgb_tables/idl14 config/rgb_tables/idl15 config/rgb_tables/idl2 config/rgb_tables/idl4 config/rgb_tables/idl5 config/rgb_tables/idl6 config/rgb_tables/isophot config/rgb_tables/light config/rgb_tables/light_b config/rgb_tables/light_green config/rgb_tables/light_r config/rgb_tables/light_r2 config/rgb_tables/lut0 config/rgb_tables/lut1 config/rgb_tables/lut2 config/rgb_tables/lut3 config/rgb_tables/lut4 config/rgb_tables/lut5 config/rgb_tables/lut6 config/rgb_tables/lut7 config/rgb_tables/lut8 config/rgb_tables/lut9 config/rgb_tables/magenta config/rgb_tables/manycol config/rgb_tables/new config/rgb_tables/pastel config/rgb_tables/pmar config/rgb_tables/qq config/rgb_tables/rainbow config/rgb_tables/rainbow1 config/rgb_tables/rainbow2 config/rgb_tables/rainbow3 config/rgb_tables/rainbow4 config/rgb_tables/rainbow4_light_green config/rgb_tables/rainbow4lut1 config/rgb_tables/rainbow4lut2 config/rgb_tables/ramp config/rgb_tables/random config/rgb_tables/random1 config/rgb_tables/random2 config/rgb_tables/random3 config/rgb_tables/random4 config/rgb_tables/random5 config/rgb_tables/random6 config/rgb_tables/real config/rgb_tables/red config/rgb_tables/red_green_blue_yellow_cyan config/rgb_tables/redlut config/rgb_tables/rgb config/rgb_tables/rgb2 config/rgb_tables/smooth config/rgb_tables/smooth1 config/rgb_tables/smooth2 config/rgb_tables/smooth3 config/rgb_tables/staircase config/rgb_tables/stairs8 config/rgb_tables/stairs9 config/rgb_tables/standard config/rgb_tables/tululb config/rgb_tables/whiteblack config/rgb_tables/yellow config/rgb_tables/yellowlut examples/binary.dat examples/disk.dat examples/gadget.dat examples/gadget_z00.dat examples/gadget_z40.dat examples/newbinary.dat examples/num.dat examples/snap.dat examples/ic/Readme examples/ic/generate_cylindrical_model_exponnential.py examples/ic/generate_cylindrical_model_miyamoto.py examples/ic/generate_galaxy.py examples/ic/generate_spherical_model_g2c.py examples/ic/generate_spherical_model_nfw.py examples/ic/generate_spherical_model_plummer.py examples/ic/plot_cylindrical_model_exponnential.py examples/ic/plot_cylindrical_model_miyamoto.py examples/ic/plot_galaxy_frequencies.py examples/ic/plot_galaxy_stability.py examples/ic/plot_galaxy_velocities.py examples/ic/plot_spherical.py examples/ic/plot_spherical_model_plummer.py examples/scripts/example01.py examples/scripts/example02.py examples/scripts/example03.py examples/scripts/example04.py examples/scripts/example05.py examples/scripts/example06.py examples/scripts/example07.py examples/scripts/example08.py examples/scripts/example09.py examples/scripts/example10.py examples/scripts/example11.py examples/scripts/example12.py examples/scripts/example13.py examples/scripts/findmax.py examples/scripts/makesnapshot.py examples/scripts/showmap.py examples/scripts/slice-p1.py examples/scripts/slice-p2.py examples/scripts/slice.py examples/scripts/testall.py examples/scripts/mkmodel_for_display.py fonts/Courier_New.ttf fonts/Courier_New_Bold.ttf fonts/helvB08.pbm fonts/helvB08.pcf fonts/helvB08.pil fonts/helvB10.pbm fonts/helvB10.pcf fonts/helvB10.pil fonts/helvB12.pbm fonts/helvB12.pcf fonts/helvB12.pil fonts/helvB14.pbm fonts/helvB14.pcf fonts/helvB14.pil fonts/helvB18.pbm fonts/helvB18.pcf fonts/helvB18.pil fonts/helvB24.pbm fonts/helvB24.pcf fonts/helvB24.pil fonts/helvBO08.pbm fonts/helvBO08.pcf fonts/helvBO08.pil fonts/helvBO10.pbm fonts/helvBO10.pcf fonts/helvBO10.pil fonts/helvBO12.pbm fonts/helvBO12.pcf fonts/helvBO12.pil fonts/helvBO14.pbm fonts/helvBO14.pcf fonts/helvBO14.pil fonts/helvBO18.pbm fonts/helvBO18.pcf fonts/helvBO18.pil fonts/helvBO24.pbm fonts/helvBO24.pcf fonts/helvBO24.pil fonts/helvO08.pbm fonts/helvO08.pcf fonts/helvO08.pil fonts/helvO10.pbm fonts/helvO10.pcf fonts/helvO10.pil fonts/helvO12.pbm fonts/helvO12.pcf fonts/helvO12.pil fonts/helvO14.pbm fonts/helvO14.pcf fonts/helvO14.pil fonts/helvO18.pbm fonts/helvO18.pcf fonts/helvO18.pil fonts/helvO24.pbm fonts/helvO24.pcf fonts/helvO24.pil fonts/helvR08.pbm fonts/helvR08.pcf fonts/helvR08.pil fonts/helvR10.pbm fonts/helvR10.pcf fonts/helvR10.pil fonts/helvR12.pbm fonts/helvR12.pcf fonts/helvR12.pil fonts/helvR14.pbm fonts/helvR14.pcf fonts/helvR14.pil fonts/helvR18.pbm fonts/helvR18.pcf fonts/helvR18.pil fonts/helvR24.pbm fonts/helvR24.pcf fonts/helvR24.pil fonts/kidprbol.ttf pNbody/Movie.py pNbody/__init__.py pNbody/cooling.py pNbody/cosmo.py pNbody/ctes.py pNbody/fortranfile.py pNbody/fourier.py pNbody/geometry.py pNbody/ic.py pNbody/io.py pNbody/libdisk.py pNbody/libgrid.py pNbody/liblog.py pNbody/libmiyamoto.py pNbody/libqt.py pNbody/libutil.py pNbody/main.py pNbody/mpi.py pNbody/palette.py pNbody/param.py pNbody/parameters.py pNbody/phot.py pNbody/plummer.py pNbody/profiles.py pNbody/pyfits.py pNbody/rec.py pNbody/talkgdisp.py pNbody/thermodyn.py pNbody/units.py pNbody/SSP/__init__.py pNbody/SSP/libSSPluminosity.py pNbody/SSP/libbruzual.py pNbody/SSP/libmaraston.py pNbody/SSP/libtill.py pNbody/SSP/libvazdekis.py scripts/addgmov scripts/combinegmov scripts/cutgmov scripts/gavi2mp4 scripts/gcmd scripts/gimage scripts/ginter scripts/gmkgmov scripts/gmov scripts/gmov2gif scripts/gmov2gmov scripts/gmov2gmova scripts/gmov2mov scripts/gmov2mov.old scripts/gmov2mpeg scripts/gmov2mpg scripts/gmov2ppm scripts/gpy scripts/gtree scripts/gwin scripts/infogmov scripts/mergegmov scripts/mkgmov scripts/pNbody_checkall scripts/pNbody_copy-defaultconfig scripts/pNbody_copy-examples scripts/pNbody_mpi scripts/pNbody_show-parameters scripts/pNbody_show-path scripts/splitgmov scripts/supgmov src/asciilib/asciilib.c src/cooling_with_metals/cooling_with_metals.c src/coolinglib/coolinglib.c src/cosmolib/cosmolib.c src/iclib/iclib.c src/mapping/mapping.c src/montecarlolib/montecarlolib.c src/myNumeric/myNumeric.c src/nbdrklib/nbdrklib.c src/nbodymodule/nbodymodule.c src/peanolib/peanolib.c src/pygsl/pygsl.c src/tessel/tessel/tessel.c src/treelib/treelib.c diff --git a/pNbody/io.pyc b/pNbody/io.pyc deleted file mode 100644 index 43d1d7f..0000000 Binary files a/pNbody/io.pyc and /dev/null differ diff --git a/pNbody/libgrid.py b/pNbody/libgrid.py index 0098e90..0a4090d 100644 --- a/pNbody/libgrid.py +++ b/pNbody/libgrid.py @@ -1,3160 +1,3171 @@ ########################################################## # grid functions ########################################################## from numpy import * from numpy import clip as numclip import mapping import libutil import treelib import myNumeric import sys ''' definition of indexes and corresponding physical values ix = int( (x-xmin)/(xmax-xmin)* nx ) x = ix*(xmax-xmin)/(nx) + xmin center of cell xc = (ix+0.5)*(xmax-xmin)/(nx) + xmin physical size of cells dx = ((arange(nx)+1)*(xmax-xmin)/(nx) + xmin ) - ((arange(nx) )*(xmax-xmin)/(nx) + xmin ) = (xmax-xmin)/nx * ones(nx) # linear density of spherical density lambda(r) = 4*pi*r**2*rho(r) here we may deal the following grids : - carthesian 2d - carthesian 3d - cylindrical 2d - cylindrical 3d - spherical 3d !!! when computing the potential, in some 1d/2d grid, !!! we do not take into account the fact that the model !!! may be assymetric !!! We should better compute potential on 3d grid !!! and take the mean. !!! the cylindrical grid is not symetric in z ''' #################################################################################################################################### # # GENERIC 1D GRID # #################################################################################################################################### class Generic_1d_Grid: def __init__(self,rmin,rmax,nr,g=None,gm=None): ''' f(rmin) = f(rmin) f(rmax) = f(rmax) ''' self.nr = int(nr) self.rmin = float(rmin) self.rmax = float(rmax) self.g = g self.gm = gm self.f = lambda r:r # by default, identity self.fm = lambda r:r # by default, identity if self.g != None and self.gm != None: self.f = lambda r: ( self.g(r)-self.g(self.rmin) )/( self.g(self.rmax)-self.g(self.rmin) ) *(self.rmax-self.rmin) + self.rmin self.fm = lambda f: self.gm( (f-self.rmin)*( self.g(self.rmax)-self.g(self.rmin) )/( self.rmax-self.rmin ) + self.g(self.rmin) ) def get_r(self,offr=0): ir = arange(self.nr) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin return self.fm(r) def get_Points(self,offr=0): ''' Return an array of points corresponding to the nodes of a 1d spherical grid To get a nt X nr array from the returned vector (pos), do x = pos[:,0] y = pos[:,0] z = pos[:,0] x.shape = (nr,np,nt) y.shape = (nr,np,nt) z.shape = (nr,np,nt) ''' ir = indices((self.nr,)) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin r = ravel(r) r = self.fm(r) x = r y = zeros(len(x)) z = zeros(len(x)) return transpose(array([x,y,z])).astype(float32) def get_MassMap(self,r,mass): ''' Return an array of points containing mass of particles ''' r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) val = ones(pos.shape).astype(float32) mass= mass.astype(float32) shape = (self.nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_GenericMap(self,r,mass,val): ''' Return an array of points containing mass of particles ''' r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) val = val.astype(float32) mass= mass.astype(float32) shape = (self.nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_MeanMap(self,r,mass,val): ''' Return an array of points containing mean of val ''' mv = self.get_GenericMap(r,mass,val) m = self.get_MassMap(r,mass) # compute mean mv = where(mv ==0,0,mv) m = where(m==0,1,m) mean = mv/m return mean,m,mv def get_SigmaMap(self,r,mass,val): ''' Return an array of points containing sigma of val ''' mv = self.get_GenericMap(r,mass,val) mv2 = self.get_GenericMap(r,mass,val*val) m = self.get_MassMap(r,mass) # compute sigma mv = where(mv ==0,0,mv) mv2 = where(mv2==0,0,mv2) m = where(m==0,1,m) sigma = mv2/m - (mv/m)**2 sigma = sqrt(numclip(sigma,0,1e10)) return sigma,m,mv,mv2 #################################################################################################################################### # # SPHERICAL 1D GRID # #################################################################################################################################### class Spherical_1d_Grid: def __init__(self,rmin,rmax,nr,g=None,gm=None): ''' f(rmin) = f(rmin) f(rmax) = f(rmax) ''' self.nr = int(nr) self.rmin = float(rmin) self.rmax = float(rmax) self.g = g self.gm = gm self.f = lambda r:r # by default, identity self.fm = lambda r:r # by default, identity if self.g != None and self.gm != None: self.f = lambda r: ( self.g(r)-self.g(self.rmin) )/( self.g(self.rmax)-self.g(self.rmin) ) *(self.rmax-self.rmin) + self.rmin self.fm = lambda f: self.gm( (f-self.rmin)*( self.g(self.rmax)-self.g(self.rmin) )/( self.rmax-self.rmin ) + self.g(self.rmin) ) def get_r(self,offr=0): ir = arange(self.nr) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin return self.fm(r) def get_Points(self,offr=0): ''' Return an array of points corresponding to the nodes of a 1d spherical grid To get a nt X nr array from the returned vector (pos), do x = pos[:,0] y = pos[:,0] z = pos[:,0] x.shape = (nr,np,nt) y.shape = (nr,np,nt) z.shape = (nr,np,nt) ''' ir = indices((self.nr,)) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin r = ravel(r) r = self.fm(r) x = r y = zeros(len(x)) z = zeros(len(x)) return transpose(array([x,y,z])).astype(float32) def get_MassMap(self,nb): ''' Return an array of points containing mass of particles ''' r = nb.rxyz() r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (self.nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_NumberMap(self,nb): ''' Return an array of points containing number of particles ''' r = nb.rxyz() r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (self.nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_ValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing val of particles ''' r = nb.rxyz() r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_MeanValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing mean val of particles ''' r = nb.rxyz() r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,) # compute zero momentum m0 = mapping.mkmap1d(pos,mass,ones(len(pos),float32),shape) # compute first momentum m1 = mapping.mkmap1d(pos,mass,val,shape) return libutil.GetMeanMap(m0,m1) def get_SigmaValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing sigma val of particles ''' r = nb.rxyz() r = self.f(r) # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) # compute values pos = transpose(array(r)).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,) # compute zero momentum m0 = mapping.mkmap1d(pos,mass,ones(len(pos),float32),shape) # compute first momentum m1 = mapping.mkmap1d(pos,mass,val,shape) # compute second momentum m2 = mapping.mkmap1d(pos,mass,val*val,shape) return libutil.GetSigmaMap(m0,m1,m2) def get_PotentialMap(self,nb,eps,UseTree=True,Tree=None): ''' Return an array of points containing potential ''' pos = self.get_Points() if UseTree: pot = nb.TreePot(pos,eps,Tree) else: pot = nb.Pot(pos,eps) # transform back into array pot.shape = (self.nr,) return pot def get_SurfaceMap(self): ''' Return an array of points containing surface (volume) of each cell. ''' r1 = self.get_r(offr=0) r2 = self.get_r(offr=1) mat = r2-r1 mat = mat.astype(float32) return mat def get_VolumeMap(self): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' r1 = self.get_r(offr=0) r2 = self.get_r(offr=1) mat = 4.0/3.0*pi*( (r2)**3 - (r1)**3 ) mat = mat.astype(float32) return mat def get_DensityMap(self,nb): ''' Return an array of points containing density in each cell ''' m = self.get_MassMap(nb) v = self.get_VolumeMap() return m/v def get_LinearDensityMap(self,nb): ''' Return an array of points containing the linear density in each cell ''' m = self.get_MassMap(nb) s = self.get_SurfaceMap() return m/s def get_AccumulatedMassMap(self,nb): ''' Return an array of points containing M(r) in each cell ''' m = self.get_MassMap(nb) mr = get_Accumulation_Along_Axis(m,axis=0) return mr def get_Interpolation(self,pos,mat,offr=0): ''' Interpolates continuous value of pos, using matrix mat ''' r = sqrt(pos[:,0]**2+pos[:,1]**2+pos[:,2]**2) r = self.f(r) ir = (r-self.rmin)/(self.rmax-self.rmin)* self.nr - offr return myNumeric.Interpolate_From_1d_Array(ir.astype(float32),mat.astype(float32)) #################################################################################################################################### # # CYLINDRICAL 2D GRID (r-z) # #################################################################################################################################### class Cylindrical_2drz_Grid: def __init__(self,rmin,rmax,nr,zmin,zmax,nz,g=None,gm=None): ''' f(rmin) = f(rmin) f(rmax) = f(rmax) ''' self.nr = int(nr) self.rmin = float(rmin) self.rmax = float(rmax) self.nz = int(nz) self.zmin = float(zmin) self.zmax = float(zmax) self.g = g self.gm = gm self.f = lambda r:r # by default, identity self.fm = lambda r:r # by default, identity if self.g != None and self.gm != None: self.f = lambda r: ( self.g(r)-self.g(self.rmin) )/( self.g(self.rmax)-self.g(self.rmin) ) *(self.rmax-self.rmin) + self.rmin self.fm = lambda f: self.gm( (f-self.rmin)*( self.g(self.rmax)-self.g(self.rmin) )/( self.rmax-self.rmin ) + self.g(self.rmin) ) def get_rz(self,offr=0,offz=0): ir = arange(self.nr) iz = arange(self.nz) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin z = (iz+offz)*(self.zmax-self.zmin)/self.nz + self.zmin r = self.fm(r) return r,z def get_Points(self,offr=0,offz=0): ''' Return an array of points corresponding to the nodes of a 2d cylindrical grid To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nr,nt) y.shape = (nr,nt) z.shape = (nr,nt) # to get r and theta r = sqrt(x**2+y**2+z**2) t = arctan2(y,x)*180/pi ''' ir,iz = indices((self.nr,self.nz)) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin z = (iz+offz)*(self.zmax-self.zmin)/self.nz + self.zmin r = ravel(r) z = ravel(z) r = self.fm(r) x = r y = zeros(len(x)) z = z return transpose(array([x,y,z])).astype(float32) def get_MassMap(self,nb,offr=0,offz=0): ''' Return an array of points containing mass of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr z = (z-self.zmin)/(self.zmax-self.zmin) - offz/self.nz # compute values pos = transpose(array([r,z,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (self.nr,self.nz) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_NumberMap(self,nb,offr=0,offz=0): ''' Return an array of points containing mass of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr z = (z-self.zmin)/(self.zmax-self.zmin) - offz/self.nz # compute values pos = transpose(array([r,z,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (self.nr,self.nz) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_PotentialMap(self,nb,eps,UseTree=True,Tree=None,AdaptativeSoftenning=False): ''' Return an array of points containing potential ''' if AdaptativeSoftenning: pos = self.get_Points() pot = zeros(len(pos),float) # define eps R1,z1 = self.get_rz(offr=0) R2,z2 = self.get_rz(offr=1) epss = R2-R1 epss = epss/epss[0]*eps pos = self.get_Points() for ir in range(self.nr): for iz in range(self.nz): #print ir*self.nz + iz, pos[ir*self.nz + iz], epss[ir] if UseTree: pot[ir*self.nz + iz] = nb.TreePot(array([pos[ir*self.nz + iz]]),epss[ir]) else: pot[ir*self.nz + iz] = nb.Pot( array([pos[ir*self.nz + iz]]),epss[ir]) else: pos = self.get_Points() if UseTree: pot = nb.TreePot(pos,eps) else: pot = nb.Pot(pos,eps) # transform back into array pot.shape = (self.nr,self.nz) return pot def get_AccelerationMap(self,nb,eps,UseTree=True,Tree=None): ''' Return an array of points containing accelerations ''' if AdaptativeSoftenning: pos = self.get_Points() acc = pos*0 # define eps R1,z1 = self.get_rz(offr=0) R2,z2 = self.get_rz(offr=1) epss = R2-R1 epss = epss/epss[0]*eps pos = self.get_Points() for ir in range(self.nr): for iz in range(self.nz): #print ir*self.nz + iz, pos[ir*self.nz + iz], epss[ir] if UseTree: acc[ir*self.nz + iz] = nb.TreeAccel(array([pos[ir*self.nz + iz]]),epss[ir]) else: acc[ir*self.nz + iz] = nb.Accel( array([pos[ir*self.nz + iz]]),epss[ir]) else: pos = self.get_Points() if UseTree: acc = nb.TreeAccel(pos,eps) else: acc = nb.Accel(pos,eps) accx = copy(acc[:,0]) accy = copy(acc[:,1]) accz = copy(acc[:,2]) # transform back into array accx.shape = (self.nr,self.nz) accy.shape = (self.nr,self.nz) accz.shape = (self.nr,self.nz) return accx,accy,accz def get_VolumeMap(self,offr=0,offz=0): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rs,zs = self.get_rz(offr,offz) def volume(ir,iz): if ir==self.nr-1: r1 = rs[ir-1] # un peu bricolage... r2 = rs[ir] else: r1 = rs[ir] r2 = rs[ir+1] if iz==self.nz-1: z1 = zs[iz-1] # un peu bricolage... z2 = zs[iz] else: z1 = zs[iz] z2 = zs[iz+1] return pi*(r2**2 - r1**2 ) * (z2-z1) # make the map mat = zeros((self.nr,self.nz)) for ir in range(self.nr): for iz in range(self.nz): mat[ir,iz]=volume(ir,iz) mat = mat.astype(float32) return mat def get_SurfaceMap(self,offr=0,offz=0): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rs,zs = self.get_rz(offr,offz) def surface(ir): if ir==self.nr-1: r1 = rs[ir-1] # un peu bricolage... r2 = rs[ir] else: r1 = rs[ir] r2 = rs[ir+1] return pi*(r2**2 - r1**2 ) # make the map mat = zeros(self.nr) for ir in range(self.nr): mat[ir]=surface(ir) mat = mat.astype(float32) return mat def get_DensityMap(self,nb,offr=0,offz=0): ''' Return an array of points containing density in each cell ''' m = self.get_MassMap(nb,offr=offr,offz=offz) v = self.get_VolumeMap(offr=offr,offz=offz) return m/v def get_SurfaceDensityMap(self,nb,offr=0,offz=0): ''' Return an array (1d) of points containing the surface density along r ''' m = self.get_MassMap(nb,offr=offr,offz=offz) m = get_Sum_Along_Axis(m) s = self.get_SurfaceMap(offr=offr,offz=offz) return m/s def get_Interpolation(self,pos,mat,offr=0,offz=0): ''' Interpolates continuous value of pos, using matrix mat ''' r = sqrt(pos[:,0]**2+pos[:,1]**2) r = self.f(r) z = pos[:,2] ir = (r-self.rmin)/(self.rmax-self.rmin)* self.nr - offr iz = (z-self.zmin)/(self.zmax-self.zmin)* self.nz - offz return myNumeric.Interpolate_From_2d_Array(ir.astype(float32),iz.astype(float32),mat.astype(float32)) def get_r_Interpolation(self,pos,mat,offr=0): ''' Interpolates continuous value of pos, using matrix mat only along first axis. ''' r = sqrt(pos[:,0]**2+pos[:,1]**2) r = self.f(r) ir = (r-self.rmin)/(self.rmax-self.rmin)* self.nr - offr return myNumeric.Interpolate_From_1d_Array(ir.astype(float32),mat.astype(float32)) #################################################################################################################################### # # CYLINDRICAL 2D GRID (r-t) # #################################################################################################################################### class Cylindrical_2drt_Grid: def __init__(self,rmin,rmax,nr,nt,z=None,g=None,gm=None): ''' f(rmin) = f(rmin) f(rmax) = f(rmax) ''' tmin = 0 tmax = 2*pi self.z = z self.nr = int(nr) self.rmin = float(rmin) self.rmax = float(rmax) self.nt = int(nt) self.tmin = float(tmin) self.tmax = float(tmax) self.g = g self.gm = gm self.f = lambda r:r # by default, identity self.fm = lambda r:r # by default, identity if self.g != None and self.gm != None: self.f = lambda r: ( self.g(r)-self.g(self.rmin) )/( self.g(self.rmax)-self.g(self.rmin) ) *(self.rmax-self.rmin) + self.rmin self.fm = lambda f: self.gm( (f-self.rmin)*( self.g(self.rmax)-self.g(self.rmin) )/( self.rmax-self.rmin ) + self.g(self.rmin) ) def get_rt(self,offr=0,offt=0): ir = arange(self.nr) it = arange(self.nt) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin t = (it+offt)*(self.tmax-self.tmin)/self.nt + self.tmin r = self.fm(r) return r,t def get_Points(self,offr=0,offt=0): ''' Return an array of points corresponding to the nodes of a 2d cylindrical grid To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nr,nt) y.shape = (nr,nt) z.shape = (nr,nt) # to get r and theta r = sqrt(x**2+y**2+z**2) t = arctan2(y,x)*180/pi ''' ir,it = indices((self.nr,self.nt)) r = (ir+offr)*(self.rmax-self.rmin)/self.nr + self.rmin t = (it+offt)*(self.tmax-self.tmin)/self.nt + self.tmin r = ravel(r) t = ravel(t) r = self.fm(r) x = r*cos(t) y = r*sin(t) z = ones(len(x))*self.z return transpose(array([x,y,z])).astype(float32) def get_MassMap(self,nb,offr=0,offt=0): ''' Return an array of points containing mass of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr t = (t-self.tmin)/(self.tmax-self.tmin) - offt/self.nt # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (self.nr,self.nt) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_NumberMap(self,nb,offr=0,offt=0): ''' Return an array of points containing mass of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr t = (t-self.tmin)/(self.tmax-self.tmin) - offt/self.nt # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (self.nr,self.nt) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_ValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing val of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr t = (t-self.tmin)/(self.tmax-self.tmin) - offt/self.nt # compute values pos = transpose(array([r,t,z])).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,self.nt) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_MeanValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing mean val of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr t = (t-self.tmin)/(self.tmax-self.tmin) - offt/self.nt # compute values pos = transpose(array([r,t,z])).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,self.nt) # compute zero momentum m0 = mapping.mkmap2d(pos,mass,ones(len(pos),float32),shape) # compute first momentum m1 = mapping.mkmap2d(pos,mass,val,shape) return libutil.GetMeanMap(m0,m1) def get_SigmaValMap(self,nb,val,offr=0,offt=0): ''' Return an array of points containing sigma val of particles ''' r = nb.rxy() r = self.f(r) t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-self.rmin)/(self.rmax-self.rmin) - offr/self.nr t = (t-self.tmin)/(self.tmax-self.tmin) - offt/self.nt # compute values pos = transpose(array([r,t,z])).astype(float32) mass= nb.mass.astype(float32) val = val.astype(float32) shape = (self.nr,self.nt) # compute zero momentum m0 = mapping.mkmap2d(pos,mass,ones(len(pos),float32),shape) # compute first momentum m1 = mapping.mkmap2d(pos,mass,val,shape) # compute second momentum m2 = mapping.mkmap2d(pos,mass,val*val,shape) return libutil.GetSigmaMap(m0,m1,m2) def get_PotentialMap(self,nb,eps,UseTree=True,Tree=None,AdaptativeSoftenning=False): ''' Return an array of points containing potential ''' if AdaptativeSoftenning: pos = self.get_Points() pot = zeros(len(pos),float) # define eps R1,t1 = self.get_rt(offr=0) R2,t2 = self.get_rt(offr=1) epss = R2-R1 epss = epss/epss[0]*eps pos = self.get_Points() for ir in range(self.nr): for it in range(self.nt): #print ir*self.nt + it, pos[ir*self.nt + it], epss[ir] if UseTree: pot[ir*self.nt + it] = nb.TreePot(array([pos[ir*self.nt + it]]),epss[ir]) else: pot[ir*self.nt + it] = nb.Pot( array([pos[ir*self.nt + it]]),epss[ir]) else: pos = self.get_Points() if UseTree: pot = nb.TreePot(pos,eps) else: pot = nb.Pot(pos,eps) # transform back into array pot.shape = (self.nr,self.nt) return pot def get_AccelerationMap(self,nb,eps,UseTree=True,Tree=None,AdaptativeSoftenning=False): ''' Return an array of points containing accelerations ''' if AdaptativeSoftenning: pos = self.get_Points() acc = pos*0 # define eps R1,t1 = self.get_rt(offr=0) R2,t2 = self.get_rt(offr=1) epss = R2-R1 epss = epss/epss[0]*eps pos = self.get_Points() for ir in range(self.nr): for it in range(self.nt): #print ir*self.nt + it, pos[ir*self.nt + it], epss[ir] if UseTree: acc[ir*self.nt + it] = nb.TreeAccel(array([pos[ir*self.nt + it]]),epss[ir]) else: acc[ir*self.nt + it] = nb.Accel( array([pos[ir*self.nt + it]]),epss[ir]) else: pos = self.get_Points() if UseTree: acc = nb.TreeAccel(pos,eps) else: acc = nb.Accel(pos,eps) accx = copy(acc[:,0]) accy = copy(acc[:,1]) accz = copy(acc[:,2]) # transform back into array accx.shape = (self.nr,self.nt) accy.shape = (self.nr,self.nt) accz.shape = (self.nr,self.nt) return accx,accy,accz def get_SurfaceMap(self,nb,offr=0,offt=0): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rs,ts = self.get_rt(offr,offt) def surface(ir): if ir==self.nr-1: r1 = rs[ir-1] # un peu bricolage... r2 = rs[ir] else: r1 = rs[ir] r2 = rs[ir+1] return pi*(r2**2 - r1**2 ) / self.nt # make the map mat = zeros(self.nr) for ir in range(self.nr): mat[ir]=surface(ir) mat = mat.astype(float32) mat = transpose( ones((self.nt,self.nr))*mat ) return mat def get_SurfaceDensityMap(self,nb,offr=0,offt=0): ''' Return an array of points containing density in each cell ''' m = self.get_MassMap(nb,offr=0,offt=0) s = self.get_SurfaceMap(nb,offr=0,offt=0) return m/s + def get_ReducedSurfaceDensityMap(self,nb,offr=0,offt=0): + ''' + Return an array of points containing density in each cell + ''' + + m = self.get_MassMap(nb,offr=0,offt=0) + m = sum(m,axis=1) + s = self.get_SurfaceMap(nb,offr=0,offt=0) + s = sum(s,axis=1) + + return m/s #################################################################################################################################### # # GRID CLASS (not used and not finished) # #################################################################################################################################### class Grid: def __init__(self,dim=1,params={},f=None,fm=None): """ Main grid object parameters xmin,xmax,nx ymin,ymax,ny zmin,zmax,nz """ # example : 1d grid self.xmin = float(params['xmin']) self.xmax = float(params['xmax']) self.nx = int(params['nx']) self.offx = 0. self.f = f self.fm= fm ix = indices((self.nx,)) # apply indx-phys transformation if self.f!=None: x = self.f(ix,self.nx,self.xmin,self.xmax) else: x = (ix+self.offx)*(self.xmax-self.xmin)/(self.nx-1) + self.xmin x = ravel(x) x = x y = zeros(len(x)) z = zeros(len(x)) self.pos = transpose(array([x,y,z])).astype(float32) self.shape = (self.nx,) def GetPoints(self): return self.pos def GetNumberMap(self,x): x = x.astype(float32) val = ones(x.shape).astype(float32) # apply phys-indx transformation and tranform between 0 and 1 if self.fm!=None: x = self.fm(x,self.nx,self.xmin,self.xmax) / (self.nx-1) else: x = (x-self.xmin)/(self.xmax-self.xmin) mat = mapping.mkmap1dw(x,val,val,self.shape) return mat def GetMassMap(self,x,mass): x = x.astype(float32) mass = mass.astype(float32) val = ones(x.shape).astype(float32) # apply phys-indx transformation and tranform between 0 and 1 if self.fm!=None: x = self.fm(x,self.nx,self.xmin,self.xmax) / (self.nx-1) else: x = (x-self.xmin)/(self.xmax-self.xmin) mat = mapping.mkmap1dw(x,mass,val,self.shape) return mat def GetSurfaceMap(self): # !!!!!!!!!!!! this is for carthesian grid !!!!!!!!!!!!!!! ix = arange(self.nx-1) + 0.5 ix = concatenate(([0],ix,[self.nx-1])) # apply indx-phys transformation if self.f!=None: ix = self.f(ix,self.nx,self.xmin,self.xmax) else: ix = (ix+self.offx)*(self.xmax-self.xmin)/(self.nx-1) + self.xmin # compute diff. s = ix[1:] - ix[:-1] return s def GetVolumeMap(self): # !!!!!!!!!!!! this is for spherical grid !!!!!!!!!!!!!!! ix = arange(self.nx-1) + 0.5 ix = concatenate(([0],ix,[self.nx-1])) # apply indx-phys transformation if self.f!=None: ix = self.f(ix,self.nx,self.xmin,self.xmax) else: ix = (ix+self.offx)*(self.xmax-self.xmin)/(self.nx-1) + self.xmin # compute volume. v = 4/3.*pi * (ix[1:]**3 - ix[:-1]**3) return v def GetSurfaceDensityMap(self,x,mass): return self.GetMassMap(x,mass)/self.GetSurfaceMap() def GetDensityMap(self,x,mass): return self.GetMassMap(x,mass)/self.GetVolumeMap() def GetPotential(self,nb,eps,force_computation=False,ErrTolTheta=0.8): Tree = nb.getTree(force_computation=force_computation,ErrTolTheta=ErrTolTheta) pot = Tree.Potential(self.pos,eps) # transform back into array #pot.shape = (nx,ny) return pot def write(self,name='grid.dat',ftype='gadget'): ''' it cant work, because one need pNbody and pNbody uses libgrid !!! ''' # create an Nbody object #nb = Nbody(status='new',p_name=name,pos=self.pos,ftype=ftype) # and save it #nb.write() pass ####################################### # general functions ####################################### def get_First_Derivative(f,x,s=None,k=2): ''' First derivative of f(x) ''' #if s!=None: # tck = interpolate.fitpack.splrep(x,f,s=s,k=k) # f = interpolate.fitpack.splev(x,tck) fp = zeros(len(x),x.dtype) fp[0 ] = (f[ 1]-f[ 0])/(x[ 1]-x[ 0]) fp[-1] = (f[-1]-f[-2])/(x[-1]-x[-2]) f1 = f[2:] f2 = f[:-2] x1 = x[2:] x2 = x[:-2] fp[1:-1] = (f1-f2)/(x1-x2) return fp def get_Mean_Along_Axis(mat,axis=0): if len(mat.shape) == 1: return mat elif len(mat.shape) == 2: a = fmod((axis+1),2) n = mat.shape[a] s = mat s = sum(s,axis=a)/n elif len(mat.shape) == 3: a = array([fmod((axis+1),3),fmod((axis+2),3)]) a.sort() a1 = a[0] a2 = a[1] n1 = mat.shape[a1] n2 = mat.shape[a2] s = mat s = sum(s,axis=a2)/n2 s = sum(s,axis=a1)/n1 return s def get_Sum_Along_Axis(mat,axis=0): if len(mat.shape) == 1: return mat elif len(mat.shape) == 2: a = fmod((axis+1),2) s = mat s = sum(s,axis=a) elif len(mat.shape) == 3: a = array([fmod((axis+1),3),fmod((axis+2),3)]) a.sort() a1 = a[0] a2 = a[1] s = mat s = sum(s,axis=a2) s = sum(s,axis=a1) return s def get_Accumulation_Along_Axis(mat,axis=0): ''' Accumulate values along an axis ''' v = get_Sum_Along_Axis(mat,axis=axis) return add.accumulate(v) def get_Integral(v,dr,ia,ib): ''' Integrate the vector v, between ia and ib. v : values of cells (must be 1 dimensional) dr : corresponding physical size of cells ia : lower real indice ib : higher real indice ''' print "WARNING : libgrid.get_Integral : you should not use this function !!!" ia = max(0,ia) ib = min(len(v),ib) if ia==ib: return 0.0 if ia>ib: raise "ia must be < ib" iap = int(ceil(ia)) ibp = int(floor(ib)) dra = iap-ia drb = ib-ibp Ia = 0.0 if dra != 0: Ia = v[iap-1] * dra Ib = 0.0 if drb != 0: Ib = v[ibp] * drb I = v[iap:ibp]*dr[iap:ibp] return sum(I)+Ia+Ib def get_Symetrisation_Along_Axis(mat,axis=1): ''' Return an array where the two half are symetrized ''' nx,ny = mat.shape odd = fmod(ny,2) if odd: mat1 = mat[:,1:(ny/2)+1] mat2 = mat[:,(ny/2)+1:] else: mat1 = mat[:,1:ny/2] mat2 = mat[:,ny/2+1:] #mat2 = myNumeric.turnup(mat2,0) mat2 = fliplr(mat2) # take the mean matm = (mat1+mat2) /2 #mat = ones(mat.shape,mat.dtype) if odd: mat[:,1:(ny/2)+1] = matm mat[:,(ny/2)+1:] = fliplr(matm) # myNumeric.turnup(matm,0) else: mat[:,1:ny/2] = matm mat[:,ny/2+1:] = fliplr(matm) # myNumeric.turnup(matm,0) return mat def get_Symetrisation_Along_Axis_Old(mat,axis=1): ''' Return an array where the two half are symetrized Old but more correct than new one ''' nx,ny = mat.shape odd = fmod(ny,2) if odd: mat1 = mat[:,0:(ny-1)/2] mat2 = mat[:,(ny-1)/2+1:] mat3 = mat[:,(ny-1)/2] else: mat1 = mat[:,0:ny/2] mat2 = mat[:,ny/2:] #mat2 = myNumeric.turnup(mat2,0) mat2 = fliplr(mat2) # take the mean matm = (mat1+mat2) /2 mat = ones(mat.shape,mat.dtype) if odd: mat[:,0:(ny-1)/2] = matm mat[:,(ny-1)/2+1:] = fliplr(matm) # myNumeric.turnup(matm,0) mat[:,(ny-1)/2] = mat3 else: mat[:,0:ny/2] = matm mat[:,ny/2:] = fliplr(matm) # myNumeric.turnup(matm,0) return mat ####################################### # carthesian grid ####################################### ############## # 2 dimensions ############## def get_xy_Of_Carthesian_2d_Grid(nx,ny,xmin,xmax,ymin,ymax,offx=0,offy=0): xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) ix = arange(nx) iy = arange(ny) x = (ix+offx)*(xmax-xmin)/nx + xmin y = (iy+offy)*(ymax-ymin)/ny + ymin return x,y def get_Points_On_Carthesian_2d_Grid(nx,ny,xmin,xmax,ymin,ymax,offx=0,offy=0): ''' Return an array of points corresponding to the center of cells af a 2d carthesian grid. To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nx,ny) y.shape = (nx,ny) z.shape = (nx,ny) ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) nx = int(nx) ny = int(ny) ix,iy = indices((nx,ny)) x = (ix+offx)*(xmax-xmin)/nx + xmin y = (iy+offy)*(ymax-ymin)/ny + ymin x = ravel(x) y = ravel(y) z = zeros(len(x)) return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax): ''' Return an array of points containing mass of particles ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) x = nb.pos[:,0] y = nb.pos[:,1] z = nb.pos[:,2] # scale between 0 and 1 x = (x-xmin)/(xmax-xmin) y = (y-ymin)/(ymax-ymin) # compute values pos = transpose(array([x,y,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nx,ny) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_NumberMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax): ''' Return an array of points containing mass of particles ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) x = nb.pos[:,0] y = nb.pos[:,1] z = nb.pos[:,2] # scale between 0 and 1 x = (x-xmin)/(xmax-xmin) y = (y-ymin)/(ymax-ymin) # compute values pos = transpose(array([x,y,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nx,ny) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_PotentialMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Carthesian_2d_Grid(nx,ny,xmin,xmax,ymin,ymax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nx,ny) return pot def get_SurfaceMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) # compute volume (surface) dx = (xmax-xmin)/nx dy = (ymax-ymin)/ny v = dx*dy # make the map mat = ones((nx,ny)) * v mat = mat.astype(float32) return mat def get_SurfaceDensityMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax) s = get_SurfaceMap_On_Carthesian_2d_Grid(nb,nx,ny,xmin,xmax,ymin,ymax) return m/s ############## # 3 dimensions ############## def get_xyz_Of_Carthesian_3d_Grid(nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax,offx=0,offy=0,offz=0): xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) zmax = float(zmax) zmin = float(zmin) ix = arange(nx) iy = arange(ny) iz = arange(nz) x = (ix+offx)*(xmax-xmin)/nx + xmin y = (iy+offy)*(ymax-ymin)/ny + ymin z = (iz+offz)*(zmax-zmin)/nz + zmin return x,y,z def get_Points_On_Carthesian_3d_Grid(nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax,offx=0,offy=0,offz=0): ''' Return an array of points corresponding to the center of cells af a 3d carthesian grid. To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nx,ny,nz) y.shape = (nx,ny,nz) z.shape = (nx,ny,nz) ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) zmax = float(zmax) zmin = float(zmin) nx = int(nx) ny = int(ny) nz = int(nz) ix,iy,iz = indices((nx,ny,nz)) x = (ix+offx)*(xmax-xmin)/nx + xmin y = (iy+offy)*(ymax-ymin)/ny + ymin z = (iz+offz)*(zmax-zmin)/nz + zmin x = ravel(x) y = ravel(y) z = ravel(z) return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax): ''' Return an array of points containing mass of particles ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) zmax = float(zmax) zmin = float(zmin) x = nb.pos[:,0] y = nb.pos[:,1] z = nb.pos[:,2] # scale between 0 and 1 x = (x-xmin)/(xmax-xmin) y = (y-ymin)/(ymax-ymin) z = (z-zmin)/(zmax-zmin) # compute values pos = transpose(array([x,y,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nx,ny,nz) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_NumberMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax): ''' Return an array of points containing mass of particles ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) zmax = float(zmax) zmin = float(zmin) x = nb.pos[:,0] y = nb.pos[:,1] z = nb.pos[:,2] # scale between 0 and 1 x = (x-xmin)/(xmax-xmin) y = (y-ymin)/(ymax-ymin) z = (z-zmin)/(zmax-zmin) # compute values pos = transpose(array([x,y,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nx,ny,nz) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_PotentialMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nx,ny,nz) return pot def get_VolumeMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' xmax = float(xmax) xmin = float(xmin) ymax = float(ymax) ymin = float(ymin) zmax = float(zmax) zmin = float(zmin) # compute volume (surface) dx = (xmax-xmin)/nx dy = (ymax-ymin)/ny dz = (zmax-zmin)/nz v = dx*dy*dz # make the map mat = ones((nx,ny,nz)) * v mat = mat.astype(float32) return mat def get_DensityMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax) v = get_VolumeMap_On_Carthesian_3d_Grid(nb,nx,ny,nz,xmin,xmax,ymin,ymax,zmin,zmax) return m/v ####################################### # cylindrical grid ####################################### ###################### # 2 dimensions r,t (h) ###################### def get_rt_Of_Cylindrical_2dh_Grid(nr,nt,rmax,offr=0,offt=0): rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi ir = arange(nr) it = arange(nt) r = (ir+offr)*(rmax-rmin)/nr + rmin t = (it+offt)*(tmax-tmin)/nt + tmin return r,t def get_Points_On_Cylindrical_2dh_Grid(nr,nt,rmax,offr=0,offt=0): ''' Return an array of points corresponding to the nodes of a 2d cylindrical grid To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nr,nt) y.shape = (nr,nt) z.shape = (nr,nt) # to get r and theta r = sqrt(x**2+y**2+z**2) t = arctan2(y,x)*180/pi ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi nr = int(nr) nt = int(nt) ir,it = indices((nr,nt)) r = (ir+offr)*(rmax-rmin)/nr + rmin t = (it+offt)*(tmax-tmin)/nt + tmin r = ravel(r) t = ravel(t) x = r*cos(t) y = r*sin(t) z = zeros(len(x)) return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) t = (t-tmin)/(tmax-tmin) # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nr,nt) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_NumberMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) t = (t-tmin)/(tmax-tmin) # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nr,nt) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_PotentialMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Cylindrical_2dh_Grid(nr,nt,rmax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nr,nt) return pot def get_SurfaceMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rmax = float(rmax) # compute volume (surface) #dr = rmax/nr #rs = arange(nr+1)*dr #v = (pi*rs[1:]**2 - pi*rs[:-1]**2 )/nt def volume(ir,it): dr = (rmax/nr) r = ir*dr return pi*((r+dr)**2 - (r)**2 )/nt # make the map mat = fromfunction(volume,(nr,nt)) mat = mat.astype(float32) return mat def get_SurfaceDensityMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax) s = get_SurfaceMap_On_Cylindrical_2dh_Grid(nb,nr,nt,rmax) return m/s ###################### # 2 dimensions r,z (v) ###################### def get_rz_Of_Cylindrical_2dv_Grid(nr,nz,rmax,zmin,zmax): rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) ir = arange(nr) iz = arange(nz) r = (ir+0.0)*(rmax-rmin)/nr + rmin z = (iz+0.0)*(zmax-zmin)/nz + zmin return r,z def get_Points_On_Cylindrical_2dv_Grid(nr,nz,rmax,zmin,zmax,offr=0,offz=0): ''' Return an array of points corresponding to the nodes of a 2d cylindrical grid To get a nt X nr array from the returned vector (pos), do x = copy(pos[:,0]) y = copy(pos[:,1]) z = copy(pos[:,2]) x.shape = (nr,nt) y.shape = (nr,nt) z.shape = (nr,nt) # to get r and theta r = sqrt(x**2+y**2+z**2) t = arctan2(y,x)*180/pi ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) ir,iz = indices((nr,nz)) r = (ir+offr)*(rmax-rmin)/nr + rmin z = (iz+offz)*(zmax-zmin)/nz + zmin r = ravel(r) z = ravel(z) x = r y = zeros(len(x)) z = z return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) z = (z-zmin)/(zmax-zmin) # compute values pos = transpose(array([r,z,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nr,nz) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_NumberMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) z = (z-zmin)/(zmax-zmin) # compute values pos = transpose(array([r,z,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nr,nz) # make the map mat = mapping.mkmap2d(pos,mass,val,shape) return mat def get_PotentialMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Cylindrical_2dv_Grid(nr,nz,rmax,zmin,zmax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nr,nz) return pot def get_AccelerationMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax,eps,Tree=None): ''' Return an array of points containing accelerations ''' pos = get_Points_On_Cylindrical_2dv_Grid(nr,nz,rmax,zmin,zmax) Tree = nb.getTree() acc = Tree.Acceleration(pos,eps) accx = copy(acc[:,0]) accy = copy(acc[:,1]) accz = copy(acc[:,2]) # transform back into array accx.shape = (nr,nz) accy.shape = (nr,nz) accz.shape = (nr,nz) return accx,accy,accz def get_VolumeMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) # compute volume (surface) #dr = rmax/nr #rs = arange(nr+1)*dr #v = (pi*rs[1:]**2 - pi*rs[:-1]**2 )/nt def volume(ir,iz): dr = (rmax-rmin)/nr dz = (zmax-zmin)/nz r = ir*dr return pi*((r+dr)**2 - (r)**2 ) * dz # make the map mat = fromfunction(volume,(nr,nz)) mat = mat.astype(float32) return mat def get_DensityMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax) v = get_VolumeMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax) return m/v def get_SurfaceDensityMap_From_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax): ''' Return an array of points containing the surface density along r ''' m = get_MassMap_On_Cylindrical_2dv_Grid(nb,nr,nz,rmax,zmin,zmax) m = get_Sum_Along_Axis(m) def surface(ir): dr = (rmax)/nr r = ir*dr return pi*((r+dr)**2 - (r)**2 ) s = fromfunction(surface,(nr,)) return m/s def get_Interpolation_On_Cylindrical_2dv_Grid(pos,mat,nr,nz,rmax,zmin,zmax,offr=0,offz=0): ''' Interpolates continuous value of pos, using matrix mat ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) r = sqrt(pos[:,0]**2+pos[:,1]**2) z = pos[:,2] ir = (r-rmin)/(rmax-rmin)* nr - offr iz = (z-zmin)/(zmax-zmin)* nz - offz return myNumeric.Interpolate_From_2d_Array(ir.astype(float32),iz.astype(float32),mat.astype(float32)) def get_r_Interpolation_On_Cylindrical_2dv_Grid(pos,mat,nr,nz,rmax,zmin,zmax,offr=0): ''' Interpolates continuous value of pos, using matrix mat only along first axis. ''' rmin = 0.0 rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) nr = int(nr) nz = int(nz) r = sqrt(pos[:,0]**2+pos[:,1]**2) ir = (r-rmin)/(rmax-rmin)* nr - offr return myNumeric.Interpolate_From_1d_Array(ir.astype(float32),mat.astype(float32)) ############## # 3 dimensions ############## def get_rtz_Of_Cylindrical_3d_Grid(nr,nt,nz,rmax,zmin,zmax,offr=0,offt=0,offz=0): rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi zmin = float(zmin) zmax = float(zmax) ir = arange(nr) it = arange(nt) iz = arange(nz) r = (ir+offr)*(rmax-rmin)/nr + rmin t = (it+offt)*(tmax-tmin)/nt + tmin z = (iz+offz)*(zmax-zmin)/nz + zmin return r,t,z def get_Points_On_Cylindrical_3d_Grid(nr,nt,nz,rmax,zmin,zmax,offr=0,offt=0,offz=0): ''' Return an array of points corresponding to the nodes of a 2d cylindrical grid To get a nt X nr array from the returned vector (pos), do x = pos[:,0] y = pos[:,0] z = pos[:,0] x.shape = (nr,nt,nz) y.shape = (nr,nt,nz) z.shape = (nr,nt,nz) # to get r and theta r = sqrt(x**2+y**2+z**2) t = arctan2(y,x)*180/pi ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi zmin = float(zmin) zmax = float(zmax) nr = int(nr) nt = int(nt) nz = int(nz) ir,it,iz = indices((nr,nt,nz)) r = (ir+offr)*(rmax-rmin)/nr + rmin t = (it+offt)*(tmax-tmin)/nt + tmin z = (iz+offz)*(zmax-zmin)/nz + zmin r = ravel(r) t = ravel(t) z = ravel(z) x = r*cos(t) y = r*sin(t) z = z return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi zmin = float(zmin) zmax = float(zmax) r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) t = (t-tmin)/(tmax-tmin) z = (t-zmin)/(zmax-zmin) # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nr,nt,nz) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_NumberMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) tmin = 0.0 tmax = 2*pi zmin = float(zmin) zmax = float(zmax) r = nb.rxy() t = nb.phi_xy() + pi z = nb.pos[:,2] # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) t = (t-tmin)/(tmax-tmin) z = (t-zmin)/(zmax-zmin) # compute values pos = transpose(array([r,t,z])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nr,nt,nz) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_PotentialMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Cylindrical_3d_Grid(nr,nt,nz,rmax,zmin,zmax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nr,nt,nz) return pot def get_VolumeMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rmax = float(rmax) zmin = float(zmin) zmax = float(zmax) # compute volume (surface) #dr = rmax/nr #rs = arange(nr+1)*dr #v = (pi*rs[1:]**2 - pi*rs[:-1]**2 )/nt def volume(ir,it,iz): dr = (rmax/nr) dz = (zmax-zmin)/nz r = ir*dr return pi*((r+dr)**2 - (r)**2 )/nt * dz # make the map mat = fromfunction(volume,(nr,nt,nz)) mat = mat.astype(float32) return mat def get_DensityMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax) v = get_VolumeMap_On_Cylindrical_3d_Grid(nb,nr,nt,nz,rmax,zmin,zmax) return m/v ####################################### # spherical grid ####################################### ############## # 1 dimension ############## def get_r_Of_Spherical_1d_Grid(nr,rmax,offr=0,f=None,fm=None): rmin = 0.0 rmax = float(rmax) if f!=None: rmin=f(rmin) rmax=f(rmax) ir = arange(nr) r = (ir+offr)*(rmax-rmin)/nr + rmin if fm!=None: r = fm(r) return r def get_Points_On_Spherical_1d_Grid(nr,rmax,offr=0,f=None,fm=None): ''' Return an array of points corresponding to the nodes of a 1d spherical grid To get a nt X nr array from the returned vector (pos), do x = pos[:,0] y = pos[:,0] z = pos[:,0] x.shape = (nr,np,nt) y.shape = (nr,np,nt) z.shape = (nr,np,nt) ''' rmin = 0.0 rmax = float(rmax) nr = int(nr) if f!=None: rmin=f(rmin) rmax=f(rmax) ir = indices((nr,)) r = (ir+offr)*(rmax-rmin)/nr + rmin r = ravel(r) if fm!=None: r = fm(r) x = r y = zeros(len(x)) z = zeros(len(x)) return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) r = nb.rxyz() # 0 -> rmax # here, we scale if f!=None: rmin = f(rmin) rmax = f(rmax) r = f(r) # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) # compute values pos = transpose(array(r)).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_GenericMap_On_Spherical_1d_Grid(nb,nr,rmax,val,f=None,fm=None): ''' Return an array of points containing mass*val ''' rmin = 0.0 rmax = float(rmax) r = nb.rxyz() # 0 -> rmax # here, we scale if f!=None: rmin = f(rmin) rmax = f(rmax) r = f(r) # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) # compute values pos = transpose(array(r)).astype(float32) val = val.astype(float32) mass= nb.mass.astype(float32) shape = (nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_NumberMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing number of particles ''' rmin = 0.0 rmax = float(rmax) r = nb.rxyz() # 0 -> rmax # here, we scale if f!=None: rmin = f(rmin) rmax = f(rmax) r = f(r) # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) # compute values pos = transpose(array(r)).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nr,) # make the map mat = mapping.mkmap1d(pos,mass,val,shape) return mat def get_PotentialMap_On_Spherical_1d_Grid(nb,nr,rmax,eps,Tree=None,f=None,fm=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Spherical_1d_Grid(nr,rmax,f=f,fm=fm) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nr,) return pot def get_SurfaceMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing surface (volume) of each cell. ''' rmin = 0.0 rmax = float(rmax) # compute volume (surface) #def volume(ir): # dr = (rmax/nr) + (ir-ir) # return dr # ## make the map #mat = fromfunction(volume,(nr,)) #mat = mat.astype(float32) r1 = get_r_Of_Spherical_1d_Grid(nr,rmax,offr=0,f=f,fm=fm) r2 = get_r_Of_Spherical_1d_Grid(nr,rmax,offr=1,f=f,fm=fm) mat = r2-r1 mat = mat.astype(float32) return mat def get_VolumeMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' rmin = 0.0 rmax = float(rmax) # compute volume (surface) #def volume(ir): # dr = (rmax/nr) # r = ir*dr # return 4.0/3.0*pi*( (r+dr)**3 - (r)**3 ) # # make the map #mat = fromfunction(volume,(nr,)) #mat = mat.astype(float32) r1 = get_r_Of_Spherical_1d_Grid(nr,rmax,offr=0,f=f,fm=fm) r2 = get_r_Of_Spherical_1d_Grid(nr,rmax,offr=1,f=f,fm=fm) mat = 4.0/3.0*pi*( (r2)**3 - (r1)**3 ) mat = mat.astype(float32) return mat def get_DensityMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Spherical_1d_Grid(nb,nr,rmax,f=f,fm=fm) v = get_VolumeMap_On_Spherical_1d_Grid(nb,nr,rmax,f=f,fm=fm) return m/v def get_LinearDensityMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing the linear density in each cell ''' m = get_MassMap_On_Spherical_1d_Grid(nb,nr,rmax,f=f,fm=fm) s = get_SurfaceMap_On_Spherical_1d_Grid(nb,nr,rmax,f=f,fm=fm) return m/s def get_AccumulatedMassMap_On_Spherical_1d_Grid(nb,nr,rmax,f=None,fm=None): ''' Return an array of points containing M(r) in each cell ''' m = get_MassMap_On_Spherical_1d_Grid(nb,nr,rmax,f=f,fm=fm) mr = get_Accumulation_Along_Axis(m,axis=0) return mr def get_Interpolation_On_Spherical_1d_Grid(pos,mat,nr,rmax,offr=0,f=None,fm=None): ''' Interpolates continuous value of pos, using matrix mat ''' rmin = 0.0 rmax = float(rmax) nr = int(nr) r = sqrt(pos[:,0]**2+pos[:,1]**2+pos[:,2]**2) # here, we scale if f!=None: rmin = f(rmin) rmax = f(rmax) r = f(r) ir = (r-rmin)/(rmax-rmin)* nr - offr return myNumeric.Interpolate_From_1d_Array(ir.astype(float32),mat.astype(float32)) ############## # 3 dimensions ############## def get_rpt_Of_Spherical_3d_Grid(nr,np,nt,rmax,offr=0,offp=0,offt=0): rmin = 0.0 rmax = float(rmax) pmin = 0.0 pmax = 2*pi tmin = -pi/2 tmax = pi/2 ir = arange(nr) ip = arange(np) it = arange(nt) r = (ir+offr)*(rmax-rmin)/nr + rmin p = (ip+offp)*(pmax-pmin)/np + pmin t = (it+offt)*(tmax-tmin)/nt + tmin return r,p,t def get_Points_On_Spherical_3d_Grid(nr,np,nt,rmax,offr=0,offp=0,offt=0): ''' Return an array of points corresponding to the nodes of a 3d spherical grid To get a nt X nr array from the returned vector (pos), do x = pos[:,0] y = pos[:,0] z = pos[:,0] x.shape = (nr,np,nt) y.shape = (nr,np,nt) z.shape = (nr,np,nt) ''' rmin = 0.0 rmax = float(rmax) pmin = 0.0 pmax = 2*pi tmin = -pi/2 tmax = pi/2 nr = int(nr) np = int(np) nt = int(nt) ir,ip,it = indices((nr,np,nt)) r = (ir+offr)*(rmax-rmin)/nr + rmin p = (ip+offp)*(pmax-pmin)/np + pmin t = (it+offt)*(tmax-tmin)/nt + tmin r = ravel(r) p = ravel(p) t = ravel(t) x = r*cos(p)*cos(t) y = r*sin(p)*cos(t) z = r *sin(t) return transpose(array([x,y,z])).astype(float32) def get_MassMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) pmin = 0.0 pmax = 2*pi tmin = -pi/2 tmax = pi/2 r = nb.rxyz() # 0 -> rmax p = nb.phi_xyz() + pi # 0 -> pi t = nb.theta_xyz() # -pi/2 -> pi/2 # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) p = (p-pmin)/(pmax-pmin) t = (t-tmin)/(tmax-tmin) # compute values pos = transpose(array([r,p,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= nb.mass.astype(float32) shape = (nr,np,nt) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_NumberMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax): ''' Return an array of points containing mass of particles ''' rmin = 0.0 rmax = float(rmax) pmin = 0.0 pmax = 2*pi tmin = -pi/2 tmax = pi/2 r = nb.rxyz() # 0 -> rmax p = nb.phi_xyz() + pi # 0 -> pi t = nb.theta_xyz() # -pi/2 -> pi/2 # scale between 0 and 1 r = (r-rmin)/(rmax-rmin) p = (p-pmin)/(pmax-pmin) t = (t-tmin)/(tmax-tmin) # compute values pos = transpose(array([r,p,t])).astype(float32) val = ones(pos.shape).astype(float32) mass= ones(pos.shape).astype(float32) shape = (nr,np,nt) # make the map mat = mapping.mkmap3d(pos,mass,val,shape) return mat def get_PotentialMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax,eps,Tree=None): ''' Return an array of points containing potential ''' pos = get_Points_On_Spherical_3d_Grid(nr,np,nt,rmax) Tree = nb.getTree() pot = Tree.Potential(pos,eps) # transform back into array pot.shape = (nr,np,nt) return pot def get_VolumeMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax): ''' Return an array of points containing corresponding physical volumes of each cell (usefull to compute density) ''' print "get_VolumeMap_On_Spherical_3d_Grid is probably wrong" print "we should add a cos(t) or somthing similar" sys.exit() rmin = 0.0 rmax = float(rmax) pmin = 0.0 pmax = 2*pi tmin = -pi/2 tmax = pi/2 # compute volume (surface) def volume(ir,ip,it): dr = (rmax/nr) r = ir*dr return 4.0/3.0*pi*( (r+dr)**3 - (r)**3 ) / np / nt # make the map mat = fromfunction(volume,(nr,np,nt)) mat = mat.astype(float32) return mat def get_DensityMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax): ''' Return an array of points containing density in each cell ''' m = get_MassMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax) v = get_VolumeMap_On_Spherical_3d_Grid(nb,nr,np,nt,rmax) return m/v diff --git a/pNbody/main.py b/pNbody/main.py index 9d9df4d..4c8d690 100644 --- a/pNbody/main.py +++ b/pNbody/main.py @@ -1,5753 +1,5775 @@ # -*- coding: utf-8 -*- # some standard modules import os,sys,string,types,glob from copy import deepcopy import warnings # array module from numpy import * from numpy import clip as numclip from numpy import random as RandomArray # module that init parameters from parameters import * # nbody python modules import io from libutil import * from palette import * import geometry as geo import fourier import param import liblog import libgrid import libdisk import libutil import nbdrklib # nbody C modules from myNumeric import * from mapping import * from nbodymodule import * # Gtools module (now integrated in nbody) #import Gtools as Gt import units import ctes import thermodyn import coolinglib import cosmo import treelib import asciilib try: import ptreelib except: pass try: import libqt except: pass try: import SM except: pass try: # all this is usefull to read files from mpi4py import MPI except: MPI = None import mpi # maybe we should send mpi instead of MPI FLOAT = float #################################################################################################################################### # # DEFAULT CLASS NBODY # #################################################################################################################################### class NbodyDefault: ''' This is the reference Nbody class. This is the constructor for the **Nbody** object. Optional arguments are: p_name : name of the file in case of multiple files, files must be included in a list ["file1","file2"] pos : positions (3xN array) vel : positions (3xN array) mass : positions (1x array) num : id of particles (1xN array) tpe : type of particles (1xN array) ftype : type of input file (binary,ascii) status : 'old' : open an old file 'new' : create a new object byteorder : 'little' or 'big' pio : parallel io : 'yes' or 'no' local : True=local object, False=global object (paralellized) Not implemeted Yet log : log file unitsfile : define the type of units by default this class initialize the following variables : self.p_name : name of the file(s) to read or write self.pos : array of positions self.vel : array of velocities self.mass : array of masses self.num : array of id self.tpe : array of types self.ftype : type of the file self.status : object status ('old' or 'new') self.byteorder : byter order ('little' or 'big') self.pio : parallel io ('yes' or 'no') self.log : log object # new variables self.nbody : local number of particles self.nbody_tot : total number of particles self.mass_tot : total mass self.npart : number of particles of each type self.npart_tot : total number of particles of each type self.spec_vars : dictionary of variables specific for the format used self.spec_vect : dictionary of vector specific for the format used ''' def __init__(self,p_name=None,pos=None,vel=None,mass=None,num=None,tpe=None,ftype=None,status='old',byteorder=sys.byteorder,pio='no',local=False,log=None,unitsfile=None): ################################# # init vars ################################# /home/leo/.pNbody/formats/gadget.py if p_name == None: status = 'new' self.set_filenames(p_name,pio=pio) self.pos = pos self.vel = vel self.mass = mass self.num = num self.tpe = tpe self.ftype = self.__class__.__name__ self.status = status self.byteorder = byteorder self.pio = pio self.log = log self.nbody = None self.nbody_tot = None self.mass_tot = None self.npart = None self.npart_tot = None self.unitsfile = unitsfile ################################# # init units ################################# self.init_units() ################################# # init other parameters ################################# self.parameters = param.Params(PARAMETERFILE,None) self.defaultparameters = self.parameters.get_dic() # log if self.log == None: self.log = liblog.Log(os.path.join(HOME,'.nbodylog'),show='yes') ################################################### # in case of an old file, open and read the file(s) ################################################### if status=='old': self.read() ################################################### # in case of a new file ################################################### elif status=='new': for i in range(len(self.p_name)): if self.p_name[i] == None: self.p_name[i] = 'file.dat' ################################################### # final initialisation ################################################### self.init() ################################################### # check consistency ################################################### # to be done ################################# # # init functions # ################################# ################################# def init(self): ################################# ''' Initialize normal and specific class variables ''' # 1) find the number of particles self.nbody = self.get_nbody() # 2) define undefined vectors if self.pos == None: self.pos = zeros((self.nbody,3),float32) self.pos = self.pos.astype(float32) else: self.pos = self.pos.astype(float32) if self.vel == None: self.vel = zeros((self.nbody,3),float32) self.vel = self.vel.astype(float32) else: self.vel = self.vel.astype(float32) if self.mass == None: self.mass = ones((self.nbody, ),float32)/self.nbody self.mass = self.mass.astype(float32) else: self.mass = self.mass.astype(float32) if self.tpe == None: self.tpe = zeros(self.nbody,int) self.tpe = self.tpe.astype(int) else: self.tpe = self.tpe.astype(int) if self.num == None: self.num = self.get_num() self.num = self.num.astype(int) else: self.num = self.num.astype(int) # 3) other variables self.nbody_tot = self.get_nbody_tot() self.mass_tot = self.get_mass_tot() self.npart = self.get_npart() self.npart_tot = self.get_npart_tot() # Init specific class variables # (may be redundant with make_specific_variables_global) self.spec_vars = self.get_default_spec_vars() list_of_vars = self.get_list_of_vars() for name in self.spec_vars.keys(): try: list_of_vars.index(name) except ValueError: setattr(self, name, self.spec_vars[name]) # Init specific class vectors self.spec_vect = self.get_default_spec_array() list_of_vect = self.get_list_of_array() for name in self.spec_vect.keys(): try: list_of_vect.index(name) except ValueError: setattr(self, name, ones(self.nbody,self.spec_vect[name][1])*self.spec_vect[name][0]) # init specific parameters self.InitSpec() # sph parameters/variables self.InitSphParameters() ################################# def InitSpec(self): ################################# """ This function allows to initialize specific parameters. It must be defined in format files. """ pass ################################# def get_format_file(self): ################################# "return the format file" return self._formatfile + + ################################# + def get_ftype(self,ftype='binary'): + ################################# + """ + get the current used format + """ + return self.ftype ################################# def set_ftype(self,ftype='binary'): ################################# """ Change the type of the file ftype : type of the file """ if mpi.NTask > 1: raise "Warning","set_ftype function is currently not suported with multi proc." new = Nbody(status='new',ftype=ftype) # now, copy all var linked to the model for name in self.get_list_of_vars(): if name != 'ftype': setattr(new, name, getattr(self,name)) # now, copy all array linked to the model for name in self.get_list_of_array(): vec = getattr(self,name) setattr(new, name, vec) # other vars new.init() return new ################################# def get_num(self): ################################# """ Compute the num variable in order to be consistent with particles types """ # compute npart_all if self.npart == None: npart = self.get_npart() else: npart = self.npart npart_all = array(mpi.mpi_allgather(npart)) return mpi.mpi_sarange(npart_all) # + 1 ################################# def get_default_spec_vars(self): ################################# ''' return specific variables default values for the class ''' return {} ################################# def get_default_spec_array(self): ################################# ''' return specific array default values for the class ''' return {} ################################# def set_pio(self,pio): ################################# """ Set parallel input/output or not io pio : 'yes' or 'no' """ self.pio = pio self.set_filenames(self.p_name_global,pio=pio) if pio=='yes': self.num_files = mpi.NTask else: self.num_files = 1 ################################# def rename(self,p_name=None): ################################# """ Rename the files p_name : new name(s) """ if p_name != None: self.set_filenames(p_name,pio=self.pio) ################################# def set_filenames(self,p_name,pio=None): ################################# """ Set the local and global names p_name : new name(s) pio : 'yes' or 'no' """ if type(p_name) == types.ListType: self.p_name_global = [] self.p_name = [] for name in p_name: if pio == 'yes': self.p_name_global.append(name) self.p_name.append("%s.%d"%(name,mpi.mpi_ThisTask())) else: self.p_name_global.append(name) self.p_name.append(name) else: if pio == 'yes': self.p_name_global = [p_name] self.p_name = ["%s.%d"%(p_name,mpi.mpi_ThisTask())] else: self.p_name_global = [p_name] self.p_name = [p_name] ################################# def get_ntype(self): ################################# """ return the number of paticles types """ return len(self.npart) ################################# def get_nbody(self): ################################# """ Return the local number of particles. """ if self.pos != None: nbody = len(self.pos) elif self.vel != None: nbody = len(self.vel) elif self.mass != None: nbody = len(self.mass) elif self.num != None: nbody = len(self.num) elif self.tpe != None: nbody = len(self.tpe) else: nbody = 0 return nbody ################################# def get_nbody_tot(self): ################################# """ Return the total number of particles. """ nbody_tot = mpi.mpi_allreduce(self.nbody) return nbody_tot ################################# def get_npart(self): ################################# """ Return the local number of particles of each types, based on the variable tpe """ npart = array([],int) n = 0 if self.tpe==None: return npart.tolist() for tpe in range(self.get_mxntpe()): np = sum( (self.tpe==tpe).astype(int) ) npart = concatenate((npart,array([np]))) n = n + np if n != self.nbody: print "get_npart : n (=%d) is different from self.nbody (=%d)"%(n,self.nbody) raise "get_npart : n is different from self.nbody" return npart.tolist() ################################# def get_npart_tot(self): ################################# """ Return the total number of particles of each types. """ npart = array(self.npart) npart_tot = mpi.mpi_allreduce(npart) npart_tot = npart_tot.tolist() return npart_tot ################################# def get_npart_all(self,npart_tot,NTask): ################################# ''' From npart_tot, the total number of particles per type, return npart_per_proc, an array where each element corresponds to the value of npart of each process. ''' if (type(npart_tot) != types.ListType) and (type(npart_tot) !=ndarray): npart_tot = array([npart_tot]) ntype = len(npart_tot) npart_all = zeros((NTask,ntype)) for i in range(len(npart_tot)): for Task in range(NTask-1,-1,-1): npart_all[Task,i] = npart_tot[i]/NTask + npart_tot[i]%NTask*(Task==0) return npart_all ################################# def get_npart_and_npart_all(self,npart): ################################# ''' From npart (usually read for the header of a file), compute : npart : number of particles in each type npart_tot : total number of particles in each type npart_all : npart for each process. ''' ################################# def get_mxntpe(self): ################################# ''' Return the max number of type for this format ''' return 6 ################################# def make_default_vars_global(self): ################################# ''' Make specific variables global ''' self.spec_vars = self.get_default_spec_vars() for name in self.spec_vars.keys(): if not self.has_var(name): setattr(self, name, self.spec_vars[name]) ################################# def set_npart(self,npart): ################################# """ Set the local number of particles of each types. This function modifies the variable self.tpe """ if sum(array(npart)) > self.nbody: raise "Error (set_npart)","sum(npart) is greater than nbody" i = 0 n0 = 0 for n in npart: self.tpe[n0:n0+n] = ones(n)*i i = i + 1 n0 = n0+n self.tpe[n0:self.nbody] = ones(self.nbody-n0)*i self.npart = self.get_npart() self.npart_tot = self.get_npart_tot() ################################# def set_tpe(self,tpe): ################################# """ Set all particles to the type tpe """ self.tpe = ones(self.nbody)*tpe self.npart = self.get_npart() self.npart_tot = self.get_npart_tot() ################################# # # parameters functions # ################################# ''' Warning, these routines are a bit bad... ''' def set_parameters(self,params): ''' Set parameters for the class ''' self.parameters = param.Params(PARAMETERFILE,None) self.parameters.params = params.params self.defaultparameters = self.parameters.get_dic() ################################# # # units functions # ################################# ''' There is several ways to set the units in pNbody In an object, the units are stored in self.localsystem_of_units which is a UnitSystem object defined in units.py We define a unit system by giving Unit_lenght, Unit_mass, Unit_time, Unit_K, Unit_mol, and Unit_C Actually only Unit_lenght, Unit_mass, Unit_time are used, all are Units object (units.py) Following Gadget2, easy ways to definde units is to give three floats, UnitVelocity_in_cm_per_s UnitMass_in_g UnitLength_in_cm This is done using the method self.set_local_system_of_units() which uses UnitVelocity_in_cm_per_s,UnitMass_in_g,UnitLength_in_cm if they are given, or read a gadget parameter file or read a pNbody unitsparameter file or use the default unitsparameter file. ''' def init_units(self): ''' This function is responsible for the units initialization. It will create : self.unitsparameters that contains parameters like - the hydrogen mass fraction, - the metalicity ionisation flag - the adiabatic index - ... and self.localsystem_of_units a UnitSystem object that really defines the system of units in the Nbody object. It uses the values : UnitLength_in_cm UnitMass_in_g UnitVelocity_in_cm_per_s All physical values computed in pNbody should use self.localsystem_of_units to be converted in other units. self.unitsparameters is usefull if other parameters needs to be known, like the adiabatic index, etc. ''' self.unitsparameters = param.Params(UNITSPARAMETERFILE,None) if self.unitsfile!=None: ############################################################## # 1) this part should be only in the gadget.py format file, no ? BOF, non # 2) we could simplify using self.set_local_system_of_units() # 3) and some options -> but this needs to update self.unitsparameters ############################################################## # if it is a gadget parameter file try: gparams = io.read_params(self.unitsfile) self.unitsparameters.set('HubbleParam', gparams['HubbleParam']) self.unitsparameters.set('UnitLength_in_cm', gparams['UnitLength_in_cm']) self.unitsparameters.set('UnitMass_in_g', gparams['UnitMass_in_g']) self.unitsparameters.set('UnitVelocity_in_cm_per_s',gparams['UnitVelocity_in_cm_per_s']) # those parameters may be in the header of the file self.unitsparameters.set('Omega0', gparams['Omega0']) self.unitsparameters.set('OmegaLambda', gparams['OmegaLambda']) self.unitsparameters.set('OmegaBaryon', gparams['OmegaBaryon']) self.unitsparameters.set('BoxSize', gparams['BoxSize']) self.unitsparameters.set('ComovingIntegrationOn', gparams['ComovingIntegrationOn']) #self.set_local_system_of_units(gadgetparameterfile=self.unitsfile) except: # try to read a pNbody units file try: self.unitsparameters = param.Params(self.unitsfile,None) #self.set_local_system_of_units(unitparameterfile=self.unitsfile) except: raise IOError(015,'format of unitsfile %s unknown ! Pease check.'%(self.unitsfile)) # define local system of units it it does not exists #if not self.has_var("localsystem_of_units"): self.set_local_system_of_units() # print info #self.localsystem_of_units.info() def set_unitsparameters(self,unitsparams): ''' Set units parameters for the class. ''' print "!!!!!! in set_unitsparameters !!!!" print "!!!!!! this is bad !!!! we should never use UNITSPARAMETERFILE" print "!!!!!! this is bad !!!! we should never use UNITSPARAMETERFILE" self.unitsparameters = param.Params(UNITSPARAMETERFILE,None) self.unitsparameters.params = unitsparams.params self.set_local_system_of_units() def set_local_system_of_units(self,params=None,UnitLength_in_cm=None,UnitVelocity_in_cm_per_s=None,UnitMass_in_g=None,unitparameterfile=None,gadgetparameterfile=None): ''' Set local system of units using UnitLength_in_cm,UnitVelocity_in_cm_per_s,UnitMass_in_g 1) if nothing is given, we use self.unitsparameters to obtain these values 2) if UnitLength_in_cm UnitVelocity_in_cm_per_s UnitMass_in_g are given, we use them 2b) if UnitLength_in_cm,UnitVelocity_in_cm_per_s,UnitMass_in_g are given in a dictionary 3) if unitparameterfile is given we read the parameters from the file (units parameter format) 4) if gadgetparameterfile is given we read the parameters from the file (gadget param format) ''' if gadgetparameterfile!=None: params = io.read_params(gadgetparameterfile) #print "Units Set From %s"%gadgetparameterfile elif unitparameterfile!=None: unitsparameters = param.Params(unitparameterfile,None) params = {} params['UnitLength_in_cm'] = unitsparameters.get('UnitLength_in_cm') params['UnitVelocity_in_cm_per_s'] = unitsparameters.get('UnitVelocity_in_cm_per_s') params['UnitMass_in_g'] = unitsparameters.get('UnitMass_in_g') print "Units Set From %s"%unitparameterfile elif params!=None: print "Units Set From %s"%params elif UnitLength_in_cm!=None and UnitVelocity_in_cm_per_s!=None and UnitMass_in_g!=None: params = {} params['UnitLength_in_cm'] = UnitLength_in_cm params['UnitVelocity_in_cm_per_s'] = UnitVelocity_in_cm_per_s params['UnitMass_in_g'] = UnitMass_in_g print "Units Set From UnitLength_in_cm,UnitVelocity_in_cm_per_s,UnitMass_in_g" else: params = {} params['UnitLength_in_cm'] = self.unitsparameters.get('UnitLength_in_cm') params['UnitVelocity_in_cm_per_s'] = self.unitsparameters.get('UnitVelocity_in_cm_per_s') params['UnitMass_in_g'] = self.unitsparameters.get('UnitMass_in_g') print "Units Set From %s (%s)"%("self.unitsparameters",self.unitsparameters.filename) # now, create the self.localsystem_of_units = units.Set_SystemUnits_From_Params(params) ################################# # # info functions # ################################# ################################# def info(self): ################################# """ Write info """ infolist = [] infolist.append("-----------------------------------") if mpi.NTask>1: infolist.append("") infolist.append("ThisTask : %s"%mpi.ThisTask.__repr__()) infolist.append("NTask : %s"%mpi.NTask.__repr__()) infolist.append("") infolist.append("particle file : %s"%self.p_name.__repr__()) infolist.append("ftype : %s"%self.ftype.__repr__()) infolist.append("mxntpe : %s"%self.get_mxntpe().__repr__()) infolist.append("nbody : %s"%self.nbody.__repr__()) infolist.append("nbody_tot : %s"%self.nbody_tot.__repr__()) infolist.append("npart : %s"%self.npart.__repr__()) infolist.append("npart_tot : %s"%self.npart_tot.__repr__()) infolist.append("mass_tot : %s"%self.mass_tot.__repr__()) infolist.append("byteorder : %s"%self.byteorder.__repr__()) infolist.append("pio : %s"%self.pio.__repr__()) if self.nbody != 0: infolist.append("") infolist.append("len pos : %s"%len(self.pos).__repr__()) infolist.append("pos[0] : %s"%self.pos[0].__repr__()) infolist.append("pos[-1] : %s"%self.pos[-1].__repr__()) infolist.append("len vel : %s"%len(self.vel).__repr__()) infolist.append("vel[0] : %s"%self.vel[0].__repr__()) infolist.append("vel[-1] : %s"%self.vel[-1].__repr__()) infolist.append("len mass : %s"%len(self.mass).__repr__()) infolist.append("mass[0] : %s"%self.mass[0].__repr__()) infolist.append("mass[-1] : %s"%self.mass[-1].__repr__()) infolist.append("len num : %s"%len(self.num).__repr__()) infolist.append("num[0] : %s"%self.num[0].__repr__()) infolist.append("num[-1] : %s"%self.num[-1].__repr__()) infolist.append("len tpe : %s"%len(self.tpe).__repr__()) infolist.append("tpe[0] : %s"%self.tpe[0].__repr__()) infolist.append("tpe[-1] : %s"%self.tpe[-1].__repr__()) if self.spec_info()!=None: infolist = infolist + self.spec_info() all_infolist = mpi.mpi_allgather(infolist) if mpi.mpi_IsMaster(): for infolist in all_infolist: for line in infolist: #print line self.log.write(line) ################################# def spec_info(self): ################################# """ Write specific info """ return None ################################# def object_info(self): ################################# """ Write class(object) info """ list_of_vars = self.get_list_of_vars() list_of_array = self.get_list_of_array() self.log.write("#############################") self.log.write("list of vars") self.log.write("#############################") for name in list_of_vars: self.log.write("%s %s"%( name,str(type(getattr(self,name))))) self.log.write("#############################") self.log.write("list of arrays") self.log.write("#############################") for name in list_of_array: self.log.write("%s %s"%(name,str(type(getattr(self,name))))) ################################# def nodes_info(self): ################################# """ Write info on nodes """ all_npart = mpi.mpi_allgather(self.npart) all_nbody = mpi.mpi_allgather(self.nbody) if mpi.mpi_IsMaster(): for Task in range(mpi.NTask): line = "Task=%4d nbody=%10d"%(Task,all_nbody[Task]) line = line + " npart= " for npart in all_npart[Task]: line = line + "%10d "%npart self.log.write(line) ################################# def memory_info(self): ################################# """ Write info on memory size of the current object (only counting arrays size) """ total_size = 0 array_size = 0 elts = self.get_list_of_array() for elt in elts: #num_of_elts = getattr(self,elt).size #byte_per_elts = getattr(self,elt).itemsize #bytes = num_of_elts*byte_per_elts bytes = getattr(self,elt).nbytes total_size = total_size + bytes array_size = array_size + bytes print "(%d) %10s %14d"%(mpi.ThisTask,elt,bytes) #elts = self.get_list_of_vars() #for elt in elts: array_size = mpi.mpi_reduce(array_size) # only the master return the info total_size = mpi.mpi_reduce(total_size) if mpi.mpi_IsMaster(): print "total size = %d octets"%(total_size) if array_size < 1024: print "total arrays size = %d octets"%(array_size) else: array_size = array_size/1024.0 if array_size < 1024: print "total arrays size = %dK"%(array_size) else: array_size = array_size/1024.0 if array_size < 1024: print "total arrays size = %dM"%(array_size) else: array_size = array_size/1024.0 if array_size < 1024: print "total arrays size = %dG"%(array_size) ################################# def print_filenames(self): ################################# """ Print files names """ self.log.write("p_name_global = %s"%str(self.p_name_global)) self.log.write("p_name = %s"%str(self.p_name)) ################################# # # list of variables functions # ################################# def get_list_of_array(self): """ Return the list of numpy vectors of size nbody. """ list_of_arrays = [] for name in dir(self): if type(getattr(self,name)) == ndarray: if len(getattr(self,name)) == self.nbody: #if (name!="nall") and (name!="nallhw") and (name!="massarr") and (name!="npart") and (name!="npart_tot"): list_of_arrays.append(name) return list_of_arrays def get_list_of_method(self): """ Return the list of instance methods (functions). """ list_of_instancemethod = [] for name in dir(self): if type(getattr(self,name)) == types.MethodType: list_of_instancemethod.append(name) return list_of_instancemethod def get_list_of_vars(self): """ Get the list of vars that are linked to the model """ list_of_allvars = dir(self) list_of_arrays = self.get_list_of_array() list_of_method = self.get_list_of_method() for name in list_of_arrays: list_of_allvars.remove(name) for name in list_of_method: list_of_allvars.remove(name) #list_of_allvars.remove('log') #list_of_allvars.remove('read_fcts') # becose these vars are linked to fcts #list_of_allvars.remove('write_fcts') # should be definitely removed return list_of_allvars def has_var(self,name): ''' Return true if the object pNbody has a variable called self.name ''' get_list_of_vars = self.get_list_of_vars() try: getattr(self,name) return True except AttributeError: return False def has_array(self,name): ''' Return true if the object pNbody has an array called self.name ''' list_of_array = self.get_list_of_array() try: list_of_array.index(name) return True except ValueError: return False def find_vars(self): ''' This function return a list of variables defined in the current object ''' elts = dir(self) lst = [] for elt in elts: exec("obj = self.%s"%(elt)) if type(obj) != types.MethodType: lst.append(elt) return lst ################################# # # check special values # ################################# def check_arrays(self): ''' check if the array contains special values like NaN or Inf ''' status = 0 for name in self.get_list_of_array(): vec = getattr(self,name) # check nan if isnan(vec).any(): msg = "array %s contains Nan !!!"%name warnings.warn(msg) status = 1 # check nan if isinf(vec).any(): msg = "array %s contains Inf !!!"%name warnings.warn(msg) status = 1 return status ################################# # # read/write functions # ################################# def read(self): """ Read the particle file(s) """ for i in range(len(self.p_name)): self.open_and_read(self.p_name[i],self.get_read_fcts()[i]) self.make_default_vars_global() def open_and_read(self,name,readfct): ''' open and read file name name : name of the input readfct : function used to read the file ''' # check p_name if self.pio=='yes' or mpi.mpi_IsMaster(): io.checkfile(name) # get size if self.pio=='yes' or mpi.mpi_IsMaster(): isize = os.path.getsize(name) # open file if self.pio=='yes' or mpi.mpi_IsMaster(): f = open(name,'r') else: f = None # read the file readfct(f) if self.pio=='yes' or mpi.mpi_IsMaster(): fsize = f.tell() else: fsize = None if self.pio=='yes' or mpi.mpi_IsMaster(): if fsize != isize: raise "ReadError","file %s not red completely"%(name) # close file if self.pio=='yes' or mpi.mpi_IsMaster(): f.close() + def get_read_fcts(self): + """ + returns the functions needed to read a snapshot file. + """ + return [] + def write(self): """ Write the particle file(s) """ for i in range(len(self.p_name)): self.open_and_write(self.p_name[i],self.get_write_fcts()[i]) def open_and_write(self,name,writefct): """ Open and write file name : name of the output writefct : function used to write the file """ if self.pio=='yes' or mpi.mpi_IsMaster(): f = open(name,'w') else: f = None writefct(f) if self.pio=='yes' or mpi.mpi_IsMaster(): f.close() + def get_write_fcts(self): + """ + returns the functions needed to write a snapshot file. + """ + return [] + + + def write_num(self,name): """ Write a num file name : name of the output """ if self.pio =='yes': f = open("%s.%d"%(name,mpi.ThisTask),'w') for n in self.num: f.write('%8i\n'%(n)) f.close() else: if mpi.mpi_IsMaster(): f = open(name,'w') for Task in range(mpi.NTask-1,-1,-1): if Task != 0: num = mpi.mpi_recv(source = Task) for n in num: f.write('%8i\n'%(n)) else: for n in self.num: f.write('%8i\n'%(n)) else: mpi.mpi_send(self.num, dest = 0) def read_num(self,name): """ Read a num file name : name of the input """ ################################# # # coordinate transformation # ################################# ################################# # positions ################################# def x(self): """ Return a 1xn float array containing x coordinate """ return self.pos[:,0] def y(self): """ Return a 1xn float array containing y coordinate """ return self.pos[:,1] def z(self): """ Return a 1xn float array containing z coordinate """ return self.pos[:,2] def rxyz(self,center=None): """ Return a 1xn float array that corresponds to the distance from the center of each particle. """ if center!=None: r = sqrt( (self.pos[:,0]-center[0])**2 + (self.pos[:,1]-center[1])**2 + (self.pos[:,2]-center[2])**2 ) else: r = sqrt( self.pos[:,0]**2 + self.pos[:,1]**2 + self.pos[:,2]**2 ) return r def phi_xyz(self): """ Return a 1xn float array that corresponds to the azimuth in spherical coordinate of each particle. """ r = self.rxyz() rxy = self.rxy() xp = self.pos[:,0]*r/rxy # x projection in the plane yp = self.pos[:,1]*r/rxy # y projection in the plane p = arctan2(yp,xp) return p def theta_xyz(self): """ Return a 1xn float array that corresponds to the elevation angle in spherical coordinate of each particle. """ r = self.rxyz() t = arcsin(self.pos[:,2]/r) return t def rxy(self): """ Return a 1xn float array that corresponds to the projected distance from the center of each particle. """ r = sqrt( self.pos[:,0]**2 + self.pos[:,1]**2) return r def phi_xy(self): """ Return a 1xn float array that corresponds to the azimuth in cylindrical coordinate of each particle. """ p = arctan2(self.pos[:,1],self.pos[:,0]) return p r = rxyz R = rxy ###################### # spherical coord ###################### def cart2sph(self,pos=None): """ Transform carthesian coodinates x,y,z into spherical coordinates r,p,t Return a 3xn float array. """ if pos!=None: x = pos[:,0] y = pos[:,1] z = pos[:,2] else: x = self.pos[:,0] y = self.pos[:,1] z = self.pos[:,2] r = self.rxyz() rxy = self.rxy() #xp = x*r/rxy # x projection in the plane #yp = y*r/rxy # y projection in the plane #p = arctan2(yp,xp) #t = arcsin(z/r) p = arctan2(y,x) t = arctan2(rxy,z) return transpose(array([r,p,t])).astype(float32) def sph2cart(self,pos=None): """ Transform spherical coordinates r,p,t into carthesian coodinates x,y,z Return a 3xn float array. """ if pos!=None: r = pos[:,0] p = pos[:,1] t = pos[:,2] else: r = self.pos[:,0] p = self.pos[:,1] t = self.pos[:,2] x = r*sin(t)*cos(p) y = r*sin(t)*sin(p) z = r*cos(t) return transpose(array([x,y,z])).astype(float32) ################################# # velocities ################################# def vx(self): """ Return a 1xn float array containing x velocity """ return self.vel[:,0] def vy(self): """ Return a 1xn float array containing y velocity """ return self.vel[:,1] def vz(self): """ Return a 1xn float array containing z velocity """ return self.vel[:,2] def vn(self): """ Return a 1xn float array that corresponds to the norm of velocities """ return sqrt(self.vel[:,0]*self.vel[:,0] + self.vel[:,1]*self.vel[:,1] + self.vel[:,2]*self.vel[:,2]) def vrxyz(self): """ Return a 1xn float array that corresponds to the radial velocity in spherical system """ r = self.rxyz() return (self.pos[:,0]*self.vel[:,0] + self.pos[:,1]*self.vel[:,1] + self.pos[:,2]*self.vel[:,2])/r def Vr(self): """ Return the radial velocies of particles The output is an 3xn float array. """ xr = sqrt(self.pos[:,0]**2+self.pos[:,1]**2) vr = (self.pos[:,0]*self.vel[:,0] + self.pos[:,1]*self.vel[:,1]) / xr return vr def Vt(self): """ Return the tangential velocies of particles The output is an 3xn float array. """ xr = sqrt(self.pos[:,0]**2+self.pos[:,1]**2) vt = (self.pos[:,0]*self.vel[:,1] - self.pos[:,1]*self.vel[:,0]) / xr return vt def Vz(self): """ Return a 1xn float array containing z velocity """ return self.vel[:,2] ###################### # cylindrical coord ###################### def vel_cyl2cart(self,pos=None,vel=None): """ Transform velocities in cylindrical coordinates vr,vt,vz into carthesian coodinates vx,vy,vz. Pos is the position of particles in cart. coord. Vel is the velocity in cylindrical coord. Return a 3xn float array. """ return vel_cyl2cart(self.pos,self.vel) def vel_cart2cyl(self): """ Transform velocities in carthesian coordinates vx,vy,vz into cylindrical coodinates vr,vz,vz. Pos is the position of particles in cart. coord. Vel is the velocity in cart. coord. Return a 3xn float array. """ return vel_cart2cyl(self.pos,self.vel) ################################# # # physical values # ################################# def get_ns(self): """ Return in an array the number of particles of each node. """ ns = mpi.mpi_allgather(self.nbody) return ns def get_mass_tot(self): """ Return the total mass of system. """ mass_tot = mpi.mpi_sum(self.mass) return mass_tot def size(self): """ Estimate the model size, using the inertial momentum """ return max(self.minert()) def cm(self): """ Return the mass center of the model. The output is an 3x1 float array. """ mtot = mpi.mpi_sum(self.mass.astype(FLOAT)) cmx = mpi.mpi_sum(self.pos[:,0].astype(float64)*self.mass.astype(FLOAT)) / mtot cmy = mpi.mpi_sum(self.pos[:,1].astype(float64)*self.mass.astype(FLOAT)) / mtot cmz = mpi.mpi_sum(self.pos[:,2].astype(float64)*self.mass.astype(FLOAT)) / mtot return array([cmx,cmy,cmz]) def get_histocenter(self,rbox=50,nb=500): """ Return the position of the higher density region in x,y,z (not good) found by the function "histocenter". rbox : size of the box nb : number of bins in each dimension """ rm = rbox/2. bins = arange(-rm,rm,float(2*rm)/float(nb)) # histograms in x,y,z (cut the tail) hx = mpi.mpi_histogram(self.pos[:,0],bins)[:-1] hy = mpi.mpi_histogram(self.pos[:,1],bins)[:-1] hz = mpi.mpi_histogram(self.pos[:,2],bins)[:-1] # max in each dim dx = bins[argmax(hx)] dy = bins[argmax(hy)] dz = bins[argmax(hz)] return array([dx,dy,dz]) def get_histocenter2(self,rbox=50,nb=64): """ Return the position of the higher density region in x,y,z (not good) found by the function "histocenter". rbox : size of the box nb : number of bins in each dimension """ # transformation -rbox->0, 0->nb/2, rbox->nb # y = (x+rbox)/(2*rbox)*nb pos = ( self.pos + [rbox,rbox,rbox] )/(2*rbox) # 0 to 1 pos = pos*[nb,nb,nb] # 0 to nb pos = pos.astype(float32) mass = self.mass.astype(float32) mat = mkmap3d(pos/nb,mass,mass,(nb,nb,nb)) # find max m = ravel(mat) arg = argmax(m) i = indices((nb,nb,nb)) # not that good ix = ravel(i[0]) # not that good iy = ravel(i[1]) # not that good iz = ravel(i[2]) # not that good ix = ix[arg] iy = iy[arg] iz = iz[arg] # transformation inverse # x = 2*rbox*(y/nb)-rbox dx = 2*rbox*(float(ix)/nb)-rbox dy = 2*rbox*(float(iy)/nb)-rbox dz = 2*rbox*(float(iz)/nb)-rbox return array([dx,dy,dz]) def cv(self): """ Return the center of the velocities of the model. The output is an 3x1 float array. """ cmx = mpi.mpi_sum(self.vel[:,0]*self.mass) / self.mass_tot cmy = mpi.mpi_sum(self.vel[:,1]*self.mass) / self.mass_tot cmz = mpi.mpi_sum(self.vel[:,2]*self.mass) / self.mass_tot return array([cmx,cmy,cmz]) def minert(self): """ Return the diagonal of the intertial momentum. """ mx = mpi.mpi_sum(self.pos[:,0]**2 *self.mass) / self.mass_tot my = mpi.mpi_sum(self.pos[:,1]**2 *self.mass) / self.mass_tot mz = mpi.mpi_sum(self.pos[:,2]**2 *self.mass) / self.mass_tot mx = sqrt(mx) my = sqrt(my) mz = sqrt(mz) return array([mx,my,mz]) def inertial_tensor(self): """ Return the inertial tensor. """ Ixx = mpi.mpi_sum(self.mass * (self.y()**2 + self.z()**2)) Iyy = mpi.mpi_sum(self.mass * (self.x()**2 + self.z()**2)) Izz = mpi.mpi_sum(self.mass * (self.x()**2 + self.y()**2)) Ixy = -mpi.mpi_sum(self.mass * self.x() * self.y()) Ixz = -mpi.mpi_sum(self.mass * self.x() * self.z()) Iyz = -mpi.mpi_sum(self.mass * self.y() * self.z()) I = array( [[Ixx,Ixy,Ixz],[Ixy,Iyy,Iyz],[Ixz,Iyz,Izz]] ) return I def x_sigma(self): """ Return the norm of the position dispersions. """ x = (self.pos - self.cm()) x2 = x[:,0]**2 + x[:,1]**2 + x[:,2]**2 x_s2 = mpi.mpi_sum( x2 *self.mass )/self.mass_tot x_s = sqrt(x_s2) return x_s def v_sigma(self): """ Return the norm of the velocity dispersions. """ v = (self.vel - self.cv()) v2 = v[:,0]**2 + v[:,1]**2 + v[:,2]**2 v_s2 = mpi.mpi_sum( v2 *self.mass )/self.mass_tot v_s = sqrt(v_s2) return v_s def dx_mean(self): """ Return the average distance between particles. """ # 1) estimate the size of the system D = self.x_sigma() # 2) estimate the # of particules per unit volume n = self.nbody_tot/D # 3) estimate the average distance between particules l = 1./n**(1./3.) return l def dv_mean(self): """ Return the average relative speed between particles. """ # 1) estimate the size of the system D = self.v_sigma() # 2) estimate the # of particules per unit volume n = self.nbody_tot/D # 3) estimate the average distance between particules l = 1./n**(1./3.) return l def Ekin(self): """ Return the total kinetic energy """ E = 0.5 * mpi.mpi_sum (self.mass * (self.vel[:,0]**2 + self.vel[:,1]**2 + self.vel[:,2]**2) ) return E def ekin(self): """ Return the total specific kinetic energy """ E = 0.5 * mpi.mpi_sum ( (self.vel[:,0]**2 + self.vel[:,1]**2 + self.vel[:,2]**2) ) return E def Epot(self,eps): """ Return the total potential energy using the softening lenght eps. eps : softening WARNING : THIS FUNCTION DO NOT WORK IN MPI MODE """ E = epot(self.pos,self.mass,eps) return E def epot(self,eps): """ Return the total specific potential energy using the softening lenght eps. eps : softening WARNING : THIS FUNCTION DO NOT WORK IN MPI MODE """ e = epot(self.pos,self.mass,eps)/self.mass_tot return e def L(self): """ Return the angular momentum in x,y,z of all particles. The output is an 3xn float array. """ l = amxyz(self.pos,self.vel,self.mass) return l def l(self): """ Return the specific angular momentum in x,y,z of all particles. The output is an 3xn float array. """ l = samxyz(self.pos,self.vel,self.mass) return l def Ltot(self): """ Return the total angular momentum. The output is an 3x1 float array. """ l = mpi.mpi_allreduce (am(self.pos,self.vel,self.mass)) #l = mpi.mpi_sum(self.L()) return l def ltot(self): """ Return the specific total angular momentum. The output is an 3x1 float array. """ l = mpi.mpi_allreduce (am(self.pos,self.vel,self.mass))/self.mass_tot #l = self.Ltot()/self.mass_tot return l def Pot(self,x,eps): """ Return the potential at a given position, using the softening lenght eps. x : position (array vector) eps : softening """ if type(x) == ndarray: p = zeros(len(x),float32) for i in range(len(x)): p[i] = mpi.mpi_allreduce ( potential(self.pos,self.mass,array(x[i],float32),eps) ) else: p = mpi.mpi_allreduce ( potential(self.pos,self.mass,array(x,float32),eps) ) return p def TreePot(self,pos,eps,Tree=None): """ Return the potential at a given position, using the softening lenght eps and using a tree. pos : position (array vector) eps : softening Tree: gravitational tree if already computed WARNING : this function do not work in parallel """ if Tree==None: self.Tree = Tree = self.getTree() pot = Tree.Potential(pos,eps) return pot def Accel(self,x,eps): """ Return the acceleration at a given position, using the softening lenght eps. x : position (array vector) eps : softening """ if type(x) == ndarray: ax = zeros(len(x),float32) ay = zeros(len(x),float32) az = zeros(len(x),float32) for i in range(len(x)): ax[i],ay[i],az[i] = acceleration(self.pos,self.mass,array(x[i],float32),eps) a = transpose(array([ax,ay,az],float32)) else: ax,ay,az = acceleration(self.pos,self.mass,array(x,float32),eps) ax = mpi.mpi_allreduce ( ax ) ay = mpi.mpi_allreduce ( ay ) az = mpi.mpi_allreduce ( az ) a = array([ax,ay,az],float32) return a def TreeAccel(self,pos,eps,Tree=None): """ Return the acceleration at a given position, using the softening lenght eps and using a tree. pos : position (array vector) eps : softening Tree: gravitational tree if already computed WARNING : this function do not work in parallel """ if Tree==None: self.Tree = Tree = self.getTree() acc = Tree.Acceleration(pos,eps) return acc def tork(self,acc): """ Return the total tork on the system due to the force acting on each particle (acc). The output is an 3xn float array. acc : 3xn float array """ trk = mpi.mpi_allreduce ( am(self.pos,array(acc,float32),self.mass) ) return trk def dens(self,r=None,nb=25,rm=50): """ Return the number density at radius r (supposing a spherical density distribution). If r is not specified, it is computed with nb and rm. The output is an n x 1 float array. !!! This routine do not use masses !!! r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) xr = sqrt(self.pos[:,0]**2 + self.pos[:,1]**2 + self.pos[:,2]**2) dens,r = histogram(xr,r) r1 = r[:-1] r2 = r[1:] dv = (4./3.)*pi*(r2**3-r1**3) dens = dens/dv # surface density # take the mean r = (r1+r2)/2 return r,mpi.mpi_allreduce(dens) def mdens(self,r=None,nb=25,rm=50): """ Return the density at radius r (supposing a spherical density distribution). If r is not specified, it is computed with nb and rm. The output is an n x 1 float array. r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) xr = sqrt(self.pos[:,0]**2 + self.pos[:,1]**2 + self.pos[:,2]**2) dens = whistogram(xr.astype(float),self.mass.astype(float),r.astype(float)) r1 = r[:-1] r2 = r[1:] dv = (4./3.)*pi*(r2**3-r1**3) dens = dens[:-1]/dv # surface density # take the mean r = (r1+r2)/2 return r,mpi.mpi_allreduce(dens) def mr(self,r=None,nb=25,rm=50): """ Return the mass inside radius r (supposing a spherical density distribution). If r is not specified, it is computed with nb and rm. The output is an n x 1 float array. r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) xr = self.rxyz() mr = whistogram(xr.astype(float),self.mass.astype(float),r.astype(float)) mr = add.accumulate(mr) return r,mpi.mpi_allreduce(mr) def Mr_Spherical(self,nr=25,rmin=0,rmax=50): """ Return the mass inside radius r (supposing a spherical density distribution). The output is 2 n x 1 float arrays. nr : number of bins (size of the output) rmin : minimal radius (this must be zero, instead it is wrong...) rmax : maximal radius """ rmin = float(rmin) rmax = float(rmax) shape = (nr,) val = ones(self.pos.shape).astype(float32) mass = self.mass.astype(float32) r = self.rxyz() r = (r-rmin)/(rmax-rmin) r = r.astype(float32) # compute the map mr = mkmap1d(r,mass,val,shape).astype(float) # compute the radii rs = arange(0.,rmax,(rmax-rmin)/nr) # sum mr = add.accumulate(mr) return rs,mpi.mpi_allreduce(mr) def sdens(self,r=None,nb=25,rm=50): """ Return the surface density at radius r. If r is not specified, it is computed with nb and rm. The output is an nx1 float array. !!! This routine do not uses masses !!! r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) xr = sqrt(self.pos[:,0]**2 + self.pos[:,1]**2) sdens,r = histogram(xr,r) r1 = r[:-1] r2 = r[1:] ds = pi*(r2**2-r1**2) sdens = sdens/ds # surface density # take the mean r = (r1+r2)/2. return r,mpi.mpi_allreduce(sdens) def msdens(self,r=None,nb=25,rm=50): """ Return the mass surface density at radius r. If r is not specified, it is computed with nb and rm. The output is an nx1 float array. r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) xr = sqrt(self.pos[:,0]**2 + self.pos[:,1]**2) sdens = whistogram(xr.astype(float),self.mass.astype(float),r.astype(float)) r1 = r[:-1] r2 = r[1:] ds = pi*(r2**2-r1**2) sdens = sdens[:-1]/ds # surface density # take the mean r = (r1+r2)/2. return r,mpi.mpi_allreduce(sdens) def sigma_z(self,r=None,nb=25,rm=50): """ Return the vertical dispertion in z at radius r. If r is not specified, it is computed with nb and rm. The output is an nx1 float array. r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) r,h = self.Histo(r,mode='sz') return r,h def sigma_vz(self,r=None,nb=25,rm=50): """ Return the vertical dispertion in z at radius r. If r is not specified, it is computed with nb and rm. The output is an nx1 float array. r : radius nb : number of bins (size of the output) rm : maximal radius """ if r!= None: r = array(r,float) else: rmax = rm dr = rm/float(nb) r = arange(0.,rm,dr) r,h = self.Histo(r,mode='svz') return r,h def zprof(self,z=None,r=2.5,dr=0.5,nb=25,zm=5.): """ Return the z-profile in a vector for a given radius !!! This routine works only if particles have equal masses !!! z : bins in z (optional) r : radius of the cut dr : width in r of the cut nb : number of bins (size of the output) zm : maximal height """ if z!= None: pass else: zmax = zm dz = 2.*zm/float(nb) z = arange(-zm,zm,dz) # select r1 = r-dr/2. r2 = r+dr/2. ann = self.selectc((self.rxy()>r1)*((self.rxy()r1[i])*((self.rxy() 1: sr.append(vr.std()) st.append(vt.std()) sz.append(vz.std()) mt.append(vt.mean()) else: sr.append(0.) st.append(0.) sz.append(0.) mt.append(0.) sr = array(sr,float) st = array(st,float) sz = array(sz,float) mt = array(mt,float) return r,sr,st,sz,mt def histovel(self,nb=100,vmin=None,vmax=None,mode='n'): """ Return or plot the histrogram of the norm of velocities or of the radial velocities. The output is a list (r,h) of 2 nx1 float arrays, where r is the radius and h the values of the histogram. nb : number of bins (size of the output) vmax : maximum velocity vmin : minimum velocity mode : 'n' (norme of the velocities) 'r' (radial velocities) """ if mode == 'r': v = (self.pos[:,0]*self.vel[:,0] + self.pos[:,1]*self.vel[:,1]) / sqrt(self.pos[:,0]**2+self.pos[:,1]**2) elif mode == 'n': v = sqrt(self.vel[:,0]**2 + self.vel[:,1]**2 + self.vel[:,2]**2) if vmax == None : vmax = mpi.mpi_max(v) if vmin == None : vmin = mpi.mpi_min(v) bins = arange(vmin,vmax,(vmax-vmin)/float(nb)) h = mpi.mpi_histogram(v,bins) return h,bins def zmodes(self,nr=32,nm=16,rm=32): """ Compute the vertical modes of a model nm = 16 : number of modes nr = 32 : number of radius rm = 50 : max radius return r : the radius used m : the modes computed m1 : the matrix of the amplitude m2 : the matrix of the phases """ ps = arange(-pi,pi+pi/(2.*nm),2*pi/(2.*nm))+pi # phases R = arange(0,nr+1,1)*float(rm)/nr # radius Rs = array([],float) r = self.rxy() m1 = array([],float) m2 = array([],float) # loop over all radius for i in range(len(R)-1): c = (r>=R[i])*(r<=R[i+1]) an = self.selectc(c) if sum(c.astype(int))<=1: #print "less than 1 particle in the coupe",R[i] amp = zeros(len(ps)/2).astype(float) m1 = concatenate((m1,amp)) m2 = concatenate((m2,amp)) continue x = an.pos[:,0] y = an.pos[:,1] z = an.pos[:,2] t = arctan2(y,x)+pi zm = [] ok = 0 for j in range(len(ps)-1): c = (t>=ps[j])*(t=R[i])*(r<=R[i+1]) an = self.selectc(c) if sum(c.astype(int))<=1: #print "less than 1 particle in the coupe",R[i] amp = zeros(len(ps)/2).astype(float) m1 = concatenate((m1,amp)) m2 = concatenate((m2,amp)) continue x = an.pos[:,0] y = an.pos[:,1] z = an.pos[:,2] t = arctan2(y,x)+pi dm = [] ok = 0 for j in range(len(ps)-1): c = (t>=ps[j])*(t [0,boxsize] centred : -> [-boxsize/2,boxsize/2] [x,y,z] : """ if boxsize==None: if self.has_var('boxsize'): boxsize = self.boxsize if boxsize != None: if mode == None: self.pos[:,0] = where((self.pos[:,0]<0.0 ),self.pos[:,0]+boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]<0.0 ),self.pos[:,1]+boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]<0.0 ),self.pos[:,2]+boxsize,self.pos[:,2]) self.pos[:,0] = where((self.pos[:,0]>boxsize),self.pos[:,0]-boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]>boxsize),self.pos[:,1]-boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]>boxsize),self.pos[:,2]-boxsize,self.pos[:,2]) elif mode=='centred': self.pos[:,0] = where((self.pos[:,0]<=-boxsize/2.),self.pos[:,0]+boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]<=-boxsize/2.),self.pos[:,1]+boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]<=-boxsize/2.),self.pos[:,2]+boxsize,self.pos[:,2]) self.pos[:,0] = where((self.pos[:,0]> boxsize/2.),self.pos[:,0]-boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]> boxsize/2.),self.pos[:,1]-boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]> boxsize/2.),self.pos[:,2]-boxsize,self.pos[:,2]) elif (type(mode)==ndarray) or (type(mode)==types.ListType): self.pos[:,0] = where((self.pos[:,0]<=mode[0]-boxsize/2.),self.pos[:,0]+boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]<=mode[1]-boxsize/2.),self.pos[:,1]+boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]<=mode[2]-boxsize/2.),self.pos[:,2]+boxsize,self.pos[:,2]) self.pos[:,0] = where((self.pos[:,0]> mode[0]+boxsize/2.),self.pos[:,0]-boxsize,self.pos[:,0]) self.pos[:,1] = where((self.pos[:,1]> mode[1]+boxsize/2.),self.pos[:,1]-boxsize,self.pos[:,1]) self.pos[:,2] = where((self.pos[:,2]> mode[2]+boxsize/2.),self.pos[:,2]-boxsize,self.pos[:,2]) def rotate_old(self,angle=0,mode='a',axis='x'): """ Rotate the positions and/or the velocities of the object around a specific axis. angle : rotation angle in radian axis : 'x' : around x 'y' : around y 'z' : around z : [x,y,z] : around this axis mode : 'p' : rotate only position 'v' : rotate only velocities 'a' : rotate both (default) """ if type(axis) == type('a'): # rotate around x,y or z if axis =='x': if mode=='p' or mode=='a': self.pos=rotx(angle,self.pos) if mode=='v' or mode=='a': self.vel=rotx(angle,self.vel) elif axis =='y': if mode=='p' or mode=='a': self.pos=roty(angle,self.pos) if mode=='v' or mode=='a': self.vel=roty(angle,self.vel) elif axis =='z': if mode=='p' or mode=='a': self.pos=rotz(angle,self.pos) if mode=='v' or mode=='a': self.vel=rotz(angle,self.vel) else: # rotate around a given axis # construction of the rotation matrix nxy = sqrt(axis[0]**2+axis[1]**2) theta_x = angle theta_z = 2.*pi - arctan2(axis[1],axis[0]) theta_y = arctan2(axis[2],nxy) if mode=='p' or mode=='a': # rot in z self.pos=rotz(theta_z,self.pos) # rot in y self.pos=roty(theta_y,self.pos) # rot in x self.pos=rotx(theta_x,self.pos) # rot in -y self.pos=roty(-theta_y,self.pos) # rot in -z self.pos=rotz(-theta_z,self.pos) if mode=='v' or mode=='a': # rot in z self.vel=rotz(theta_z,self.vel) # rot in y self.vel=roty(theta_y,self.vel) # rot in x self.vel=rotx(theta_x,self.vel) # rot in -y self.vel=roty(-theta_y,self.vel) # rot in -z self.vel=rotz(-theta_z,self.vel) def rotate(self,angle=0,axis=[1,0,0],point=[0,0,0],mode='a'): """ Rotate the positions and/or the velocities of the object around a specific axis defined by a vector and an point. angle : rotation angle in radian axis : direction of the axis point : center of the rotation mode : 'p' : rotate only position 'v' : rotate only velocities 'a' : rotate both (default) """ if axis=='x': axis = array([1,0,0],float) elif axis=='y': axis = array([0,1,0],float) elif axis=='z': axis = array([0,0,1],float) x = self.pos v = self.vel # center point x = x-point # construction of the rotation matrix norm = sqrt(axis[0]**2 + axis[1]**2 + axis[2]**2) if norm == 0: return x sn = sin(-angle/2.) e0 = cos(-angle/2.) e1 = axis[0]*sn/norm e2 = axis[1]*sn/norm e3 = axis[2]*sn/norm a = zeros((3,3),float) a[0,0] = e0**2 + e1**2 - e2**2 - e3**2 a[1,0] = 2.*(e1*e2 + e0*e3) a[2,0] = 2.*(e1*e3 - e0*e2) a[0,1] = 2.*(e1*e2 - e0*e3) a[1,1] = e0**2 - e1**2 + e2**2 - e3**2 a[2,1] = 2.*(e2*e3 + e0*e1) a[0,2] = 2.*(e1*e3 + e0*e2) a[1,2] = 2.*(e2*e3 - e0*e1) a[2,2] = e0**2 - e1**2 - e2**2 + e3**2 a = a.astype(float) # multiply x and v if mode=='p': x = dot(x,a) elif mode=='v': v = dot(v,a) else: x = dot(x,a) v = dot(v,a) # decenter point x = x+point self.pos = x.astype(float32) self.vel = v.astype(float32) def rotateR(self,R,mode='a'): """ Rotate the model using the matrix R mode : 'p' : only position 'v' : only velocities 'a' : both (default) """ # multiply x and v if mode=='p': self.pos = dot(self.pos,R) elif mode=='v': self.vel = dot(self.vel,R) else: self.pos = dot(self.pos,R) self.vel = dot(self.vel,R) def get_rotation_matrix_to_align_with_main_axis(self): """ Get the rotation matrix used to rotate the object in order to align it's main axis with the axis of its inertial tensor. """ # compute inertial tensor I = self.inertial_tensor() # find eigen values and vectors val,vec =linalg.eig(I) l1 = val[0] l2 = val[1] l3 = val[2] a1 = vec[:,0] a2 = vec[:,1] a3 = vec[:,2] # find Rm such that Rm*1,0,0 = a1 # that Rm*0,1,0 = a2 # that Rm*0,0,1 = a3 Rm = transpose(array([a1,a2,a3])) return Rm def align_with_main_axis(self,mode='a'): """ Rotate the object in order to align it's major axis with the axis of its inertial tensor. mode : 'p' : only position 'v' : only velocities 'a' : both (default) """ # find rotation matrix R = self.get_rotation_matrix_to_align_with_main_axis() # apply it self.rotateR(R,mode) def align(self,axis,mode='a',sgn='+',fact=None): """ Rotate the object in order to align the axis 'axis' with the z axis. axis : [x,y,z] mode : 'p' : only position 'v' : only velocities 'a' : both (default) sgn : '+' : normal rotation '-' : reverse sense of rotation fact : int : factor to increase the angle """ n = [axis[1], -axis[0],0.] theta = arccos(axis[2]/sqrt(axis[0]**2+axis[1]**2+axis[2]**2)) if sgn =='-': theta = -theta if fact != None: theta = theta*fact self.rotate(angle=theta,mode=mode,axis=n) def align2(self,axis1=[1,0,0],axis2=[0,0,1],point=[0,0,0]): ''' Rotate the object in order to align the axis 'axis' with the z axis. axis1 : [x,y,z] axis2 : [x,y,z] point : [x,y,z] ''' a1 = array(axis1,float) a2 = array(axis2,float) a3 = array([0,0,0],float) a3[0] = a1[1]*a2[2] - a1[2]*a2[1] a3[1] = a1[2]*a2[0] - a1[0]*a2[2] a3[2] = a1[0]*a2[1] - a1[1]*a2[0] n1 = sqrt(a1[0]**2 + a1[1]**2 + a1[2]**2) n2 = sqrt(a2[0]**2 + a2[1]**2 + a2[2]**2) angle = arccos(inner(a1,a2)/(n1*n2)) self.rotate(angle=angle,axis=a3,point=point) def spin(self,omega=None,L=None,j=None,E=None): """ Spin the object with angular velocity "omega" (rigid rotation). Omega is a 1 x 3 array object If L (total angular momentum) is explicitely given, compute Omega from L (1 x 3 array object). omega : angular speed (array vector) L : desired angular momentum j : desired energy fraction in rotation E : Total energy (without rotation) """ # do nothing if L==None and omega==None and j==None: pass # use j and E (spin around z axis) if j!=None: if E==None: "spin : print you must give E" else: if (j > 1): self.log.write("spin : j must be less than 1") sys.exit() Erot = j*E/(1-j) omega = sqrt(2*Erot/mpi_sum(self.mass*self.rxy()**2) ) omega = array([0,0,omega],float32) self.vel=spin(self.pos,self.vel,omega) # use omega elif L==None and omega!=None: omega = array(omega,float32) self.vel=spin(self.pos,self.vel,omega) # use L # Pfenniger 93 elif L!=None: L = array(L,float32) aamx = L[0] aamy = L[1] aamz = L[2] x = self.pos[:,0] y = self.pos[:,1] z = self.pos[:,2] vx = self.vel[:,0] vy = self.vel[:,1] vz = self.vel[:,2] m = self.mass Ixy = sum(m*x*y) Iyz = sum(m*y*z) Izx = sum(m*z*x) Ixx = sum(m*x*x) Iyy = sum(m*y*y) Izz = sum(m*z*z) Axx = Iyy+Izz Ayy = Izz+Ixx Azz = Ixx+Iyy Axy = -Ixy Ayz = -Iyz Azx = -Izx D = Axx*Ayy*Azz + 2*Axy*Ayz*Azx - Axx*Ayz**2 - Ayy*Azx**2 - Azz*Axy**2 DLX = sum(m*(y*vz-z*vy))-aamx DLY = sum(m*(z*vx-x*vz))-aamy DLZ = sum(m*(x*vy-y*vx))-aamz Bxx = Ayy*Azz - Ayz**2 Byy = Azz*Axx - Azx**2 Bzz = Axx*Ayy - Axy**2 Bxy = Azx*Ayz - Axy*Azz Byz = Axy*Azx - Ayz*Axx Bzx = Ayz*Axy - Azx*Ayy omega = array([0,0,0],float32) omega[0] = -(Bxx*DLX + Bxy*DLY + Bzx*DLZ)/D omega[1] = -(Bxy*DLX + Byy*DLY + Byz*DLZ)/D omega[2] = -(Bzx*DLX + Byz*DLY + Bzz*DLZ)/D self.vel=spin(self.pos,self.vel,omega) ################################# # # selection of particles # ################################# def selectc(self,c,local=False): """ Return an N-body object that contain only particles where the corresponding value in c is not zero. c is a nx1 Nbody array. c : the condition vector local : local selection (True) or global selection (False) """ new = Nbody(status='new',ftype=self.ftype[6:],local=local) # now, copy all var linked to the model for name in self.get_list_of_vars(): setattr(new, name, getattr(self,name)) # here, we create ptype on the fly (used to create new.npart) #self.ptype = array([],int) #for i in range(len(self.npart)): # self.ptype = concatenate( (self.ptype,ones(self.npart[i])*i) ) # now, copy and compress all array linked to the model for name in self.get_list_of_array(): vec = getattr(self,name) setattr(new, name, compress(c,vec,axis=0)) # now, compute new.npart #new.npart = array([],int) #for i in range(len(self.npart)): # c = (new.tpe==i) # npart_i = sum(c.astype(int)) # new.npart = concatenate( (new.npart, npart_i ) ) # check #if len(new.pos)!= sum(new.npart): # pass # other vars new.init() return new def selecti(self,i,local=False): """ Return an N-body object that contain only particles having their index (not id) in i. i : vector containing indexes local : local selection (True) or global selection (False) """ new = Nbody(status='new',ftype=self.ftype[6:],local=local) # now, copy all var linked to the model for name in self.get_list_of_vars(): setattr(new, name, getattr(self,name)) # here, we create ptype on the fly (used to create new.npart) #self.ptype = array([],int) #for i in range(len(self.npart)): # self.ptype = concatenate( (self.ptype,ones(self.npart[i])*i) ) # now, copy and compress all array linked to the model for name in self.get_list_of_array(): vec = getattr(self,name) setattr(new, name, vec[i]) # now, compute new.npart #new.npart = array([],int) #for i in range(len(self.npart)): # c = (new.tpe==i) # npart_i = sum(c.astype(int)) # new.npart = concatenate( (new.npart, npart_i ) ) # check #if len(new.pos)!= sum(new.npart): # pass # other vars new.init() return new def select(self,i=0): """ Return an N-body object that contain only particles of type i """ import types if type(i)== int: return self.selectc(self.tpe==i) elif type(i)==types.ListType: types = i for j in types: c = c * (j==self.tpe) return self.selectc(c) else: return self def sub(self,n1=0,n2=None): """ Return an N-body object that have particles whith indicies in the range [n1:n2]. n1 : number of the first particule n2 : number of the last particule Note : the first particle is 0 """ if n1 == None: n1 = 0 if n2 == None: n2 = self.nbody if n2 <= n1: n2 = n1+1 num = arange(self.nbody) return self.selectc((num>=n1)*(num<=n2)) def reduc(self,n,mass=False): """ Return an N-body object that contain a fraction 1/n of particles. n : inverse of the fraction of particule to be returned """ c = where(fmod(arange(self.nbody),n).astype(int)==0,1,0) nb = self.selectc(c) if mass: nb.mass = nb.mass * n return nb def selectp(self,lst=None,file=None,reject=False,local=False,from_num=True): """ Return an N-body object that contain only particles with specific number id. The list of id's is given either by lst (nx1 int array) or by the name ("file") of a file containing the list of id's. lst : vector list (integer) reject : True/False : if True, reject particles in lst (default = False) local : local selection (True) or global selection (False) frum_num : if True, use self.num to select particules if False, use arange(self.nbody) """ if lst != None: lst = array(lst,int) if file != None: lst = [] f = open(file) while 1: try: lst.append(int(string.split(f.readline())[0])) except: break f.close() lst = array(lst,int) # 1) sort the list ys = sort(lst) # 2) sort index in current file if from_num: xs = sort(self.num) zs = take(arange(self.nbody),argsort(self.num)) # sort 0,1,2,n following xs (or self.num) else: xs = arange(self.nbody) # 3) apply mask on sorted lists (here, getmask need xs and ys to be sorted) m = getmask(xs.astype(int),ys.astype(int)) if reject: m = logical_not(m) # 4) revert mask, following zs inverse transformation if from_num: c = take(m,argsort(zs)) else: c = m new = self.selectc(c,local=local) return new def getindex(self,num): """ Return an array of index of a particle from its specific number id. The array is empty if no particle corresponds to the specific number id. num : Id of the particle """ idx = compress((self.num == num),arange(self.nbody)) if len(idx)==1: return idx[0] else: return idx ################################# # # add particles # ################################# def append(self,solf,do_not_sort=False): """ Add to the current N-body object, particles form the N-body object "new". solf : Nbody object """ if solf.ftype != self.ftype: raise "append Error","files have different type" return if solf.get_list_of_array() != self.get_list_of_array(): raise "append Error","files have different arrays" return # loop over all types self_npart = self.npart solf_npart = solf.npart if len(self_npart) != len(self_npart): print "append : files have different mxnpart !" sys.exit() # add array linked to the model names = self.get_list_of_array() for name in names: vec1 = getattr(self,name) vec2 = getattr(solf,name) ''' vec = array([],float32) if vec1.ndim == 1: vec.shape = (0,) else: vec.shape = (0,3) # here, we guarantee the order of particles according to npart for i in arange(len(self_npart)): e11 = sum((arange(len(self_npart)) < i) * self_npart,0) e21 = sum((arange(len(solf_npart)) < i) * solf_npart,0) vec = concatenate((vec,vec1[e11:e11+self_npart[i]],vec2[e21:e21+solf_npart[i]])) ''' vec = concatenate((vec1,vec2)) setattr(self, name, vec) # here, we sort the particles, according to tpe if do_not_sort: pass else: sequence = self.tpe.argsort() for name in names: vec = getattr(self,name) vec = take(vec,sequence,axis=0) setattr(self, name, vec) self.nbody = self.nbody + solf.nbody self.npart = self.get_npart() # needed by self.get_num() self.npart_tot = self.get_npart_tot() # needed by self.get_num() self.num = self.get_num() self.init() def __add__(self,solf,do_not_sort=False): # first copy self new = deepcopy(self) # now, add solf new.append(solf,do_not_sort) return new ################################# # # sort particles # ################################# def sort(self): ''' sort particles according to their num variable ''' new = Nbody(status='new',ftype=self.ftype[6:]) sequence = argsort(self.num) # now, copy all var linked to the model for name in self.get_list_of_vars(): setattr(new, name, getattr(self,name)) # add array linked to the model for name in self.get_list_of_array(): setattr(new, name, take(getattr(self,name),sequence,axis=0)) new.num = new.get_num() new.init() return new def sort_type(self): ''' Contrary to sort, this fonction sort particles respecting their type. ''' new = Nbody(status='new',ftype=self.ftype[6:]) # now, copy all var linked to the model for name in self.get_list_of_vars(): setattr(new, name, getattr(self,name)) # add array linked to the model for name in self.get_list_of_array(): #vec = take(getattr(self,name),sequence,axis=0) vec = array([],float32) vec1 = getattr(self,name) if vec1.ndim == 1: vec.shape = (0,) else: vec.shape = (0,3) # loop over all types npart = self.npart for i in arange(len(npart)): e11 = sum((arange(len(npart)) < i) * npart) sequence = argsort(self.num[e11:e11+npart[i]]) vec = concatenate((vec,take(vec1[e11:e11+npart[i]],sequence,axis=0))) setattr(new, name, vec) new.num = new.get_num() new.init() return new ################################# # # Tree and SPH functions # ################################# def InitSphParameters(self,DesNumNgb=33,MaxNumNgbDeviation=3): self.DesNumNgb = DesNumNgb self.MaxNumNgbDeviation = MaxNumNgbDeviation self.Density = None self.Hsml = None if not self.has_var('Tree'): self.Tree = None def setTreeParameters(self,Tree,DesNumNgb,MaxNumNgbDeviation): if Tree==None: self.Tree = Tree = self.getTree() if DesNumNgb==None: DesNumNgb = self.DesNumNgb else: self.DesNumNgb = DesNumNgb if MaxNumNgbDeviation==None: MaxNumNgbDeviation = self.MaxNumNgbDeviation else: self.MaxNumNgbDeviation = MaxNumNgbDeviation return Tree,DesNumNgb,MaxNumNgbDeviation def getTree(self,force_computation=False,ErrTolTheta=0.8): ''' Return a Tree object ''' if self.Tree!=None and force_computation==False: return self.Tree else: print "create the tree : ErrTolTheta=",ErrTolTheta # decide if we use tree or ptree npart = array(self.npart) if mpi.mpi_NTask()>1: print "%d : use ptree"%(mpi.mpi_ThisTask()) self.Tree = ptreelib.Tree(npart=npart,pos=self.pos,vel=self.vel,mass=self.mass,num=self.num,tpe=self.tpe) else: self.Tree = treelib.Tree(npart=npart,pos=self.pos,vel=self.vel,mass=self.mass,ErrTolTheta=ErrTolTheta) return self.Tree def get_rsp_approximation(self,DesNumNgb=None,MaxNumNgbDeviation=None,Tree=None): ''' Return an aproximation of rsp, based on the tree. ''' Tree,DesNumNgb,MaxNumNgbDeviation = self.setTreeParameters(Tree,DesNumNgb,MaxNumNgbDeviation) return Tree.InitHsml(DesNumNgb,MaxNumNgbDeviation) def ComputeSph(self,DesNumNgb=None,MaxNumNgbDeviation=None,Tree=None): ''' Compute self.Density and self.Hsml using sph approximation ''' Tree,DesNumNgb,MaxNumNgbDeviation = self.setTreeParameters(Tree,DesNumNgb,MaxNumNgbDeviation) if self.Hsml==None: if not self.has_array('rsp'): self.Hsml = self.get_rsp_approximation(DesNumNgb,MaxNumNgbDeviation,Tree) else: self.Hsml=self.rsp self.Density,self.Hsml = Tree.Density(self.pos,self.Hsml,DesNumNgb,MaxNumNgbDeviation) def ComputeDensityAndHsml(self,pos=None,Hsml=None,DesNumNgb=None,MaxNumNgbDeviation=None,Tree=None): ''' Compute Density and Hsml (for a specific place) ''' Tree,DesNumNgb,MaxNumNgbDeviation = self.setTreeParameters(Tree,DesNumNgb,MaxNumNgbDeviation) if pos==None: pos = self.pos if Hsml==None: Hsml = ones(len(pos)).astype(float32) Density,Hsml = Tree.Density(pos,Hsml,DesNumNgb,MaxNumNgbDeviation) return Density,Hsml def SphEvaluate(self,val,pos=None,vel=None,hsml=None,DesNumNgb=None,MaxNumNgbDeviation=None,Tree=None): ''' Return an sph evaluation of the variable var ''' Tree,DesNumNgb,MaxNumNgbDeviation = self.setTreeParameters(Tree,DesNumNgb,MaxNumNgbDeviation) if pos == None: pos = self.pos if vel == None: vel = self.vel if hsml == None: if self.Hsml==None: if not self.has_array('rsp'): self.Hsml = self.get_rsp_approximation(DesNumNgb,MaxNumNgbDeviation,Tree) else: self.Hsml=self.rsp hsml = self.Hsml if self.Density==None: if not self.has_array('rho'): self.Density = self.SphDensity(DesNumNgb,MaxNumNgbDeviation,Tree) else: self.Density=self.rho if type(val) == ndarray: val = Tree.SphEvaluate(pos,hsml,self.Density,val,DesNumNgb,MaxNumNgbDeviation) else: if val =='div': val = Tree.SphEvaluateDiv(pos,vel,hsml,self.Density,DesNumNgb,MaxNumNgbDeviation) elif val =='rot': val = Tree.SphEvaluateRot(pos,vel,hsml,self.Density,DesNumNgb,MaxNumNgbDeviation) elif val =='ngb': val = Tree.SphEvaluateNgb(pos,hsml,DesNumNgb,MaxNumNgbDeviation) return val ################################# # # sph functions # ################################# def weighted_numngb(self,num): ''' num = particle where to compute weighted_numngb see Springel 05 ''' def wk1(hinv3,u): KERNEL_COEFF_1=2.546479089470 KERNEL_COEFF_2=15.278874536822 wk = hinv3 * (KERNEL_COEFF_1 + KERNEL_COEFF_2 * (u - 1) * u * u) return wk def wk2(hinv3,u): KERNEL_COEFF_5=5.092958178941 wk = hinv3 * KERNEL_COEFF_5 * (1.0 - u) * (1.0 - u) * (1.0 - u) return wk def getwk(r,h): # we do not exclude the particle itself u = r/h hinv3 = 1./h**3 wk = where((u<0.5),wk1(hinv3,u),wk2(hinv3,u)) wk = where((r1: list_of_array = self.get_list_of_array() # loop over all particles type npart = self.npart new_npart = npart for i in range(len(npart)): #if i==0: nparts = mpi.mpi_allgather(npart[i]) nparts = array(nparts) if mpi.mpi_IsMaster(): ex_table = mpi.mpi_GetExchangeTable(nparts) ex_table = mpi.mpi_bcast(ex_table,0) else: ex_table = None ex_table = mpi.mpi_bcast(ex_table,0) # send particles for toTask in range(mpi.NTask): if ex_table[mpi.ThisTask,toTask] > 0: n_elt = ex_table[mpi.ThisTask,toTask] #print "%d send %d to %d"%(mpi.ThisTask,n_elt,toTask) # first_elt = first elt of the current block first_elt = sum((arange(len(new_npart)) < i) * new_npart) # update npart new_npart[i] = new_npart[i] - n_elt # loop over all vect erd,mass,num,pos,rho,rsp,u,vel for name in list_of_array: vec = getattr(self,name) sub_vec = vec[first_elt:first_elt+n_elt] if len(sub_vec) != n_elt: print "redistribute error :" print "node %d should send len=%d got len=%d"%(mpi.ThisTask,n_elt,len(sub_vec)) sys.exit() mpi.mpi_send(name,toTask) mpi.mpi_send(sub_vec,toTask) #self.pos = concatenate( (self.pos[:first_elt],self.pos[first_elt+n_elt:]) ) setattr(self, name, concatenate( (vec[:first_elt],vec[first_elt+n_elt:]) ) ) # recieve particles for fromTask in range(mpi.NTask): if ex_table[fromTask,mpi.ThisTask] > 0: n_elt = ex_table[fromTask,mpi.ThisTask] #print "%d get %d from %d"%(mpi.ThisTask,n_elt,fromTask) # first_elt = first elt of the current block first_elt = sum((arange(len(new_npart)) < i) * new_npart) # update npart new_npart[i] = new_npart[i] + n_elt # loop over all vect for name in list_of_array: # first, check name send_name = mpi.mpi_recv(fromTask) if send_name != name: raise "Task %d FromTask %d, %s != %s"%(mpi.mpi_ThisTask(),fromTask,send_name,name) vec = getattr(self,name) sub_vec = mpi.mpi_recv(fromTask) if len(sub_vec) != n_elt: print "redistribute error :" print "node %d should recive len=%d got len=%d"%(mpi.ThisTask,n_elt,len(sub_vec)) sys.exit() #self.pos = concatenate( (vec[:first_elt],sub_vec,vec[first_elt:]) ) setattr(self, name, concatenate( (vec[:first_elt],sub_vec,vec[first_elt:]) ) ) self.init() def ExchangeParticles(self): ''' Exchange particles betwee procs, using peano-hilbert decomposition computed in ptree ''' if self.Tree==None: self.Tree = self.getTree() # get num and procs from the Tree num,procs = self.Tree.GetExchanges() # compute the transition table T H,bins = histogram(procs,arange(mpi.mpi_NTask())) T = mpi.mpi_AllgatherAndConcatArray(H) T.shape = (mpi.mpi_NTask(),mpi.mpi_NTask()) # loop over all numpy vectors list_of_array = self.get_list_of_array() # loop over all vect for name in list_of_array: if name != "num": setattr(self, name, mpi.mpi_ExchangeFromTable(T,procs,num,getattr(self,name),copy(self.num)) ) # do num at the end self.num = mpi.mpi_ExchangeFromTable(T,procs,num,self.num,copy(self.num)) self.init() def SendAllToAll(self): ''' Send all particles to all nodes at the end of the day, all nodes have the same nbody object ''' nbs = [] for i in xrange(mpi.NTask-1): prev = (mpi.ThisTask-i-1)%mpi.NTask next = (mpi.ThisTask+i+1)%mpi.NTask nbs.append(mpi.mpi_sendrecv(self,dest=next,source=prev)) for nbi in nbs: self = self + nbi return self ################################# # # specific parallel functions # ################################# def gather_pos(self): ''' Gather in a unique array all positions of all nodes. ''' return self.gather_vec(self.pos) def gather_vel(self): ''' Gather in a unique array all velocites of all nodes. ''' return self.gather_vec(self.vel) def gather_mass(self): ''' Gather in a unique array all mass of all nodes. ''' return self.gather_vec(self.mass) def gather_num(self): ''' Gather in a unique array all num of all nodes. ''' return self.gather_vec(self.num) def gather_vec(self,vec): ''' Gather in a unique array all vectors vec of all nodes. ''' # here, we assume that we have a vector npart # giving the number of particles per type vec_all = array([],vec.dtype) if vec.ndim==1: vec_all.shape = (0,) else: vec_all.shape = (0,vec.shape[1]) i1 = 0 npart = self.npart for i in range(len(npart)): i2 = i1 + npart[i] if (i1!=i2): vec_all = concatenate((vec_all,mpi.mpi_AllgatherAndConcatArray(vec[i1:i2]))) i1 = i1 + npart[i] return vec_all ################################# # # graphical operations # ################################# def display(self,*arg,**kw): ''' Display the model ''' if kw.has_key('palette'): palette = kw['palette'] else: palette = None if kw.has_key('save'): save = kw['save'] else: save = None if kw.has_key('marker'): marker = kw['marker'] else: marker = None params = extract_parameters(arg,kw,self.defaultparameters) mat,matint,mn_opts,mx_opts,cd_opts = self.Map(params) if mpi.mpi_IsMaster(): if save != None: if os.path.splitext(save)[1] == ".fits": io.WriteFits(transpose(mat).astype(float32),save, None) return if palette!=None: mplot(matint,palette=palette,save=save,marker=marker) else: mplot(matint,save=save,marker=marker) def show(self,*arg,**kw): ''' Display the model this is an alias to display ''' self.display(*arg,**kw) def Map(self,*arg,**kw): ''' Return 2 final images (float and int) ''' params = extract_parameters(arg,kw,self.defaultparameters) mn_opts = [] mx_opts = [] cd_opts = [] if self.nbody==0 and mpi.mpi_NTask()==1: mat = zeros(params['shape'],float32) matint = mat.astype(int) mn_opts.append(params['mn']) mx_opts.append(params['mx']) cd_opts.append(params['cd']) return mat,matint,mn_opts,mx_opts,cd_opts # compute map mat = self.CombiMap(params) # set ranges matint,mn_opt,mx_opt,cd_opt = set_ranges(mat,scale=params['scale'],cd=params['cd'],mn=params['mn'],mx=params['mx']) mn_opts.append(mn_opt) mx_opts.append(mx_opt) cd_opts.append(cd_opt) # add contour if params['l_color'] != 0: matint = contours(mat,matint,params['l_n'],params['l_min'],params['l_max'],params['l_kx'],params['l_ky'],params['l_color'],params['l_crush']) # add box and ticks if params['b_weight'] != 0: matint = add_box(matint,shape=params['shape'],size=params['size'],center=None,box_opts=(params['b_weight'],params['b_xopts'],params['b_yopts'],params['b_color'])) return mat,matint,mn_opts,mx_opts,cd_opts def CombiMap(self,*arg,**kw): ''' Return an image in form of a matrix (nx x ny float array). Contrary to ComputeMap, CombiMap compose different output of ComputeMap. pos : position of particles (moment 0) sr : dispertion in r (with respect to xp) svr : dispertion in vr vxyr : mean velocity in the plane svxyr: dispertion in vxy vtr : mean tangential velocity in the plane svtr : dispertion in vt szr : ratio sigma z/sigma r ''' params = extract_parameters(arg,kw,self.defaultparameters) mode = params['mode'] #if mode == 'pos': # mat = self.ComputeMap(params) if mode == 'm': mat = self.ComputeMap(params) elif mode == 'sr': mat = self.ComputeSigmaMap(params,mode1='r',mode2='r2') elif mode == 'svr': mat = self.ComputeSigmaMap(params,mode1='vr',mode2='vr2') elif mode == 'svxyr': mat = self.ComputeSigmaMap(params,mode1='vxyr',mode2='vxyr2') elif mode == 'svtr': mat = self.ComputeSigmaMap(params,mode1='vtr',mode2='vtr2') elif mode == 'szr': # could be simplified m0 = self.ComputeMap(params,mode='m') m1 = self.ComputeMap(params,mode='vr') m2 = self.ComputeMap(params,mode='vr2') m1 = where(m0==0,0,m1) m2 = where(m0==0,0,m2) m0 = where(m0==0,1,m0) mat = m2/m0 - (m1/m0)**2 mat_sz = sqrt(numclip(mat,0,1e10)) m0 = self.ComputeMap(params,mode='m') m1 = self.ComputeMap(params,mode='vxyr') m2 = self.ComputeMap(params,mode='vxyr2') m1 = where(m0==0,0,m1) m2 = where(m0==0,0,m2) m0 = where(m0==0,1,m0) mat = m2/m0 - (m1/m0)**2 mat_sr = sqrt(numclip(mat,0,1e10)) mat_sz = where(mat_sr==0,0,mat_sz) mat_sr = where(mat_sr==0,1,mat_sr) mat = mat_sz/mat_sr elif mode == 'lum': mat = self.ComputeMap(params,mode='lum') else: mat = self.ComputeMeanMap(params,mode1=mode) return mat def ComputeMeanMap(self,*arg,**kw): """ Compute the mean map of an observable. """ params = extract_parameters(arg,kw,self.defaultparameters) if kw.has_key('mode1'): mode1 = kw['mode1'] else: raise "ComputeMeanMap :","you must give parameter mode1" m0 = self.ComputeMap(params,mode='0') m1 = self.ComputeMap(params,mode=mode1) m1 = where(m0==0,0,m1) m0 = where(m0==0,1,m0) mat = m1/m0 return mat def ComputeSigmaMap(self,*arg,**kw): """ Compute the sigma map of an observable. """ params = extract_parameters(arg,kw,self.defaultparameters) if kw.has_key('mode1'): mode1 = kw['mode1'] else: raise "ComputeMeanMap","you must give parameter mode1" if kw.has_key('mode2'): mode2 = kw['mode2'] else: raise "ComputeMeanMap","you must give parameter mode2" m0 = self.ComputeMap(params,mode='0') m1 = self.ComputeMap(params,mode=mode1) m2 = self.ComputeMap(params,mode=mode2) m1 = where(m0==0,0,m1) m2 = where(m0==0,0,m2) m0 = where(m0==0,1,m0) mat = m2/m0 - (m1/m0)**2 mat = sqrt(numclip(mat,0,1e10)) return mat def ComputeMap(self,*arg,**kw): ''' Return an image in form of a matrix (nx x ny float array) obs : position of observer x0 : eye position xp : focal position alpha : angle of the head view : 'xy' 'xz' 'yz' eye : 'right' 'left' dist_eye : distance between eyes mode : mode of map space : pos or vel persp : 'on' 'off' clip : (near,far) size : (maxx,maxy) cut : 'yes' 'no' frsp : factor for rsp shape : shape of the map ''' params = extract_parameters(arg,kw,self.defaultparameters) obs = params['obs'] x0 = params['x0'] xp = params['xp'] alpha = params['alpha'] mode = params['mode'] view = params['view'] r_obs = params['r_obs'] eye = params['eye'] dist_eye = params['dist_eye'] foc = params['foc'] space = params['space'] persp = params['persp'] clip = params['clip'] size = params['size'] shape = params['shape'] cut = params['cut'] frsp = params['frsp'] filter_name = params['filter_name'] filter_opts = params['filter_opts'] # 0) if getvaltype(mode)=='normal': val = getval(self,mode=mode,obs=obs) # 1) get observer position if obs==None: obs = geo.get_obs(x0=x0,xp=xp,alpha=alpha,view=view,r_obs=r_obs) # 2) expose the model # !!! as in self.expose we use Nbody() this must be called by each Task nb,obs = self.expose(obs,eye,dist_eye,foc=foc,space=space) if self.nbody > 0: # 3) compute val if getvaltype(mode)=='in projection': val = getval(nb,mode=mode,obs=obs) # 4) projection transformation if persp == 'on': zp = - nb.pos[:,2] # save dist obs-point pos = geo.frustum(nb.pos,clip,size) else: pos = geo.ortho(nb.pos,clip,size) # 5) keep only particles in 1:1:1 if not self.has_array('rsp'): # bad !!! self.rsp = None if cut=='yes': if self.rsp!= None: if params['rendering']=='map': pos,(mass,rsp,val,zp) = geo.boxcut(pos,[self.mass,self.rsp,val,zp]) else: pos,(mass,rsp,val,zp) = geo.boxcut_segments(pos,[self.mass,self.rsp,val,zp]) else: if params['rendering']=='map': pos,(mass,val,zp) = geo.boxcut(pos,[self.mass,val,zp]) else: pos,(mass,val,zp) = geo.boxcut_segments(pos,[self.mass,val,zp]) rsp = None else: mass = self.mass rsp = self.rsp if len(pos)!=0: # 6) scale rsp and scale mass if frsp!= 0: if (rsp==None) or (sum(rsp)==0): rsp = ones(len(pos),float32) if persp == 'on': fact = 1/(zp+clip[0]) # rsp is distance dependant... rsp = rsp * fact rsp = rsp.astype(float32) # mass is distance dependant... mass = mass*fact**2 mass = mass.astype(float32) rsp = rsp*frsp # multiply with the factor self.log.write( "rsp : min = %10.5f max = %10.5f mean = %10.5f"%(min(rsp),max(rsp),rsp.mean()) ) rsp = numclip(rsp,0,100) self.log.write( "rsp : min = %10.5f max = %10.5f mean = %10.5f"%(min(rsp),max(rsp),rsp.mean()) ) rsp = rsp.astype(float32) else: rsp = None # 7) viewport transformation : (x,y) -> ((0,1),(0,1)) pos = geo.viewport(pos,shape=None) pos = pos.astype(float32) # 8) render : map or lines if params['rendering']=='map': if rsp != None: mat = mkmap2dsph(pos,mass,val,rsp,shape) else: mat = mkmap2d(pos,mass,val,shape) elif params['rendering']=='polygon': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_polygon(mat,pos[:,0],pos[:,1],1) elif params['rendering']=='lines': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_lines(mat,pos[:,0],pos[:,1],1) elif params['rendering']=='segments': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_segments(mat,pos[:,0],pos[:,1],1,zp) elif params['rendering']=='points': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_points(mat,pos[:,0],pos[:,1],1) elif params['rendering']=='polygon2': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_polygonN(mat,pos[:,0],pos[:,1],1,2) elif params['rendering']=='polygon4': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_polygonN(mat,pos[:,0],pos[:,1],1,4) elif params['rendering']=='polygon10': pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_polygonN(mat,pos[:,0],pos[:,1],1,10) elif params['rendering'][:8]=='polygon#': n = int(params['rendering'][8:]) pos = pos * array([params['shape'][0],params['shape'][1] ,0]) mat = zeros(params['shape'],float32) mat = draw_polygonN(mat,pos[:,0],pos[:,1],1,n) else: # compute a map if rsp != None: mat = mkmap2dsph(pos,mass,val,rsp,shape) else: mat = mkmap2d(pos,mass,val,shape) # there is no particles (after 5) else: mat = zeros(params['shape'],float32) # there is no particles else: mat = zeros(params['shape'],float32) # 9) sum mat over all proc mat = mpi.mpi_allreduce(mat) # could be more efficient if only the master get the final mat # 10) filter matrix if mpi.mpi_IsMaster(): if params['filter_name'] != None: mat = apply_filter(mat,name=filter_name,opt=filter_opts) return mat def ComputeObjectMap(self,*arg,**kw): ''' * * * IN DEVELOPPEMENT : allow to draw an object like a box, a grid... * * * Return an image in form of a matrix (nx x ny float array) obs : position of observer x0 : eye position xp : focal position alpha : angle of the head view : 'xy' 'xz' 'yz' eye : 'right' 'left' dist_eye : distance between eyes mode : mode of map space : pos or vel persp : 'on' 'off' clip : (near,far) size : (maxx,maxy) cut : 'yes' 'no' frsp : factor for rsp shape : shape of the map ''' # here, nb must represent a geometric object # ob # expose : -> must use ob instead of self # --> we can give explicitely pos and vel # ensuite, le nb est le bon params = extract_parameters(arg,kw,self.defaultparameters) obs = params['obs'] x0 = params['x0'] xp = params['xp'] alpha = params['alpha'] mode = params['mode'] view = params['view'] r_obs = params['r_obs'] eye = params['eye'] dist_eye = params['dist_eye'] foc = params['foc'] space = params['space'] persp = params['persp'] clip = params['clip'] size = params['size'] shape = params['shape'] cut = params['cut'] frsp = params['frsp'] filter_name = params['filter_name'] filter_opts = params['filter_opts'] # 0) if getvaltype(mode)=='normal': val = getval(self,mode=mode,obs=obs) # 1) get observer position if obs==None: obs = geo.get_obs(x0=x0,xp=xp,alpha=alpha,view=view,r_obs=r_obs) # 2) expose the model # !!! as in self.expose we use Nbody() this must be called by each Task nb,obs = self.expose(obs,eye,dist_eye,foc=foc,space=space) if self.nbody > 0: # 3) compute val if getvaltype(mode)=='in projection': val = getval(nb,mode=mode,obs=obs) # 4) projection transformation if persp == 'on': zp = - nb.pos[:,2] # save dist obs-point pos = geo.frustum(nb.pos,clip,size) else: pos = geo.ortho(nb.pos,clip,size) # 5) keep only particles in 1:1:1 if not self.has_array('rsp'): # bad !!! self.rsp = None if cut=='yes': if self.rsp!= None: pos,(mass,rsp,val,zp) = geo.boxcut(pos,[self.mass,self.rsp,val,zp]) else: pos,(mass,val,zp) = geo.boxcut(pos,[self.mass,val,zp]) rsp = None else: mass = self.mass rsp = self.rsp if len(pos)!=0: # 6) scale rsp and scale mass if frsp!= 0: if (rsp==None) or (sum(rsp)==0): rsp = ones(len(pos),float32) if persp == 'on': fact = 1/((zp-clip[0])+ 2*clip[0]) # rsp is distance dependant... rsp = rsp * fact rsp = rsp.astype(float32) # mass is distance dependant... mass = mass*fact**2 mass = mass.astype(float32) rsp = rsp*frsp # multiply with the factor self.log.write( "rsp : min = %10.5f max = %10.5f mean = %10.5f"%(min(rsp),max(rsp),rsp.mean()) ) rsp = numclip(rsp,0,100) self.log.write( "rsp : min = %10.5f max = %10.5f mean = %10.5f"%(min(rsp),max(rsp),rsp.mean()) ) rsp = rsp.astype(float32) else: rsp = None # 7) viewport transformation : (x,y) -> ((0,1),(0,1)) pos = geo.viewport(pos,shape=None) pos = pos.astype(float32) # 8) get the map #if rsp != None: # mat = mkmap2dsph(pos,mass,val,rsp,shape) #else: # mat = mkmap2d(pos,mass,val,shape) # # empty matrix mat = zeros(params['shape'],float32) #for po in pos: # i = int(po[0]*params['shape'][0]) # j = int(po[1]*params['shape'][1]) # mat[i,j]=255 pos = pos * [ params['shape'][0],params['shape'][1],1 ] x0 = pos[0][0] y0 = pos[0][1] x1 = pos[1][0] y1 = pos[1][1] mat = libutil.draw_line(mat,x0,x1,y0,y1,255) #mat = libutil.draw_cube(mat,pos,255) # there is no particles (after 5) else: mat = zeros(params['shape'],float32) # there is no particles else: mat = zeros(params['shape'],float32) # 9) sum mat over all proc mat = mpi.mpi_allreduce(mat) # may be inefficient, better use reduce ? # 10) filter matrix if mpi.mpi_IsMaster(): if params['filter_name'] != None: mat = apply_filter(mat,name=filter_name,opt=filter_opts) return mat def expose(self,obs,eye=None,dist_eye=None,foc=None,space='pos',pos=None,vel=None): """ Rotate and translate the object in order to be seen as if the observer was in x0, looking at a point in xp. obs : observer matrix eye : 'right' or 'left' dist_eye : distance between eyes (separation = angle) space : pos or vel foc : focal """ # create a minimal copy of self if pos!=None and vel!=None: obj = Nbody(status='new',p_name='none',pos=pos,vel=vel,mass=None,ftype='default') else: obj = Nbody(status='new',p_name='none',pos=self.pos,vel=self.vel,mass=None,ftype='default') if space=='vel': obj.pos = self.vel # first : put x0 at the origin obj.translate(- obs[0]) obs = obs - obs[0] # second : anti-align e1 with z obj.align2( axis1=obs[1],axis2=[0,0,-1]) obs = geo.align(obs,axis1=obs[1],axis2=[0,0,-1]) # third : align e3 with y obj.align2( axis1=obs[3],axis2=[0,1,0]) obs = geo.align(obs,axis1=obs[3],axis2=[0,1,0]) # fourth if eye is defined if eye=='right': if foc==None or foc==0: # simple translation (wee look at infini) obj.translate([-dist_eye,0,0]) # not /2 for compatibility with glups else: Robs = foc phi = -arctan(dist_eye) obj.rotate( angle=-phi,axis=[0,1,0],point=[0,0,-Robs]) elif eye=='left': if foc==None or foc==0: # simple translation (wee look at infini) obj.translate([+dist_eye,0,0]) # not /2 for compatibility with glups else: Robs = foc phi = -arctan(dist_eye) obj.rotate( angle=+phi,axis=[0,1,0],point=[0,0,-Robs]) return obj,obs ''' def getvxy(self,shape=(256,256),size=(30.,30.),center=(0.,0.,0.),view='xz',vn=8.,vmax=0.1,color=1): self.log.write( "the result may be incertain (in development)" ) # choice of the view if view=='xz': view=1 elif view=='xy': view=2 elif view=='yz': view=3 elif view!='xz'and view!='xy'and view!='yz': view=1 dx = mapone(self.pos,self.mass,self.vel[:,0],shape,size,center,view) * vn/vmax dy = - mapone(self.pos,self.mass,self.vel[:,2],shape,size,center,view) * vn/vmax # mask mask = fromfunction(lambda x,y: (fmod(x,vn) + fmod(y,vn))==0 ,shape) # points de depart x0 = indices(shape)[0] + int(vn/2.) y0 = indices(shape)[1] + int(vn/2.) # points d'arrivee x1 = x0 + dx.astype(int) y1 = y0 + dy.astype(int) # truncation x1 = numclip(x1,0,shape[0]) y1 = numclip(y1,0,shape[1]) # compress mask = mask*(x1!=x0)*(y1!=y0) mask = ravel(mask) x0 = compress(mask,ravel(x0)) x1 = compress(mask,ravel(x1)) y0 = compress(mask,ravel(y0)) y1 = compress(mask,ravel(y1)) # trace lines mat = zeros(shape,float32) color = array(color,int8)[0] for i in range(len(x0)): create_line(mat,x0[i],y0[i],x1[i],y1[i],color) create_line(mat,x0[i],y0[i],x0[i]+1,y0[i]+1,color) create_line(mat,x0[i],y0[i],x0[i]+1,y0[i] ,color) create_line(mat,x0[i],y0[i],x0[i] ,y0[i]+1,color) return mat.astype(int8) ''' ################################# # # 1d histograms routines # ################################# ########################### def Histo(self,bins,mode='m',space='R'): ########################### histo = self.CombiHisto(bins,mode=mode,space=space) # take the mean bins1 = bins[:-1] bins2 = bins[1:] bins = (bins1+bins2)/2. return bins,mpi.mpi_allreduce(histo) ########################### def CombiHisto(self,bins,mode='m',space='R'): ########################### if mode == 'm': histo = self.ComputeHisto(bins,mode='0',space=space) elif mode == 'sz': histo = self.ComputeSigmaHisto(bins,mode1='z',mode2='z2',space=space) elif mode == 'svz': histo = self.ComputeSigmaHisto(bins,mode1='vz',mode2='vz2',space=space) elif mode == 'svt': histo = self.ComputeSigmaHisto(bins,mode1='vt',mode2='vt2',space=space) elif mode == 'svr': histo = self.ComputeSigmaHisto(bins,mode1='vr',mode2='vr2',space=space) elif mode == 'vt': histo = self.ComputeMeanHisto(bins,mode1='vt',space=space) elif mode == 'vr': histo = self.ComputeMeanHisto(bins,mode1='vr',space=space) elif mode == 'vz': histo = self.ComputeMeanHisto(bins,mode1='vz',space=space) else: print "unknown mode %s"%(mode) return histo ################################# def ComputeMeanHisto(self,bins,mode1,space): ################################# """ Compute the mean map of an observable. """ h0 = self.ComputeHisto(bins,mode='0',space=space) h1 = self.ComputeHisto(bins,mode=mode1,space=space) h1 = where(h0==0,0,h1) h0 = where(h0==0,1,h0) h = h1/h0 return h ################################# def ComputeSigmaHisto(self,bins,mode1,mode2,space): ################################# """ Compute the histogram of an observable. """ h0 = self.ComputeHisto(bins,mode='0',space=space) h1 = self.ComputeHisto(bins,mode=mode1,space=space) h2 = self.ComputeHisto(bins,mode=mode2,space=space) h1 = where(h0==0,0,h1) h2 = where(h0==0,0,h2) h0 = where(h0==0,1,h0) h = h2/h0 - (h1/h0)**2 h = sqrt(numclip(h,0,1e10)) return h ################################# def ComputeHisto(self,bins,mode,space): ################################# ''' Compute and histogram ''' # set space if space == 'R': x = self.rxy() elif space == 'r': x = self.rxyz() # set mode if mode == 'm' or mode=='0': v = self.mass elif mode == 'z': v = self.mass*self.z() elif mode == 'z2': v = self.mass*self.z()**2 elif mode == 'vz': v = self.mass*self.vz() elif mode == 'vz2': v = self.mass*self.vz()**2 elif mode == 'vt': v = self.mass*self.Vt() elif mode == 'vt2': v = self.mass*self.Vt()**2 elif mode == 'vr': v = self.mass*self.Vr() elif mode == 'vr2': v = self.mass*self.Vr()**2 else: print "unknown mode %s"%(mode) histo = whistogram(x.astype(float),v.astype(float),bins.astype(float)) return histo ############################################ # # Routines to get velocities from positions # ############################################ def Get_Velocities_From_Virial_Approximation(self,select=None,vf=1.,eps=0.1,UseTree=True,Tree=None,ErrTolTheta=0.5): ''' This routine does not work ! Do not use it, or check ! ''' if select!=None: nb_sph = self.select(select) else: nb_sph = self # build the Tree for nb self.getTree(force_computation=True,ErrTolTheta=ErrTolTheta) # compute potential pot = 0.5*self.TreePot(nb_sph.pos,eps) # virial approximation to get the velocities sigmasp = sqrt(-pot/3*vf) # compute accel acc = self.TreeAccel(nb_sph.pos,eps) pot = (acc[:,0]*nb_sph.pos[:,0] + acc[:,1]*nb_sph.pos[:,1] + acc[:,2]*nb_sph.pos[:,2]) # virial approximation to get the velocities sigmas = sqrt(-pot/3*vf) # avoid negative values sigmas = where( (pot>0),sigmasp, sigmas ) # generate velocities vx = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vy = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vz = sigmas*RandomArray.standard_normal([nb_sph.nbody]) nb_sph.vel = transpose(array([vx,vy,vz])).astype(float32) return nb_sph def Get_Velocities_From_AdaptativeSpherical_Grid(self,select=None,eps=0.1,n=1000,UseTree=True,Tree=None,phi=None,ErrTolTheta=0.5): ''' Computes velocities using the jeans equation in spherical coordinates. An adaptative grid is set automatically. ''' if select!=None: nb_sph = self.select(select) else: nb_sph = self # create the adaptiative grid and compute rho r = nb_sph.rxyz() a = r.argsort() x = take(nb_sph.pos[:,0],a) y = take(nb_sph.pos[:,1],a) z = take(nb_sph.pos[:,2],a) mass = take(nb_sph.mass,a) r = sqrt( x**2 + y**2 + z**2 ) n_bins = int((nb_sph.nbody+1)/n + 1) rs = [] rsmin = [] rsmax = [] rhos = [] ns = [] for i in xrange(n_bins): jmin = i*n jmax = i*n + n jmin = min(jmin,nb_sph.nbody-1) jmax = min(jmax,nb_sph.nbody-1) if jmin!=jmax: rr = r[jmin:jmax] mm = mass[jmin:jmax] rmean = rr.mean() rmin = rr.min() rmax = rr.max() rs.append(rmean) rsmin.append(rmin) rsmax.append(rmax) # density rho = sum(mm)/(4/3.*pi*(rmax**3-rmin**3)) rhos.append(rho) # number ns.append(len(rr)) r = array(rs) rsmin = array(rsmin) rsmax = array(rsmax) rho = array(rhos) dr = rsmax-rsmin nn = array(ns) # build the Tree for nb self.getTree(force_computation=True,ErrTolTheta=ErrTolTheta) # compute potential x = r y = zeros(len(r)) z = zeros(len(r)) pos = transpose(array([x,y,z])).astype(float32) phi = self.TreePot(pos,eps) # compute sigma sigma = libdisk.get_1d_Sigma_From_Rho_Phi(rho=rho,phi=phi,r=r,dr=dr) # generate velocities for all particles sigmas = lininterp1d(nb_sph.rxyz().astype(float32),r.astype(float32),sigma.astype(float32)) vx = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vy = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vz = sigmas*RandomArray.standard_normal([nb_sph.nbody]) nb_sph.vel = transpose(array([vx,vy,vz])).astype(float32) # here we should limit the speed according to max speed phis = lininterp1d(nb_sph.rxyz().astype(float32),r.astype(float32),phi.astype(float32)) vm = 0.95*sqrt(-2*phis) vn = nb_sph.vn() vf = where(vn>vm,vm/vn,1) vf.shape = (len(vf),1) nb_sph.vel = nb_sph.vel * vf # other info phi dphi = libgrid.get_First_Derivative(phi,r) vcirc = libdisk.Vcirc(r,dphi) stats = {} stats['r'] = r stats['nn'] = nn stats['phi'] = phi stats['rho'] = rho stats['sigma'] = sigma stats['vc'] = vcirc return nb_sph,phi,stats def Get_Velocities_From_Spherical_Grid(self,select=None,eps=0.1,nr=128,rmax=100.0,UseTree=True,Tree=None,phi=None,ErrTolTheta=0.5,g=None,gm=None,NoDispertion=False,omega=None): ''' Computes velocities using the jeans equation in spherical coordinates. ''' if select!=None: nb_sph = self.select(select) else: nb_sph = self # build the Tree for nb self.getTree(force_computation=True,ErrTolTheta=ErrTolTheta) # create the grid G = libgrid.Spherical_1d_Grid(rmin=0,rmax=rmax,nr=nr,g=g,gm=gm) if phi==None: phi = G.get_PotentialMap(self,eps=eps,UseTree=UseTree) r = G.get_r() rho = G.get_DensityMap(nb_sph) nn = G.get_NumberMap(nb_sph) # dr dr = G.get_r(offr=1)-G.get_r(offr=0) # compute sigma sigma = libdisk.get_1d_Sigma_From_Rho_Phi(rho=rho,phi=phi,r=r,dr=dr) # correct sigma in case of rotation (we assume the rotation around z) if omega!=None: print "add rotation" e_jeans = 0.5*sigma*sigma e_rot = 0.5*r**2 * omega**2 e = e_jeans - e_rot if (e<0).any(): print "at some radius the kinetic specifig energy is less than zero\nYou should decrease omega." e = where(e<0,0,e) sigma = sqrt(2*e) # generate velocities for all particles sigmas = G.get_Interpolation(nb_sph.pos,sigma) if NoDispertion: vx = sigmas*ones(nb_sph.nbody) vy = sigmas*ones(nb_sph.nbody) vz = sigmas*ones(nb_sph.nbody) else: vx = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vy = sigmas*RandomArray.standard_normal([nb_sph.nbody]) vz = sigmas*RandomArray.standard_normal([nb_sph.nbody]) nb_sph.vel = transpose(array([vx,vy,vz])).astype(float32) # do not spin # add rotation #if omega!=None: # nb_sph.spin(omega=array([0,0,omega])) # here we should limit the speed according to max speed phis = G.get_Interpolation(nb_sph.pos,phi) vm = 0.95*sqrt(-2*phis) vn = nb_sph.vn() vf = where(vn>vm,vm/vn,1) vf.shape = (len(vf),1) nb_sph.vel = nb_sph.vel * vf # other info phi dphi = libgrid.get_First_Derivative(phi,r) vcirc = libdisk.Vcirc(r,dphi) stats = {} stats['r'] = r stats['nn'] = nn stats['phi'] = phi stats['rho'] = rho stats['sigma'] = sigma stats['vc'] = vcirc return nb_sph,phi,stats def Get_Velocities_From_Cylindrical_Grid(self,select='disk',disk=('gas','disk'),eps=0.1,nR=32,nz=32,nt=2,Rmax=100,zmin=-10,zmax=10,params=[None,None,None],UseTree=True,Tree=None,Phi=None,ErrTolTheta=0.5,AdaptativeSoftenning=False,g=None,gm=None,NoDispertion=False): ''' Computes velocities using the jeans equation in cylindrical coordinates. ''' mode_sigma_z = params[0] mode_sigma_r = params[1] mode_sigma_p = params[2] if params[0]==None: mode_sigma_z = {"name":"jeans","param":None} if params[1]==None: mode_sigma_r = {"name":"toomre","param":1.0} if params[2]==None: mode_sigma_p = {"name":"epicyclic_approximation","param":None} nb_cyl = self.select(select) # current component nb_dis = self.select(disk) # disk component, for Q computation # build the Tree for nb self.getTree(force_computation=True,ErrTolTheta=ErrTolTheta) # create the grid G = libgrid.Cylindrical_2drz_Grid(rmin=0,rmax=Rmax,nr=nR,zmin=zmin,zmax=zmax,nz=nz,g=g,gm=gm) R,z = G.get_rz() #################################### # compute Phi in a 2d rz grid #################################### # here, we could use Acc instead Phi = G.get_PotentialMap(self,eps=eps,UseTree=UseTree,AdaptativeSoftenning=AdaptativeSoftenning) Phi = libgrid.get_Symetrisation_Along_Axis(Phi,axis=1) #Accx,Accy,Accz = libgrid.get_AccelerationMap_On_Cylindrical_2dv_Grid(self,nR,nz,Rmax,zmin,zmax,eps=eps) #Ar = sqrt(Accx**2+Accy**2) #################################### # compute Phi (z=0) in a 2d rt grid #################################### Grt = libgrid.Cylindrical_2drt_Grid(rmin=0,rmax=Rmax,nr=nR,nt=nt,z=0,g=g,gm=gm) Accx,Accy,Accz = Grt.get_AccelerationMap(self,eps=eps,UseTree=UseTree,AdaptativeSoftenning=AdaptativeSoftenning) Ar = sqrt(Accx**2+Accy**2) Ar = sum(Ar,axis=1)/nt Rp,tp = Grt.get_rt() Phi0 = Phi[:,nz/2] # not used dPhi0 = Ar d2Phi0 = libgrid.get_First_Derivative(dPhi0,R) # density rho = G.get_DensityMap(nb_cyl,offz=-0.5) rho = libgrid.get_Symetrisation_Along_Axis(rho,axis=1) # number per bin nn = G.get_NumberMap(nb_cyl,offz=-0.5) Sden = G.get_SurfaceDensityMap(nb_cyl,offz=-0.5) Sdend = G.get_SurfaceDensityMap(nb_dis,offz=-0.5) # compute frequencies (in the plane) kappa = libdisk.Kappa(R,dPhi0,d2Phi0) omega = libdisk.Omega(R,dPhi0) vcirc = libdisk.Vcirc(R,dPhi0) nu = libdisk.Nu(z,Phi) # compute sigma_z if mode_sigma_z['name']=='jeans': R1,z1 = G.get_rz(offz=0) R2,z2 = G.get_rz(offz=1) dz = z2-z1 sigma_z = libdisk.get_2d_Sigma_From_Rho_Phi(rho=rho,Phi=Phi,z=z,dz=dz) sigma_z2 = sigma_z**2 elif mode_sigma_z['name']=='surface density': """sigma_z2 = pi*G*Sden*Hz""" print "mode sigma z : 'surface density', not implemented yet" sys.exit() # compute sigma_r if mode_sigma_r['name']=='epicyclic_approximation': beta2 = mode_sigma_r['param'] f = where( kappa**2>0 , (1./beta2) * (nu**2/kappa**2) , 1.0 ) f.shape = (nR,1) sigma_r2 = sigma_z2 * f elif mode_sigma_r['name']=='isothropic': sigma_r2 = sigma_z2 elif mode_sigma_r['name']=='toomre': Q = mode_sigma_r['param'] Gg = 1.0 sr = where(kappa>0,Q*3.36*Gg*Sdend/kappa,sigma_z[:,nz/2]) sr.shape = (nR,1) sigma_r2 = ones((nR,nz)) * sr sigma_r2 = sigma_r2**2 elif mode_sigma_r['name']=='constant': sr = mode_sigma_r['param'] sigma_r2 = ones((nR,nz)) * sr sigma_r2 = sigma_r2**2 # compute sigma_p if mode_sigma_p['name']=='epicyclic_approximation': f = where( omega**2>0 , (1/4.0) * (kappa**2/omega**2) , 1.0 ) f.shape = (nR,1) sigma_p2 = sigma_r2 * f elif mode_sigma_p['name']=='isothropic': sigma_p2 = sigma_z2 notok = True count = 0 while notok: count = count + 1 print "compute vm" # compute vm sr2 = sigma_r2[:,nz/2] # should not be only in the plane sp2 = sigma_p2[:,nz/2] # should not be only in the plane vc = vcirc # should not be only in the plane T1 = vc**2 T2 = + sr2 - sp2 T3 = where(Sden>0,R/Sden,0)* libgrid.get_First_Derivative(Sden*sr2,R) vm2 = T1 + T2 + T3 # if vm2 < 0 c = (vm2<0) if sum(c)>0: print "Get_Velocities_From_Cylindrical_Grid : vm2 < 0 for %d elements"%(sum(c)) ''' vm2 = where(c,0,vm2) dsr2 = where(c,(T1+T2+T3)/2.,0) # energie qu'il faut retirer a sr dsp2 = where(c,(T1+T2+T3)/2.,0) # energie qu'il faut retirer a sp # take energy from sigma_r and sigma_p sigma_r2 = transpose(transpose(sigma_r2) + dsr2) sigma_p2 = transpose(transpose(sigma_p2) + dsp2) ''' E = sr2 + sp2 + vm2 if sum(E<0) != 0: print "-----------------------------------------------------" for i in range(len(R)): print R[i],vc[i]**2,sr2[i],sp2[i],vm2[i],sigma_z[i,nz/2]**2 print "-----------------------------------------------------" print "Get_Velocities_From_Cylindrical_Grid : we are in trouble here..." raise "E<0" vm2 = where(c,E/3.,vm2) sr2 = where(c,E/3.,sr2) sp2 = where(c,E/3.,sp2) sigma_r2 = transpose(ones((nz,nR)) * sr2) sigma_p2 = transpose(ones((nz,nR)) * sp2) if count > 0: notok = False else: notok = False # old implementation #vm2 = where(c,T1,vm2) #dsr2 = where(c,-(T2+T3)/2.,0) #dsp2 = where(c,-(T2+T3)/2.,0) # check again c = (vm2<0).astype(int) nvm = sum(c) if sum(nvm)>0: print "WARNING : %d cells still have vm<0 !!!"%(nvm) print "Vc^2 < 0 !!!" vm2 = where(c,0,vm2) vm = where(vm2>0,sqrt(vm2),0) # generate velocities for all particles sigma_r2s = G.get_Interpolation(nb_cyl.pos,sigma_r2) sigma_p2s = G.get_Interpolation(nb_cyl.pos,sigma_p2) sigma_z2s = G.get_Interpolation(nb_cyl.pos,sigma_z2) sigma_rs = where(sigma_r2s>0,sqrt(sigma_r2s),0) sigma_ps = where(sigma_p2s>0,sqrt(sigma_p2s),0) sigma_zs = where(sigma_z2s>0,sqrt(sigma_z2s),0) vcircs = G.get_r_Interpolation(nb_cyl.pos,vcirc) vms = G.get_r_Interpolation(nb_cyl.pos,vm) if NoDispertion: vr = sigma_rs vp = sigma_ps*0 + vms vz = sigma_zs else: vr = sigma_rs*RandomArray.standard_normal([nb_cyl.nbody]) vp = sigma_ps*RandomArray.standard_normal([nb_cyl.nbody]) + vms vz = sigma_zs*RandomArray.standard_normal([nb_cyl.nbody]) vel = transpose(array([vr,vp,vz])).astype(float32) nb_cyl.vel = libutil.vel_cyl2cart(nb_cyl.pos,vel) # here we should limit the speed according to max speed phis = G.get_Interpolation(nb_cyl.pos,Phi) vmax = 0.95*sqrt(-2*phis) vn = nb_cyl.vn() vf = where(vn>vmax,vmax/vn,1) vf.shape = (len(vf),1) nb_cyl.vel = nb_cyl.vel * vf # some output sr = sigma_r = sqrt(sigma_r2[:,nz/2]) sp = sigma_p = sqrt(sigma_p2[:,nz/2]) sz = sigma_z = sqrt(sigma_z2[:,nz/2]) Q = where((Sden>0),sr*kappa/(3.36*Sdend),0) stats = {} stats['R'] = R stats['z'] = z stats['vc'] = vc stats['vm'] = vm stats['sr'] = sr stats['sp'] = sp stats['sz'] = sz stats['kappa'] = kappa stats['omega'] = omega stats['nu'] = nu stats['Sden'] = Sden stats['Sdend'] = Sdend stats['Q'] = Q #stats['Ar'] = Ar stats['rho'] = rho stats['phi'] = Phi stats['nn'] = nn stats['sigma_z']= sqrt(sigma_z2) stats['Phi0'] = Phi0 stats['dPhi0'] = dPhi0 stats['d2Phi0'] = d2Phi0 return nb_cyl,Phi,stats ############################################ # # evolution routines # ############################################ def IntegrateUsingRK(self,tstart=0,dt=1,dt0=1e-5,epsx=1.e-13,epsv=1.e-13): """ Integrate the equation of motion using RK78 integrator. tstart : initial time dt : interval time dt0 : inital dt epsx : position precision epsv : velocity precision tmin,tmax,dt,dtout,epsx,epsv,filename """ tend = tstart + dt self.pos,self.vel,self.atime,self.dt = nbdrklib.IntegrateOverDt(self.pos.astype(float),self.vel.astype(float),self.mass.astype(float),tstart,tend,dt,epsx,epsv) self.pos = self.pos.astype(float32) self.vel = self.vel.astype(float32) ################################# # # Thermodynamic functions # ################################# def U(self): ''' Return the gas specific energy of the model. The output is an nx1 float array. ''' return self.u def Rho(self): ''' Return the gas density of the model. The output is an nx1 float array. ''' try: a3inv = 1./self.atime**3 except: a3inv = 1. if self.unitsparameters.get('HubbleParam') == 1: self.log.write("assuming non cosmological simulation") a3inv = 1. try: rho = self.rho*a3inv return rho except: return self.rho def T(self): ''' Return the gas temperature of the model. The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') xi = self.unitsparameters.get('xi') ionisation = self.unitsparameters.get('ionisation') mu = thermodyn.MeanWeight(xi,ionisation) mh = ctes.PROTONMASS.into(self.localsystem_of_units) k = ctes.BOLTZMANN.into(self.localsystem_of_units) thermopars = {"k":k,"mh":mh,"mu":mu,"gamma":gamma} T = where((self.u>0),thermodyn.Tru(self.Rho(),self.u,thermopars),0) return T def MeanWeight(self): ''' Return the mean weight of a model, taking into account heating by UV source. The output is an nx1 float array. ''' xi = self.unitsparameters.get('xi') Redshift = 1./self.atime - 1. UnitDensity_in_cgs = self.localsystem_of_units.get_UnitDensity_in_cgs() UnitEnergy_in_cgs = self.localsystem_of_units.get_UnitEnergy_in_cgs() UnitMass_in_g = self.localsystem_of_units.get_UnitMass_in_g() HubbleParam = self.hubbleparam # 0) convert into cgs Density = self.rho.astype(float) *UnitDensity_in_cgs * (HubbleParam*HubbleParam) / self.atime**3 Egyspec = self.u.astype(float) *UnitEnergy_in_cgs/UnitMass_in_g # 1) compute mu MeanWeight,Lambda = coolinglib.cooling(Egyspec,Density,xi,Redshift) return MeanWeight.astype(float32) def Tmu(self): ''' Return the gas temperature of the model. The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') mh = ctes.PROTONMASS.into(self.localsystem_of_units) k = ctes.BOLTZMANN.into(self.localsystem_of_units) T = (gamma-1) *self.MeanWeight().astype(float) *mh/k * self.u return T.astype(float32) def A(self): ''' Return the gas entropy of the model. The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') xi = self.unitsparameters.get('xi') ionisation = self.unitsparameters.get('ionisation') mu = thermodyn.MeanWeight(xi,ionisation) mh = ctes.PROTONMASS.into(self.localsystem_of_units) k = ctes.BOLTZMANN.into(self.localsystem_of_units) thermopars = {"k":k,"mh":mh,"mu":mu,"gamma":gamma} A = where((self.u>0),thermodyn.Aru(self.Rho(),self.u,thermopars),0) return A def P(self): ''' Return the gas pressure of the model. The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') xi = self.unitsparameters.get('xi') ionisation = self.unitsparameters.get('ionisation') mu = thermodyn.MeanWeight(xi,ionisation) mh = ctes.PROTONMASS.into(self.localsystem_of_units) k = ctes.BOLTZMANN.into(self.localsystem_of_units) thermopars = {"k":k,"mh":mh,"mu":mu,"gamma":gamma} P = where((self.u>0),thermodyn.Pru(self.Rho(),self.u,thermopars),0) return P def Tcool(self,coolingfile=None): ''' Return the cooling time of the model. The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') xi = self.unitsparameters.get('xi') ionisation = self.unitsparameters.get('ionisation') metalicity = self.unitsparameters.get('metalicity') hubbleparam= self.unitsparameters.get('HubbleParam') if coolingfile==None: coolingfile = self.unitsparameters.get('coolingfile') mu = thermodyn.MeanWeight(xi,ionisation) mh = ctes.PROTONMASS.into(self.localsystem_of_units) k = ctes.BOLTZMANN.into(self.localsystem_of_units) try: if self.hubbleparam != hubbleparam: self.log.write("Warning (Tcool): using hubbleparam=%f, but self.hubbleparam=%f"%(hubbleparam,self.hubbleparam)) except: pass thermopars = {"k":k,"mh":mh,"mu":mu,"gamma":gamma,"Xi":xi,"metalicity":metalicity,"hubbleparam":hubbleparam} tc,c = thermodyn.CoolingTime(self.Rho(),self.u,self.localsystem_of_units,thermopars,coolingfile) tc = where(c,tc,0) return tc def Ne(self): ''' Return the electron density of the model. The output is an nx1 float array. ''' xi = self.unitsparameters.get('xi') ionisation = self.unitsparameters.get('ionisation') mh = ctes.PROTONMASS.into(self.localsystem_of_units) thermopars = {"mh":mh,"Xi":xi,"ionisation":ionisation} ne = thermodyn.ElectronDensity(self.Rho(),thermopars) return ne def S(self): ''' Return the `entropy` of the model, defined as S = T * Ne^(1-gamma) The output is an nx1 float array. ''' gamma = self.unitsparameters.get('gamma') s = self.T()*self.Ne()**(1.-gamma) return s def Lum(self): ''' Return the luminosty of the model, defined as Lum = m*u/Tcool = m*Lambda/rho The output is an nx1 float array. ''' Lum = self.mass*self.u/self.Tcool() return Lum #################################################################################################################################### # # NBODY REDIRECTOR # #################################################################################################################################### if FORMATSDIR != None: formatsfiles = glob.glob(os.path.join(FORMATSDIR,'*.py')) def Nbody(*arg,**kw): """ The aim of this function is simply to return to the right class """ # default value nb = None if not kw.has_key('ftype'): kw['ftype']='default' # find the right file formatfile = os.path.join(FORMATSDIR,"%s.py"%kw['ftype']) try: formatsfiles.index(formatfile) except ValueError: print "format %s is unknown !"%kw['ftype'] print "%s does not exists"%formatfile sys.exit() # this does not work, because NbodyDefault is unknown #sys.path.append(FORMATSDIR) #__import__(kw['ftype'],) # instead, we use execfile(formatfile) # check class name class_name = "Nbody_%s"%kw['ftype'] try: dir().index(class_name) except ValueError: print "format %s is unknown !"%kw['ftype'] sys.exit() ## not very good class_name = eval(class_name) nb = class_name(*arg,**kw) nb._formatfile = formatfile return nb diff --git a/pNbody/main.pyc b/pNbody/main.pyc deleted file mode 100644 index 79e788f..0000000 Binary files a/pNbody/main.pyc and /dev/null differ diff --git a/pNbody/mpi.py b/pNbody/mpi.py index 8b1aae4..1a54f5a 100644 --- a/pNbody/mpi.py +++ b/pNbody/mpi.py @@ -1,1153 +1,1154 @@ # -*- coding: iso-8859-1 -*- from numpy import * import types import libutil import sys class myMPI: def __init__(self): pass try: from mpi4py import MPI - comm = MPI.COMM_WORLD - ThisTask = comm.Get_rank() - NTask = comm.Get_size() - Procnm = MPI.Get_processor_name() - Rank = ThisTask + comm = MPI.COMM_WORLD + comm_self = MPI.COMM_SELF + ThisTask = comm.Get_rank() + NTask = comm.Get_size() + Procnm = MPI.Get_processor_name() + Rank = ThisTask except: - MPI = myMPI() - MPI.SUM = sum - ThisTask = 0 - NTask = 1 - Rank = ThisTask - Procnm = "" + MPI = myMPI() + MPI.SUM = sum + ThisTask = 0 + NTask = 1 + Rank = ThisTask + Procnm = "" def mpi_IsMaster(): return ThisTask==0 def mpi_Rank(): return ThisTask def mpi_ThisTask(): return ThisTask def mpi_NTask(): return NTask def mpi_barrier(): if NTask>1: comm.barrier() else: pass ##################################### # # Output functions # ##################################### def mpi_iprint(msg,mode=None): ''' Synchronized print, including info on node. ''' msg = "[%d] : %s"%(ThisTask,msg) pprint(msg) def mpi_pprint(msg): ''' Synchronized print. ''' if NTask>1: MPI.pprint(msg) else: print msg def mpi_rprint(msg): ''' Rooted print. ''' if NTask>1: MPI.rprint(msg) else: print msg ##################################### # # communication functions # ##################################### def mpi_send(x,dest): ''' Send x to node dest. When there is only one is defined, it does nothing. ''' if NTask>1: return comm.send(x, dest = dest) else: pass def mpi_recv(source): ''' Return a variable sent by node \var{source}. When there is only one is defined, it does nothing. ''' if NTask>1: return comm.recv(source = source) else: pass def mpi_sendrecv(x,dest,source): ''' ''' if NTask>1: return comm.sendrecv(x,dest=dest,source = source) else: pass def mpi_reduce(x,root=0,op=MPI.SUM): """ Reduce x from all node only for root. When there is only one is defined, the function return x. """ if NTask>1: sx = comm.reduce(x,op=op,root=root) else: sx = x return sx def mpi_allreduce(x,op=MPI.SUM): """ Reduce x from all node for all nodes. When there is only one is defined, the function return x. """ if NTask>1: sx = comm.allreduce(x,op=op) else: sx = x return sx def mpi_bcast(x,root=0): ''' Broadcast from node root the variable x. When there is only one is defined, it simplay returns x. ''' if NTask>1: x = comm.bcast(x,root=root) else: x = x return x def mpi_gather(x,root=0): """ Gather x from all nodes to node dest. Returns a list. """ if NTask>1: x = comm.gather(x,root=root) else: x = [x] return x def mpi_allgather(x): """ Gather x from all to all. Returns a list. """ if NTask>1: x = comm.allgather(x) else: x = [x] return x def mpi_AllgatherAndConcatArray(vec): ''' AllGather array vec and concatenate it in a unique array (concatenation order is reversed). ''' if NTask>1: vec_all = array([],vec.dtype) if vec.ndim==1: vec_all.shape = (0,) else: vec_all.shape = (0,vec.shape[1]) if ThisTask==0: for Task in range(NTask-1,-1,-1): if Task != 0: v = comm.recv(source=Task) vec_all = concatenate((vec_all,v)) else: vec_all = concatenate((vec_all,vec)) else: comm.send(vec, dest = 0) # send to everybody xi = comm.bcast(vec_all) return xi else: return vec ##################################### # # array functions # ##################################### def mpi_sum(x): """ Sum elements of array x. """ if NTask>1: sx = comm.allreduce(sum(x),op=MPI.SUM) else: sx = sum(x) return sx def mpi_min(x): """ Minimum element of array x. """ if NTask>1: mx = comm.allreduce(min(x),op=MPI.MIN) else: mx = min(x) return mx def mpi_max(x): """ Maximum element of array x. """ if NTask>1: mx = comm.allreduce(max(x),op=MPI.MAX) else: mx = max(x) return mx def mpi_mean(x): """ Mean of elements of array x. """ if NTask>1: sm = comm.allreduce(sum(x),op=MPI.SUM) sn = comm.allreduce(len(x),op=MPI.SUM) mn = sm/sn else: mn = x.mean() return mn def mpi_len(x): """ Lenght of array x. """ if NTask>1: ln = comm.allreduce(len(x),op=MPI.SUM) else: ln = len(x) return ln def mpi_arange(n): """ Create an integer array containing elements from 0 to n spreaded over all nodes. """ if NTask>1: ns = array(comm.allgather(n)) c = ((NTask-1-ThisTask) > arange(len(ns)-1,-1,-1) ) i = sum(ns*c) v = arange(i,i+ns[ThisTask]) else: v = arange(n) return v def mpi_sarange(npart_all): """ Create an integer array containing elements from 0 to n, spreaded over all nodes. The repartition of elements and type of elements over nodes is given by the array npart_all """ npart = npart_all[ThisTask] if NTask>1: v = array([],int) o = 0 for i in arange(len(npart)): # number of particles of this type for this proc n = npart_all[ThisTask,i] n_tot = sum(npart_all[:,i]) ns = array(comm.allgather(n)) c = ((NTask-1-ThisTask) > arange(len(ns)-1,-1,-1) ) j = sum(ns*c) vj = arange(o+j,o+j+ns[ThisTask]) v = concatenate((v,vj)) o = o + n_tot else: n = sum(ravel(npart)) v = arange(n) return v def mpi_argmax(x): """ Find the arument of the amximum value in x. idx = (p,i) : where i = index in proc p """ if NTask>1: # 1) get all max and all argmax maxs = array(comm.allgather(max(x))) args = array(comm.allgather(argmax(x))) # 2) choose the right one p = argmax(maxs) # proc number i = args[p] # index in proc p else: p = 0 i = argmax(x) return p,i def mpi_argmin(x): """ Find the arument of the maximum value in x. idx = (p,i) : where i = index in proc p """ if NTask>1: # 1) get all mins and all argmins mins = array(comm.allgather(min(x))) args = array(comm.allgather(argmin(x))) # 2) choose the right one p = argmin(mins) # proc number i = args[p] # index in proc p else: p = 0 i = argmin(x) return p,i def mpi_getval(x,idx): """ Return the value of array x corresponding to the index idx. idx = (p,i) : where i = index in proc p equivalent to x[i] from proc p """ p = idx[0] i = idx[1] if NTask>1: # send to everybody xi = comm.bcast(x[i]) else: xi = x[i] return xi def mpi_histogram(x,bins): """ Return an histogram of vector x binned using binx. """ # histogram n = searchsorted(sort(x),bins) n = concatenate([n,[len(x)]]) hx = n[1:]-n[:-1] if NTask>1: hx = comm.allreduce(hx,op=MPI.SUM) return hx ##################################### # # exchange function # ##################################### ################################# def mpi_find_a_toTask(begTask,fromTask,ex_table,delta_n): ################################# ''' This function is used to find recursively an exange table ''' for toTask in range(begTask,NTask): if delta_n[toTask] > 0: # find a node who accept particles if delta_n[toTask] + delta_n[fromTask] >= 0: # can give all delta = -delta_n[fromTask] else: # can give only a fiew delta = +delta_n[toTask] ex_table[fromTask,toTask] = +delta ex_table[toTask,fromTask] = -delta delta_n[fromTask] = delta_n[fromTask] +delta delta_n[toTask] = delta_n[toTask] -delta if delta_n[fromTask] == 0: # everything has been distributed return ex_table else: return mpi_find_a_toTask(begTask+1,fromTask,ex_table,delta_n) ################################# def mpi_GetExchangeTable(n_i): ################################# ''' This function returns the exchange table ''' # n_i : initial repartition ntot = sum(n_i) if ntot==0: return zeros((NTask,NTask)) # final repartition n_f = zeros(NTask) for Task in range(NTask-1,-1,-1): n_f[Task] = ntot/NTask + ntot%NTask*(Task==0) # delta delta_n = n_f - n_i # find who gives what to who ex_table = zeros((NTask,NTask)) for fromTask in range(NTask): if delta_n[fromTask] < 0: # the node gives particles ex_table = mpi_find_a_toTask(0,fromTask,ex_table,delta_n) return ex_table ################################# def mpi_ExchangeFromTable(T,procs,ids,vec,num): ################################# ''' Exchange an array according to a transfer array T T : exchange table procs : list of processor (from Tree.GetExchanges()) ids : list of id (from Tree.GetExchanges()) vec : vector to exchange num : id correspondings to particles ''' # now, we have to send / recv SendPart = T[NTask-1-ThisTask,:] RecvPart = T[:,ThisTask] new_procs = array([]) new_vec = array([]) if vec.ndim == 1: #new_vec.shape = (0,3) pass elif vec.ndim == 2: new_vec.shape = (0,3) # send/recv (1) for i in range(NTask): if i > ThisTask: # here, we send to i N = T[NTask-1-ThisTask,i] comm.send(N,i) if N>0: to_procs = compress((procs==i),procs,axis=0) to_ids = compress((procs==i),ids,axis=0) to_vec = libutil.compress_from_lst(vec,num,to_ids) comm.send(to_vec,i) vec = libutil.compress_from_lst(vec,num,to_ids,reject=True) num = libutil.compress_from_lst(num,num,to_ids,reject=True) elif i < ThisTask: N = comm.recv(i) if N > 0: new_vec = concatenate((new_vec,comm.recv(i))) # send/recv (1) for i in range(NTask): if i < ThisTask: # here, we send to i N = T[NTask-1-ThisTask,i] comm.send(N,i) if N>0: to_procs = compress((procs==i),procs,axis=0) to_ids = compress((procs==i),ids,axis=0) to_vec = libutil.compress_from_lst(vec,num,to_ids) comm.send(to_vec,i) vec = libutil.compress_from_lst(vec,num,to_ids,reject=True) num = libutil.compress_from_lst(num,num,to_ids,reject=True) elif i > ThisTask: N = comm.recv(i) if N > 0: new_vec = concatenate((new_vec,comm.recv(i))) # check c = (new_procs!=ThisTask).astype(int) if sum(c)>0: print "here we are in trouble" sys.exit() return concatenate((vec,new_vec)) # !!!!!!!!!!!!!!!!!!!!!!! this has to be changed .... ##################################### # # io functions # ##################################### ################################# def mpi_ReadAndSendBlock(f,data_type,shape=None,byteorder=sys.byteorder,split=None): ################################# ''' Read and brodcast a binary block. data_type = int,float32,float or data_type = array shape = tuple ''' ##################### # read first header ##################### if ThisTask==0: try: nb1 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb1.byteswap(True) nb1 = nb1[0] except IndexError: raise "ReadBlockError" ##################### # read a tuple ##################### if type(data_type) == types.TupleType: if ThisTask==0: data = [] for tpe in data_type: if type(tpe) == int: val = f.read(tpe) else: bytes = dtype(tpe).itemsize val = fromstring(f.read(bytes),tpe) if sys.byteorder != byteorder: val.byteswap(True) val = val[0] data.append(val) # send the data if NTask > 1: if ThisTask==0: for Task in range(1,NTask): comm.send(data, dest = Task) else: data = comm.recv(source = 0) ##################### # read an array ##################### else: if split: if ThisTask==0: bytes_per_elt = data_type.bytes if shape != None: ndim = shape[1] else: ndim = 1 nelt = nb1/bytes_per_elt/ndim nleft = nelt nread = nelt/NTask for Task in range(NTask-1,-1,-1): if nleft < 2*nread and nleft > nread: nread = nleft nleft = nleft-nread data = f.read(nread*bytes_per_elt*ndim) shape = (nread,ndim) if Task==ThisTask: # this should be last data = fromstring(data,data_type) if sys.byteorder != byteorder: data.byteswap(True) data.shape=shape else: comm.send(data, dest = Task) comm.send(shape, dest = Task) else : data = comm.recv(source = 0) shape = comm.recv(source = 0) data = fromstring(data,data_type) if sys.byteorder != byteorder: data.byteswap(True) data.shape=shape else: if ThisTask==0: data = fromstring(f.read(nb1),data_type) if sys.byteorder != byteorder: data.byteswap(True) # send the data if NTask > 1: if ThisTask==0: for Task in range(1,NTask): comm.send(data, dest = Task) else: data = comm.recv(source = 0) ##################### # read second header ##################### if ThisTask==0: nb2 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb2.byteswap(True) nb2 = nb2[0] if nb1 != nb2: raise "ReadBlockError" # reshape if needed if split: # already done before pass else: if shape != None: data.shape=shape return data ################################# def mpi_ReadAndSendArray(f,data_type,shape=None,byteorder=sys.byteorder,npart=None): ################################# ''' Read and Brodcast a binary block assuming it contains an array. ''' oneD = False if len(shape)==1: shape = (shape[0],1) oneD = True if shape != None: n_dim = shape[1] data = array([],data_type) data.shape = (0,shape[1]) n_elts = shape[0] else: raise "mpi_ReadAndSendArray : var shape must be defined here." #n_dim = 1 #data = array([],data_type) #data.shape = (0,) nbytes_per_elt = dtype(data_type).itemsize * n_dim n_left = nbytes_per_elt * n_elts # check npart if npart != None: ntype = len(npart) else: ntype = 1 npart = array([n_elts]) # this may be wrong in 64b, as nb1 is not equal to the number of bytes in block #if sum(npart,0) != n_elts: # raise "We are in trouble here : sum(npart)=%d n_elts=%d"%(sum(npart,0),n_elts) ##################### # read first header ##################### if ThisTask==0: try: nb1 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb1.byteswap(True) nb1 = nb1[0] # this may be wrong in 64b, as nb1 is not equal to the number of bytes in block #if nb1 != n_left: # raise "We are in trouble here : nb1=%d n_left=%d"%(nb1,n_left) except IndexError: raise "ReadBlockError" ##################### # read data ##################### if ThisTask == 0: for i in range(ntype): n_write = libutil.get_n_per_task(npart[i],NTask) for Task in range(NTask-1,-1,-1): sdata = f.read(nbytes_per_elt*n_write[Task]) # send block for a slave if Task != 0: comm.send(sdata, dest = Task) # keep block for the master else: sdata = fromstring(sdata,data_type) if shape!=None: sdata.shape = (n_write[Task],shape[1]) data = concatenate((data,sdata)) else: for i in range(ntype): sdata = comm.recv(source = 0) sdata = fromstring(sdata,data_type) if shape!=None: ne = len(sdata)/shape[1] sdata.shape = (ne,shape[1]) data = concatenate((data,sdata)) if oneD: data.shape = (data.shape[0],) ##################### # read second header ##################### if ThisTask==0: nb2 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb2.byteswap(True) nb2 = nb2[0] if nb1 != nb2: raise "ReadBlockError" return data ################################# def mpi_GatherAndWriteArray(f,data,byteorder=sys.byteorder,npart=None): ################################# ''' Gather and array and write it in a binary block. data = array shape = tuple ''' # check npart if npart != None: pass else: npart = array([len(data)]) data_nbytes = mpi_allreduce(data.nbytes) if ThisTask==0: nbytes = array([data_nbytes],int32) if sys.byteorder != byteorder: nbytes.byteswap(True) # write header 1 f.write(nbytes.tostring()) # loop over particles type for i in range(len(npart)): n_write = libutil.get_n_per_task(npart[i],NTask) if ThisTask==0: for Task in range(NTask-1,-1,-1): if Task != 0: sdata = comm.recv(source = Task) else: i1 = sum(npart[:i]) i2 = i1+npart[i] sdata = data[i1:i2] if sys.byteorder != byteorder: sdata.byteswap(True) f.write(sdata.tostring()) else: i1 = sum(npart[:i]) i2 = i1+npart[i] comm.send(data[i1:i2], dest = 0) if ThisTask==0: # write header 2 f.write(nbytes.tostring()) ################################# def mpi_OldReadAndSendArray(f,data_type,shape=None,skip=None,byteorder=sys.byteorder,nlocal=None): ################################# ''' Read and Brodcast a binary block assuming it contains an array. The array is splitted acroding to the variable nlocal. data_type = array type shape = tuple nlocal : array NTask x Npart array NTask ''' ##################### # read first header ##################### if ThisTask==0: try: nb1 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb1.byteswap(True) nb1 = nb1[0] except IndexError: raise "ReadBlockError" ##################### # read an array ##################### if nlocal.ndim != 2: raise "ReadAndSendArray error","nlocal must be of rank 2" else: ntype = nlocal.shape[1] if shape != None: ndim = shape[1] data = array([],data_type) data.shape = (0,shape[1]) else: ndim = 1 data = array([],data_type) data.shape = (0,) nbytes_per_elt = data_type.bytes * ndim if ThisTask == 0: for i in range(ntype): for Task in range(NTask-1,-1,-1): sdata = f.read(nbytes_per_elt*nlocal[Task,i]) # send block for a slave if Task != 0: comm.send(sdata, dest = Task) # keep block for the master else: sdata = fromstring(sdata,data_type) if shape!=None: sdata.shape = (nlocal[ThisTask,i],shape[1]) data = concatenate((data,sdata)) else: for i in range(ntype): sdata = comm.recv(source = 0) sdata = fromstring(sdata,data_type) if shape!=None: sdata.shape = (nlocal[ThisTask,i],shape[1]) data = concatenate((data,sdata)) ##################### # read second header ##################### if ThisTask==0: nb2 = fromstring(f.read(4),int32) if sys.byteorder != byteorder: nb2.byteswap(True) nb2 = nb2[0] if nb1 != nb2: raise "ReadBlockError" return data ################################# def mpi_OldGatherAndWriteArray(f,data,byteorder=sys.byteorder,nlocal=None): ################################# ''' Gather and array and write it in a binary block. data = array shape = tuple ''' if nlocal.ndim != 2: raise "ReadAndSendArray error","nlocal must be of rank 2" else: ntype = nlocal.shape[1] data_size = mpi_allreduce(data.size()) if ThisTask==0: nbytes = array([data.type().bytes*data_size],int32) if sys.byteorder != byteorder: nbytes.byteswap(True) # write header 1 f.write(nbytes.tostring()) # loop over particles type i1 = 0 for i in range(ntype): if ThisTask==0: for Task in range(NTask-1,-1,-1): if Task != 0: sdata = comm.recv(source = Task) else: i2 = i1 + nlocal[Task,i] sdata = data[i1:i2] i1 = i2 if sys.byteorder != byteorder: sdata.byteswap(True) f.write(sdata.tostring()) else: i2 = i1 + nlocal[ThisTask,i] comm.send(data[i1:i2], dest = 0) i1 = i2 if ThisTask==0: # write header 2 f.write(nbytes.tostring()) diff --git a/pNbody/mpi.pyc b/pNbody/mpi.pyc deleted file mode 100644 index 2274624..0000000 Binary files a/pNbody/mpi.pyc and /dev/null differ diff --git a/pNbody/units.py b/pNbody/units.py index 98d3d9d..c8b692a 100644 --- a/pNbody/units.py +++ b/pNbody/units.py @@ -1,587 +1,588 @@ import copy ################################################################################ class Units: ################################################################################ ''' Units ''' def __init__(self,symbol,factor=1.,power=1,ulist=[]): ''' ''' self.symbol = symbol self.factor = float(factor) self.power = int(power) self.ulist = ulist self.GatherBaseUnits() def set_symbol(self,symbol): self.ulist = [copy.deepcopy(self)] self.power = 1 self.factor = 1.0 self.symbol = symbol self.GatherBaseUnits() def GatherBaseUnits(self): ''' create the self.bases list (recursively) ''' self.bases = [] for unit in self.ulist: if len(unit.ulist)==0: # a base unit self.bases.append(unit) # add the unit else: # go deeper in the tree unit.GatherBaseUnits() self.bases = self.bases + unit.bases self.ReduceBaseUnits() def ReduceBaseUnits(self): allunits = {} factorTot = 1.0 for base in self.bases: factorTot = factorTot * base.factor if allunits.has_key(base.symbol): allunits[base.symbol] = allunits[base.symbol] + base.power else: allunits[base.symbol] = base.power self.bases = [] flag = 1 for key in allunits.keys(): factor = 1.0 if flag: factor = factorTot flag = 0 power = allunits[key] symbol = key ulist = [] self.bases.append(Units(symbol,factor,power,ulist)) def str_base(self): string = '' if self.factor != 1.0: string = string + self.factor.__str__() if len(string) > 0: string = string + ' ' if self.symbol == None: string = string + '-' else: string = string + self.symbol if self.power != 1: string = string + '^%d'%self.power return string def __str__(self): string = '' # (1) display using symbol if self.factor != 0: string = string+self.str_base() str1 = "[%s]"%string string = '' # (2) display using (base units) for base in self.bases: if len(string) > 0: string = string + ' ' string = string + base.str_base() str2 = "[%s]"%string if str2 != '[]': str1 = "%s\t%s"%(str1,str2) return str1 def __mul__(self,y): x = copy.deepcopy(self) y = copy.deepcopy(y) # multiply with a scalar if type(y)== int or type(y) == float : x.factor = x.factor*y x.mul_list(x.ulist,y) x.GatherBaseUnits() return x # multiply with a unit elif type(y)==type(x): factor = x.factor * y.factor if x.symbol == y.symbol: # same units x.symbol = x.symbol x.power = x.power + y.power x.ulist = x.ulist + y.ulist x.GatherBaseUnits() return x else: # different unit symbol = None power = 1 factor = 1.0 if x.ulist == []: # case of a base unit x.ulist = [copy.deepcopy(x)] if y.ulist == []: # case of a base unit y.ulist = [copy.deepcopy(y)] ulist = x.ulist + y.ulist return Units(symbol,factor,power,ulist) def __div__(self,y): x = copy.deepcopy(self) y = copy.deepcopy(y) # divide with a scalar if type(y) == int or type(y) == float: x.factor = x.factor/y x.div_list(x.ulist,y) x.GatherBaseUnits() return x # divide with a unit elif type(y)==type(x): factor = x.factor / y.factor if x.symbol == y.symbol: # same units x.symbol = None x.power = 0 x.ulist = [] x.GatherBaseUnits() return x else: # different unit symbol = None power = 1 factor = 1.0 if x.ulist == []: # case of a base unit x.ulist = [copy.deepcopy(x)] if y.ulist == []: # case of a base unit y.ulist = [copy.deepcopy(y)] y.pow_list(y.ulist,-1) y.GatherBaseUnits() ulist = x.ulist + y.ulist return Units(symbol,factor,power,ulist) def __rdiv__(self,y): x = copy.deepcopy(self) # divide with a scalar if type(y)==int or type(y)==float: return y*x**-1 else: raise "TypeError","unsupported operand type(s) for / or div()" def __pow__(self,y): x = copy.deepcopy(self) y = copy.deepcopy(y) # power with a scalar if type(y)==int or type(y)==float: x.factor = x.factor**y x.power = x.power*y x.pow_list(x.ulist,y) x.GatherBaseUnits() return x else: raise "TypeError","unsupported operand type(s) for ** or pow()" def mul_list(self,ulist,y): for u in ulist: if len(u.ulist)==0: # a base unit u.factor = u.factor*y else: # go deeper in the tree u.mul_list(u.ulist,y) u.GatherBaseUnits() def div_list(self,ulist,y): for u in ulist: if len(u.ulist)==0: # a base unit u.factor = u.factor/y else: # go deeper in the tree u.mul_list(u.ulist,y) u.GatherBaseUnits() def pow_list(self,ulist,y): for u in ulist: if len(u.ulist)==0: # a base unit u.factor = u.factor**y u.power = u.power*y else: # go deeper in the tree u.pow_list(u.ulist,y) u.GatherBaseUnits() __rmul__ = __mul__ ################################################################################ class UnitSystem: ################################################################################ ''' Units system ''' def __init__(self,UnitSysName,UnitLst): ''' UnitDic = {'lenght':UnitLenght,'mass':UnitMass,'time':UnitTime} The units used to define a system of units must be "basic", i.e., not based on several units ''' UnitDic = {} for u in UnitLst: u.GatherBaseUnits() if len(u.bases) > 1: raise "UnitSystemError","%s is not a pure Unit."%(u.symbol) elif len(u.bases) == 1: UnitDic[u.bases[0].symbol] = u.bases[0] elif len(u.bases) == 0: # base unit UnitDic[u.symbol] = u dic_of_factors = {} dic_of_powers = {} for UnitType in UnitDic.keys(): u = UnitDic[UnitType] if u.bases ==[]: # base unit dic_of_factors[UnitType] = 1/u.factor dic_of_powers[UnitType] = u.power else: dic_of_factors[UnitType] = 1/u.bases[0].factor dic_of_powers[UnitType] = u.bases[0].power self.UnitDic = UnitDic self.dic_of_factors = dic_of_factors # value of base unit in current units self.dic_of_powers = dic_of_powers # not usefull self.UnitLength = self.UnitDic['m'] self.UnitMass = self.UnitDic['kg'] self.UnitTime = self.UnitDic['s'] self.UnitVelocity = self.UnitDic['m']/self.UnitDic['s'] self.UnitSpecEnergy = self.UnitVelocity**2 self.UnitEnergy = self.UnitSpecEnergy*self.UnitMass self.UnitDensity = self.UnitMass/self.UnitLength**3 self.UnitSurfaceDensity = self.UnitMass/self.UnitLength**2 + self.UnitSurface = self.UnitLength**2 def get_UnitLength_in_cm(self): f = self.UnitDic['m'].factor c = PhysCte(f,Unit_m) return c.into(cgs) def get_UnitMass_in_g(self): f = self.UnitDic['kg'].factor c = PhysCte(f,Unit_kg) return c.into(cgs) def get_UnitTime_in_s(self): f = self.UnitDic['s'].factor c = PhysCte(f,Unit_s) return c.into(cgs) def get_UnitVelocity_in_cm_per_s(self): return self.get_UnitLength_in_cm()/self.get_UnitTime_in_s() def get_UnitEnergy_in_cgs(self): return self.get_UnitMass_in_g()*self.get_UnitVelocity_in_cm_per_s()**2 def get_UnitDensity_in_cgs(self): return self.get_UnitMass_in_g()/self.get_UnitLength_in_cm()**3 def info(self): ''' print some info ''' print "UnitLength_in_cm =%g"%(self.get_UnitLength_in_cm()) print "UnitVelocity_in_cm_per_s =%g"%(self.get_UnitVelocity_in_cm_per_s()) print "UnitMass_in_g =%g"%(self.get_UnitMass_in_g()) def convertionFactorTo(self,newUnits): ''' return the conversion factor to obtain the new units ''' f = 1.0 # loop over all base units of the new Units if newUnits.bases==[]: baselist = [newUnits] else: baselist = newUnits.bases for base in baselist: factor = base.factor symbol = base.symbol power = base.power # multiply by the right factor f = f /( factor * (self.dic_of_factors[symbol])**power ) return f def into(self,newUnits): ''' return into the new units ''' return PhysCte(self.convertionFactorTo(newUnits),newUnits) ################################################################################ class PhysCte: ################################################################################ ''' Physical constant ''' def __init__(self,value,Unit): #super(PhysCte,self).__init__(value) self.value = value self.Unit = Unit def factor_to_base(self): if self.Unit.bases == []: factor = self.Unit.factor else: factor = self.Unit.bases[0].factor return factor def __str__(self): return float.__str__(float(self.value)) + ' %s'%(self.Unit) def into(self,SystemName): # here, we could return an object f = self.value if self.Unit.bases==[]: unit = self.Unit symbol = unit.symbol f = f * unit.factor * (SystemName.dic_of_factors[symbol])**unit.power else: for unit in self.Unit.bases: symbol = unit.symbol f = f * unit.factor * (SystemName.dic_of_factors[symbol])**unit.power return float(f) def __float__(self): return float(self.value) def __mul__(self,y): pass def __add__(self,y): pass def __sub__(self,y): pass def __div__(self,y): pass def __pow__(self,y): pass ################################################################################ # define some units ################################################################################ # base units Unit_m = Units('m') Unit_kg = Units('kg') Unit_s = Units('s') Unit_mol = Units('mol') Unit_C = Units('C') Unit_K = Units('K') # other lenght units Unit_cm = 0.01*Unit_m Unit_cm.set_symbol('cm') Unit_km = 1e3*Unit_m Unit_km.set_symbol('km') Unit_Mm = 1e6*Unit_m Unit_Mm.set_symbol('Mm') Unit_Gm = 1e9*Unit_m Unit_Gm.set_symbol('Gm') Unit_kpc= 1000*3.085e18*Unit_cm Unit_kpc.set_symbol('kpc') Unit_pc= 3.085e18*Unit_cm Unit_pc.set_symbol('pc') Unit_ua= 1.495978e11*Unit_m Unit_ua.set_symbol('ua') # other mass units Unit_g = 0.001*Unit_kg Unit_g.set_symbol('g') Unit_Ms = 1.9891e33*Unit_g Unit_Ms.set_symbol('Ms') Unit_Msol = 1.9891e33*Unit_g Unit_Msol.set_symbol('Ms') Unit_Mg = 2.23e11*Unit_Ms Unit_Mg.set_symbol('Mg') Unit_Mt = 5.9742e24*Unit_kg Unit_Mt.set_symbol('Mt') Unit_Mj = 317.893*Unit_Mt Unit_Mj.set_symbol('Mj') # other time units Unit_h = 3600*Unit_s Unit_h.set_symbol('h') Unit_yr = 31536000*Unit_s Unit_yr.set_symbol('yr') Unit_kyr= 1e3*Unit_yr Unit_kyr.set_symbol('kyr') Unit_Myr= 1e6*Unit_yr Unit_Myr.set_symbol('Myr') Unit_dy= 86400*Unit_s Unit_dy.set_symbol('days') Unit_hr= 3600*Unit_s Unit_hr.set_symbol('hr') # other speed units Unit_kmh= Unit_km/Unit_h Unit_kmh.set_symbol('kmh') Unit_kms= Unit_km/Unit_s Unit_kms.set_symbol('kms') # other units Unit_ms = Unit_m/Unit_s Unit_ms2 = Unit_ms**2 Unit_J = Unit_kg * Unit_ms2 Unit_J.set_symbol('J') Unit_erg = Unit_g * (Unit_cm/Unit_s)**2 Unit_erg.set_symbol('erg') Unit_N = Unit_kg * Unit_m/(Unit_s)**2 Unit_N.set_symbol('N') Unit_Pa = Unit_N/Unit_m**2 Unit_Pa.set_symbol('Pa') Unit_G = Unit_N*Unit_m**2/Unit_kg**2 Unit_d = Unit_kg/Unit_m**3 ################################################################################ # define some common unit systems ################################################################################ mks = UnitSystem('mks',[Unit_m, Unit_kg, Unit_s , Unit_K, Unit_mol, Unit_C]) cgs = UnitSystem('cgs',[Unit_cm, Unit_g, Unit_s , Unit_K, Unit_mol, Unit_C]) gal = UnitSystem('gal',[Unit_kpc,Unit_Mg, Unit_Myr, Unit_K, Unit_mol, Unit_C]) ''' constantes are now defined in ctes.py ################################################## # define some const. in mks ################################################## slight = PhysCte(2.99792458e8,Unit_m/Unit_s) G = PhysCte(6.6732e-11,Unit_G) q_electron = PhysCte(1.6022e-19,Unit_C) planck = PhysCte(6.6262e-34,Unit_J*Unit_s) boltzmann = PhysCte(1.3807e-23,Unit_J/Unit_K) m_electron = PhysCte(9.1095e-31,Unit_kg) m_proton = PhysCte(1.6726e-27,Unit_kg) m_neutron = PhysCte(1.6750e-27,Unit_kg) Na = PhysCte(6.0220e+23,Unit_mol) R = PhysCte(8.3144e+00,Unit_J/Unit_mol) Av = PhysCte(6.828e-50 ,Unit_Pa*Unit_m**6) Bv = PhysCte(4.419e-29,Unit_m**3) ''' ################################# def Set_SystemUnits_From_Params(params): ################################# ''' return a system of units from given parameters params is a dictionary that must constains at least params['UnitVelocity_in_cm_per_s'] params['UnitMass_in_g'] params['UnitLength_in_cm'] ''' UnitVelocity_in_cm_per_s = params['UnitVelocity_in_cm_per_s'] UnitMass_in_g = params['UnitMass_in_g'] UnitLength_in_cm = params['UnitLength_in_cm'] UnitTime_in_s = UnitLength_in_cm / UnitVelocity_in_cm_per_s # now from the params, define a system of units Unit_lenght = Unit_cm * UnitLength_in_cm Unit_mass = Unit_g * UnitMass_in_g Unit_time = Unit_s * UnitTime_in_s localsystem = UnitSystem('local',[Unit_lenght, Unit_mass, Unit_time , Unit_K, Unit_mol, Unit_C]) return localsystem