Page MenuHomec4science
Files F67753375

main.py
No OneTemporary
Actions

File Metadata

Created
Mon, Jun 24, 04:06

main.py
View Options

# -*- 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:
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):
"""
This is the constructor for the \texttt{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 function 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
"""
#################################
# init vars
#################################
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)
if self.vel == None:
self.vel = zeros((self.nbody,3),float32)
self.vel = self.vel.astype(float32)
#pass
if self.mass == None:
self.mass = ones((self.nbody, ),float32)/self.nbody
self.mass = self.mass.astype(float32)
#pass
if self.tpe == None:
self.tpe = zeros(self.nbody,int)
self.tpe = self.tpe.astype(int)
#pass
if self.num == None:
self.num = self.get_num()
self.num = self.num.astype(int)
#pass
# 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_vect()
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 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_vect(self):
#################################
'''
return specific vector 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()
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)
print "UnitLength_in_cm =%g"%(params['UnitLength_in_cm'])
print "UnitVelocity_in_cm_per_s =%g"%(params['UnitVelocity_in_cm_per_s'])
print "UnitMass_in_g =%g"%(params['UnitMass_in_g'])
#################################
#
# 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 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 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()<r2)))
prof,z = histogram(ann.pos[:,2],z)
z1 = z[:-1]
z2 = z[1:]
ds = z2-z1
ds1 = pi*(r2**2-r1**2)
prof = prof/(ds*ds1)
z = z[:-1]
return z,mpi.mpi_allreduce(prof)
def sigma(self,r=None,nb=25.,rm=50.):
"""
Return the 3 velocity dispersion (in cylindrical coordinates) and the mean azimuthal velocity curve.
If r is not specified, it is computed with nb and rm.
The output is a list (r,sr,st,sz,mt) of 5 $n\times 1$ float arrays,
where r is the radius, sr the radial velocity dispersion, st, the azimuthal velocity dispersion,
sz, the vertical velocity dispersion and mt, the mean azimuthal velocity curve.
!!! This routine works only if particles have equal masses !!!
r : radius where to compute the values
nb : number of bins (size of the output)
rm : maximal radius
return : r,sr,st,sz,mt
"""
if r!= None:
r1 = r[:-1]
r2 = r[1:]
r = (r1+r2)/2.0
else:
rmax = rm
dr = rm/float(nb)
r1 = arange(0,rmax,dr)
r2 = arange(dr,rmax+dr,dr)
r = (r1+r2)/2.
sr = []
sz = []
st = []
mt = []
for i in range(len(r1)):
#print i,len(r1)
ann = self.selectc((self.rxy()>r1[i])*((self.rxy()<r2[i])))
x = ann.pos[:,0]
y = ann.pos[:,1]
vx = ann.vel[:,0]
vy = ann.vel[:,1]
vz = ann.vel[:,2]
xr = sqrt(x**2+y**2)
vr = (x*vx + y*vy) / xr
vt = (x*vy - y*vx) / xr
# stand dev.
if len(vr) > 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<ps[j+1])
if sum(c.astype(int))<=1:
break
zm.append(compress(c,z).mean())
else:
ok = 1
if not ok:
#print "less than 1 particle in the sub coupe",R[i]
amp = zeros(len(ps)/2).astype(float)
m1 = concatenate((m1,amp))
m2 = concatenate((m2,amp))
continue
ps = ps.astype(float)
zm = array(zm,float)
t = t.astype(float)
z = z.astype(float)
# fourier decomposition
f, amp, phi = fourier.fourier(ps,zm)
m1 = concatenate((m1,amp))
m2 = concatenate((m2,phi))
Rs = concatenate((Rs,array([(R[i]+R[i+1])/2.])))
m = f*2*pi
m1 = reshape(m1,(nr,nm))
m2 = reshape(m2,(nr,nm))
m1 = fliplr(m1,0)
m2 = fliplr(m2,0)
return Rs,m,m1,m2
def dmodes(self,nr=32,nm=16,rm=32):
"""
Compute the density 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
dm = []
ok = 0
for j in range(len(ps)-1):
c = (t>=ps[j])*(t<ps[j+1])
if sum(c.astype(int))<=1:
break
dm.append(sum(c.astype(int)))
else:
ok = 1
if not ok:
#print "less than 1 particle in the sub coupe",R[i]
amp = zeros(len(ps)/2).astype(float)
m1 = concatenate((m1,amp))
m2 = concatenate((m2,amp))
Rs = concatenate((Rs,array([(R[i]+R[i+1])/2.])))
continue
ps = ps.astype(float)
dm = array(dm,float)
t = t.astype(float)
z = z.astype(float)
# fourier decomposition
f, amp, phi = fourier.fourier(ps,dm)
phi = -phi
m1 = concatenate((m1,amp))
m2 = concatenate((m2,phi))
Rs = concatenate((Rs,array([(R[i]+R[i+1])/2.])))
'''
#an.rotate(pi,axis='y')
#an.show(view='xy',size=(20,20))
print (R[i]+R[i+1])/2.,phi[2]+2*pi
g = SM.plot()
g.erase()
g.limits(0,360,dm)
g.box()
g.connect(ps*180/pi,dm)
fdm = amp[2]*cos(ps*2-phi[2])
#fdm = amp[0]*cos(ps*0+phi[0])
#for j in range(1,16):
# fdm = fdm + amp[j]*cos(ps*j+phi[j])
g.ctype('red')
g.connect(ps*180/pi,fdm)
g.ctype('black')
g.show()
'''
m = f*2*pi
m1 = reshape(m1,(nr,nm))
m2 = reshape(m2,(nr,nm))
m1 = fliplr(m1,0)
m2 = fliplr(m2,0)
return Rs,m,m1,m2
def getRadiusInCylindricalGrid(self,z,Rmax,nr=32,nt=32):
'''
Compute the radius in cells of a cylindrical grid
'''
irs = arange(nr)
its = arange(nt)
Radi = zeros((nt,nr),float)
Rs = irs*Rmax/float(nr)
ts = its*2*pi/float(nt)
# use to compute values at the center of the cell
dr = Rs[1]-Rs[0]
dt = ts[1]-ts[0]
for ir in irs:
for it in its:
R = ir*Rmax/float(nr) + dr/2.
#t = it*2*pi/float(nt) + dt/2.
Radi[it,ir]= R
return Radi
def getAccelerationInCylindricalGrid(self,eps,z,Rmax,nr=32,nt=32,UseTree=False):
'''
Compute the Acceleration in cells of a cylindrical grid
'''
irs = arange(nr)
its = arange(nt)
Accx = zeros((nt,nr),float)
Accy = zeros((nt,nr),float)
Accz = zeros((nt,nr),float)
Rs = irs*Rmax/float(nr)
ts = its*2*pi/float(nt)
# use to compute values at the center of the cell
dr = Rs[1]-Rs[0]
dt = ts[1]-ts[0]
# build the tree
if UseTree:
if self.Tree==None:
self.Tree = self.getTree()
for ir in irs:
for it in its:
R = ir*Rmax/float(nr) + dr/2.
t = it*2*pi/float(nt) + dt/2.
x = R*cos(t)
y = R*sin(t)
z = 0.0
if UseTree:
a = self.TreeAccel(array([[x,y,z]],float32),eps,Tree=None)
Accx[it,ir]= a[0][0]
Accy[it,ir]= a[0][1]
Accz[it,ir]= a[0][2]
else:
a = self.Accel([x,y,z],eps)
Accx[it,ir]= a[0]
Accy[it,ir]= a[1]
Accz[it,ir]= a[2]
return Accx,Accy,Accz
def getPotentialInCylindricalGrid(self,eps,z,Rmax,nr=32,nt=32,UseTree=False):
'''
Compute the potential in cells of a cylindrical grid
'''
irs = arange(nr)
its = arange(nt)
Phis = zeros((nt,nr),float)
Rs = irs*Rmax/float(nr)
ts = its*2*pi/float(nt)
# build the tree
if UseTree:
if self.Tree==None:
self.Tree = self.getTree()
# use to compute values at the center of the cell
dr = Rs[1]-Rs[0]
dt = ts[1]-ts[0]
for ir in irs:
for it in its:
R = ir*Rmax/float(nr) + dr/2.
t = it*2*pi/float(nt) + dt/2.
x = R*cos(t)
y = R*sin(t)
z = 0.0
if UseTree:
P = self.TreePot(array([[x,y,z]],float32),eps,Tree=None)
Phis[it,ir]= P[0]
else:
Phis[it,ir]= self.Pot([x,y,z],eps)
return Phis
def getSurfaceDensityInCylindricalGrid(self,Rmax,nr=32,nt=32):
'''
Compute the surface density in cells of a cylindrical grid
'''
# r and t between 1 and 2
r = self.rxy()/Rmax
t = (self.phi_xy()+pi)/(2*pi)
pos = transpose(array([t,r,r],float32))
shape = (nt,nr)
Sdens = GetMassMap(pos,self.mass,shape)
# divide by the suface of a cell
Rs = arange(nr+1)*Rmax/float(nr)
R1 = Rs[:-1]
R2 = Rs[1:]
S = pi*(R2**2-R1**2)/float(nt)
return Sdens/S
def getNumberParticlesInCylindricalGrid(self,Rmax,nr=32,nt=32):
'''
Compute the number of particles in cells of a cylindrical grid
'''
# r and t between 1 and 2
r = self.rxy()/Rmax
t = (self.phi_xy()+pi)/(2*pi)
pos = transpose(array([t,r,r],float32))
shape = (nt,nr)
Num = GetMassMap(pos,ones(len(pos)).astype(float32),shape)
return Num
def getRadialVelocityDispersionInCylindricalGrid(self,Rmax,nr=32,nt=32):
'''
Compute the radial velocity dispersion in cells of a cylindrical grid
'''
# r and t between 1 and 2
r = self.rxy()/Rmax
t = (self.phi_xy()+pi)/(2*pi)
pos = transpose(array([t,r,r],float32))
shape = (nt,nr)
Sigmar = GetSigmaValMap(pos,self.mass,self.Vr(),shape)
return Sigmar
#################################
#
# geometrical operations
#
#################################
def cmcenter(self):
"""
Move the N-body object in order
to center the mass center at the origin.
"""
self.pos = (self.pos - self.cm()).astype(float32)
def cvcenter(self):
"""
Center the center of velocities at the origin.
"""
self.vel = (self.vel - self.cv()).astype(float32)
def histocenter(self,rbox=50,nb=500):
"""
Move the N-body object in order to center the higher
density region found near the mass center.
The higher density region is determined whith density histograms.
rbox : box dimension, where to compute the histograms
nb : number of bins for the histograms
"""
self.pos = (self.pos - self.get_histocenter(rbox=rbox,nb=nb)).astype(float32)
def histocenter2(self,rbox=50,nb=64):
"""
Move the N-body object in order to center the higher
density region found near the mass center.
The higher density region is determined whith density histograms.
rbox : box dimension, where to compute the histograms
nb : number of bins for the histograms
"""
self.pos = (self.pos - self.get_histocenter2(rbox=rbox,nb=nb)).astype(float32)
def hdcenter(self):
"""
Move the N-body object in order to center the higher
density region found.
"""
if self.has_array('rho'):
idx = mpi.mpi_argmax(self.rho)
self.pos = (self.pos - mpi.mpi_getval(self.pos,idx)).astype(float32)
def translate(self,dx,mode='p'):
"""
Translate the positions or the velocities of the object.
dx : shift (array vector)
mode : 'p' : translate positions
'v' : translate velocities
"""
if mode=='p':
self.pos = (self.pos + dx).astype(float32)
elif mode=='v':
self.vel = (self.vel + dx).astype(float32)
def rebox(self,boxsize=None,mode=None):
"""
Translate the positions of the object in order that all particles
being contained in a box of size boxsize.
boxsize : size of the box
if boxsize is not defined, we try first to see if self.boxsize
is defined.
"""
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])
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 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((r<h),wk,0)
return wk
i = self.getindex(num)
h = self.rsp[i]
r = sqrt( (self.pos[:,0]-self.pos[i,0])**2 + (self.pos[:,1]-self.pos[i,1])**2 + (self.pos[:,2]-self.pos[i,2])**2 )
# compute wk for these particle
wk = getwk(r,h)
NORM_COEFF=4.188790204786 # 4/3.*pi
return NORM_COEFF *sum(wk)*h**3
def real_numngb(self,num):
'''
number of particles wor wich r<h
'''
i = self.getindex(num)
h = self.rsp[i]
r = sqrt( (self.pos[:,0]-self.pos[i,0])**2 + (self.pos[:,1]-self.pos[i,1])**2 + (self.pos[:,2]-self.pos[i,2])**2 )
n = sum(where((r<h),1,0))
return n
def usual_numngb(self,num):
'''
usual way to compute the number of neighbors
'''
i = self.getindex(num)
h = self.rsp[i]
r = sqrt( (self.pos[:,0]-self.pos[i,0])**2 + (self.pos[:,1]-self.pos[i,1])**2 + (self.pos[:,2]-self.pos[i,2])**2 )
c = (r<h)
M1 = sum(c*self.mass)/sum(c.astype(int)) # mean mass
M2 = 4/3.*pi*h**3 * self.rho[i] # total mass
n = M2/M1
return n
#################################
#
# redistribution of particles
#
#################################
def redistribute(self):
"""
This function redistribute particles amoung all nodes in order to
have a similar number of particles per nodes
"""
if mpi.NTask>1:
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):
'''
does not work well
'''
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):
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):
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):
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)
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
####################################################################################################################################
#
# other classes
#
####################################################################################################################################
#if FORMATSFILE != None:
# execfile(FORMATSFILE)
#if USERFORMATSFILE != None:
# execfile(USERFORMATSFILE)
#if USERFORMATSDIR != None:
# formatsfiles = glob.glob(os.path.join(USERFORMATSDIR,'*.py'))
if FORMATSDIR != None:
formatsfiles = glob.glob(os.path.join(FORMATSDIR,'*.py'))
for format in formatsfiles:
execfile(format)
####################################################################################################################################
#
# NBODY REDIRECTOR
#
####################################################################################################################################
ELTS = dir()
def get_known_formats():
'''
return the name of known Nbody formats
'''
formats = []
for elt in ELTS:
if elt[:6]=='Nbody_':
formats.append(elt)
return formats
def Nbody(*arg,**kw):
"""
The aim of this function is simply to return to the right class
"""
# default value
nb = None
# gather all known classes
formats = get_known_formats()
# check ftype
if kw.has_key('ftype'):
if kw['ftype'] != None:
format = "Nbody_%s"%kw['ftype']
class_name = eval(format)
nb = class_name(*arg,**kw)
else:
format = 'Nbody_default'
class_name = eval(format)
nb = class_name(*arg,**kw)
if nb == None:
self.log.write( "you should give the right format ftype" )
sys.exit()
'''
# check ftype
if kw.has_key('ftype'):
if kw['ftype'] != None:
format = "Nbody_%s"%kw['ftype']
# check if format is known
try:
formats.index(format)
except ValueError:
print "%s format is unknown"%format
format = None
class_name = eval(format)
nb = class_name(*arg,**kw)
# if format exists, try to create the object
if format != None:
class_name = eval(format)
try:
nb = class_name(*arg,**kw)
return nb
except "ReadBlockError":
pass
except "ReadAsciiError":
pass
except "ReadError":
pass
except "IOError":
print "IOError"
# except:
# pass
print "not a %s format"%(format)
# test all known Nbody class
for format in formats:
#print "try to open as %s"%format
class_name = eval(format)
try:
nb = class_name(*arg,**kw)
break
except "ReadBlockError":
pass
except "ReadAsciiError":
pass
except "ReadError":
pass
except "IOError":
print "IOError"
# except:
# pass
if nb == None:
print "unreconized format"
sys.exit()
'''
return nb

Event Timeline

c4science · Help