Page MenuHomec4science
Files F74457143

__init__.py
No OneTemporary
Actions

File Metadata

Created
Sat, Jul 27, 22:35

__init__.py
View Options

'''
Some usefull function used to plot
using matplotlib (pylab)
version 4.0
'''
from numpy import *
from pylab import *
import matplotlib.colors as colors
import string
import sys
import os
import io
try:
from scipy import interpolate
is_scipy = True
except ImportError:
is_scipy = False
####################################################################################################################################
#
# Init
#
####################################################################################################################################
font = {'fontname' : 'Courier',
'color' : 'r',
'fontweight' : 'bold',
'fontsize' : 11}
rc('text', usetex=True)
rc('xtick.major', size=8)
rc('xtick.minor', size=4)
rc('xtick.major', pad=6)
rc('ytick.major', size=8)
rc('ytick.minor', size=4)
rc('ytick.major', pad=6)
rc('figure', facecolor='w')
####################################################################################################################################
#
# some usefull functions
#
####################################################################################################################################
def CleanString(s):
if type(s)!=str:
return s
else:
return string.replace(s, "_", "")
####################################################################################################################################
#
# some usefull objects
#
####################################################################################################################################
class DataPoints():
def __init__(self,x,y,yerr=None,color='k',linestyle='-',pointmarker='.',label='',tpe='line',linewidth=1):
self.x = x
self.y = y
self.yerr = yerr
self.x0 = x
self.y0 = y
self.color = color
self.linestyle = linestyle
self.linewidth = linewidth
self.pointmarker = pointmarker
self.label = label
self.tpe = tpe
def reduc(self,rf,n=None,mn=None,mx=None,dx=None):
if dx!=None:
n = int((max(self.x)-min(self.x))/dx)
yh,xh=histogram(self.x,n,(min(self.x),max(self.x)),weights=self.y)
else:
if mn==None and mx==None and n==None:
yh,xh=histogram(self.x,int(len(self.x)/float(rf)),(min(self.x),max(self.x)),weights=self.y)
else:
yh,xh=histogram(self.x,n,(mn,mx),weights=self.y)
xh = xh[:-1]
self.x = xh
self.y = yh
def integrate(self):
self.y = add.accumulate(self.y)
def interpolate(self,xi):
if not is_scipy:
raise "module scipy needed for function interpolate in DataPoints !"
n = len(self.x)
s1 = n-sqrt(2*n)
s2 = n+sqrt(2*n)
tck = interpolate.fitpack.splrep(self.x,self.y,s=s2,k=2)
yi = interpolate.fitpack.splev(xi,tck)
self.xi = xi
self.yi = yi
def derive(self):
dx = self.x[1:]-self.x[0:-1]
dy = self.y[1:]-self.y[0:-1]
self.x = self.x[1:]
self.y = dy/dx
c = isfinite(self.y)
self.x = compress(c,self.x)
self.y = compress(c,self.y)
def xy(self,):
return self.x,self.y
def get_x(self,):
return self.x
def get_y(self,):
return self.y
class Colors():
'''
Return the next color
'''
def __init__(self,palette=None,n=256,clist=None):
'''
Initialize with two modes :
1) give number + eventually palette name
2) give the list of colors
'''
if clist != None:
self.ls = clist
else:
if palette==None:
palette = GetPalette()
if type(palette)==str:
palette = GetPalette(palette)
self.ls = []
# set colors
for i in xrange(n):
f = i*255./n
self.ls.append(SetColor(f,palette))
self.i = 0
def set(self,i):
try:
self.i = self.ls.index(i)
except:
pass
def current(self):
return self.ls[self.i]
def next(self):
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return self.ls[self.i]
def get(self):
ls = self.ls[self.i]
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return ls
class LineStyles():
'''
Return the next linestyle
'''
def __init__(self,n=4):
self.ls = ['-','--','-.',':']
self.i = 0
def set(self,i):
try:
self.i = self.ls.index(i)
except:
pass
def current(self):
return self.ls[self.i]
def next(self):
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return self.ls[self.i]
def get(self):
ls = self.ls[self.i]
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return ls
class PointMakers():
'''
Return the next point marker
'''
def __init__(self):
self.ls = ['.',',','o','v','^','<','>','1','2','3','4','s','p','*','h','H','+','x','D','d','|','_']
self.i = 0
def set(self,i):
try:
self.i = self.ls.index(i)
except:
pass
def current(self):
return self.ls[self.i]
def next(self):
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return self.ls[self.i]
def get(self):
ls = self.ls[self.i]
self.i = self.i+1
if self.i==len(self.ls):
self.i = 0
return ls
def plot_data(data):
if data.yerr is None:
plot(data.x,data.y,color=data.color)
else:
errorbar(data.x,data.y,data.yerr,color=data.color)
####################################################################################################################################
#
# OPTIONS FACILITIES
#
####################################################################################################################################
#################################
def check_files_number(args):
#################################
if len(args) == 0:
raise "you must specify at least a filename"
#################################
def add_postscript_options(parser):
#################################
'''
This function allow to add postscript options to a parser object
'''
parser.add_option("-p",
action="store",
dest="ps",
type="string",
default = None,
help="postscript filename",
metavar=" FILE")
return parser
#################################
def add_ftype_options(parser):
#################################
'''
This function allow to add color options to a parser object
'''
parser.add_option("-t",
action="store",
dest="ftype",
type="string",
default = 'gadget',
help="type of the file",
metavar=" TYPE")
return parser
#################################
def add_reduc_options(parser):
#################################
'''
This function allow to reduc the number of particles
'''
parser.add_option("--reduc",
action="store",
dest="reduc",
type="int",
default=None,
help="reduc from a factor n")
return parser
#################################
def do_reduc_options(nb,opt,verbose=True):
#################################
if type(opt.reduc) == int:
if verbose:
print "reducing %s"%(opt.reduc)
nb = nb.reduc(opt.reduc)
return nb
#################################
def add_center_options(parser):
#################################
'''
This function allow to center the model
'''
parser.add_option("--center",
action="store",
dest="center",
type="string",
default=None,
help="center the model (histocenter,hdcenter)")
return parser
#################################
def do_center_options(nb,opt,verbose=True):
#################################
# center the model
if opt.center=='hdcenter':
if verbose:
print "centering %s"%(opt.center)
nb.hdcenter()
elif opt.center=='histocenter':
if verbose:
print "centering %s"%(opt.center)
nb.histocenter()
elif opt.center=='cmcenter':
if verbose:
print "centering %s"%(opt.center)
nb.cmcenter()
elif opt.center=='rebox':
if verbose:
print "centering %s"%(opt.center)
nb.rebox(mode='centred')
return nb
#################################
def add_select_options(parser):
#################################
'''
This function allow to select particles from the model
'''
parser.add_option("--select",
action="store",
dest="select",
type="string",
default=None,
help="select particles from the model ('gas','sph','sticky',...)")
return parser
#################################
def do_select_options(nb,opt,verbose=True):
#################################
# select
if opt.select!=None:
if verbose:
print "select %s"%(opt.select)
nb = nb.select(opt.select)
return nb
#################################
def add_info_options(parser):
#################################
'''
This function allow to select particles from the model
'''
parser.add_option("--info",
action="store_true",
dest="info",
default=False,
help="print info on the model")
return parser
#################################
def do_info_options(nb,opt,verbose=True):
#################################
# select
if opt.info:
nb.info()
return nb
#################################
def add_cmd_options(parser):
#################################
'''
This function allow to execute a command on the model
'''
parser.add_option("--cmd",
action="store",
dest="cmd",
type="string",
default=None,
help="python command 'nb = nb.selectc((nb.T()>10))'")
return parser
#################################
def do_cmd_options(nb,opt,verbose=True):
#################################
# cmd
if (opt.cmd!= None) and (opt.cmd!= 'None'):
if verbose:
print "exec : %s"%opt.cmd
exec(opt.cmd)
return nb
#################################
def add_display_options(parser):
#################################
'''
This function allow to display the model
'''
parser.add_option("--display",
action="store",
dest="display",
type="string",
default=None,
help="display the model")
return parser
#################################
def do_display_options(nb,opt,verbose=True):
#################################
# display
if (opt.display!= None) and (opt.display!= 'None'):
nb.display(obs=None,view=opt.display,marker='cross',pt=None,xp=None)
return nb
#################################
def add_limits_options(parser,xmin=None,xmax=None,ymin=None,ymax=None,zmin=None,zmax=None,wmin=None,wmax=None):
#################################
'''
This function allow to add limits options to a parser object
'''
parser.add_option("--xmin",
action="store",
dest="xmin",
type="float",
default=xmin,
help="min value in x")
parser.add_option("--xmax",
action="store",
dest="xmax",
type="float",
default=xmax,
help="max value in x")
parser.add_option("--ymin",
action="store",
dest="ymin",
type="float",
default=ymin,
help="min value in y")
parser.add_option("--ymax",
action="store",
dest="ymax",
type="float",
default=ymax,
help="max value in y")
parser.add_option("--zmin",
action="store",
dest="zmin",
type="float",
default=zmin,
help="min value in z")
parser.add_option("--zmax",
action="store",
dest="zmax",
type="float",
default=zmax,
help="max value in z")
parser.add_option("--wmin",
action="store",
dest="wmin",
type="float",
default=wmin,
help="min value in w")
parser.add_option("--wmax",
action="store",
dest="wmax",
type="float",
default=wmax,
help="max value in w")
return parser
#################################
def add_color_options(parser):
#################################
'''
This function allow to add color options to a parser object
'''
parser.add_option("-c",
action="store",
dest="colors",
type="string",
default = None,
help="colors",
metavar=" 0,64,192")
return parser
#################################
def add_labels_options(parser):
#################################
'''
This function allow to add labels options to a parser object
'''
parser.add_option("--labels",
action="store",
dest="labels",
type="string",
default = None,
help="labels",
metavar=''' "['name1','name2',...]" ''')
parser.add_option("--labx",
action="store",
dest="labx",
type="float",
default = 0.7,
help="box labels position x",
metavar="FLOAT")
parser.add_option("--laby",
action="store",
dest="laby",
type="float",
default = 0.7,
help="box labels position y",
metavar="FLOAT")
parser.add_option("--labdx",
action="store",
dest="labdx",
type="float",
default = 0.85,
help="box labels width : width = (1-labx)*labdx",
metavar="FLOAT")
parser.add_option("--labdy",
action="store",
dest="labdy",
type="float",
default = 0.85,
help="box labels height : height = (1-laby)*labdy",
metavar="FLOAT")
parser.add_option("--labex",
action="store",
dest="labex",
type="float",
default = 0.1,
help="labels offset in x with respect to the label box dxlab = dxbox*labex",
metavar="FLOAT")
parser.add_option("--labbox",
action="store",
dest="labbox",
type="int",
default = 1,
help="draw labels box or not",
metavar="INT")
return parser
#################################
def add_units_options(parser):
#################################
'''
This function allow to add postscript options to a parser object
'''
parser.add_option("--UnitLength_in_cm",
action="store",
dest="UnitLength_in_cm",
type="float",
default=None,
help="UnitLength in cm")
parser.add_option("--UnitMass_in_g",
action="store",
dest="UnitMass_in_g",
type="float",
default=None,
help="UnitMass in g")
parser.add_option("--UnitVelocity_in_cm_per_s",
action="store",
dest="UnitVelocity_in_cm_per_s",
type="float",
default=None,
help="UnitVelocity in cm per s")
parser.add_option("--param",
action="store",
dest="GadgetParameterFile",
type="string",
default = None,
help="Gadget parameter file",
metavar=" FILE")
return parser
#################################
def add_log_options(parser):
#################################
'''
This function allow to add postscript options to a parser object
'''
parser.add_option("--log",
action="store",
dest="log",
type="string",
default=None,
help="Use log to plot values (x,y,xy)")
return parser
#################################
def add_histogram_options(parser):
#################################
'''
This function allow to add histogram options to a parser object
'''
parser.add_option("--histogram",
action="store",
dest="histogram",
type="string",
default = 'none',
help="add histogram (none,add,only)")
return parser
#################################
def add_legend_options(parser):
#################################
'''
This function allow to add legend options to a parser object
'''
parser.add_option("--legend",
action="store_true",
dest="legend",
default = False,
help="add a legend")
parser.add_option("--legend_txt",
action="store",
dest="legend_txt",
type="string",
default = None,
help="legend text",
metavar=" LIST of STRINGS")
parser.add_option("--legend_loc",
action="store",
type="string",
dest="legend_loc",
default = None,
help="legend location 'upper right'... ")
return parser
#################################
def do_legend_options(opt,verbose=True,c=','):
#################################
if opt.legend_txt!=None:
opt.legend_txt = string.split(opt.legend_txt,c)
opt.legend=True
return opt
####################################################################################################################################
#
# GRAPH FACILITIES
#
####################################################################################################################################
#################################
def FindLimits(xmin,xmax,x,log,c=None):
#################################
if c!=None:
x = compress(c,x)
if log!=None:
if string.find(log,'x') != -1:
x = log10(x)
if xmin!=None:
xmin=log10(xmin)
if xmax!=None:
xmax=log10(xmax)
if xmin == None:
xmin = min(x)
if xmax == None:
xmax = max(x)
if xmin==xmax:
xmin = xmin-0.05*xmin
xmax = xmax+0.05*xmax
else:
xmin = xmin-0.05*(xmax-xmin)
xmax = xmax+0.05*(xmax-xmin)
# compute condition vector
c = (x>=xmin)*(x<=xmax)
if log!=None:
if string.find(log,'x') != -1:
xmin = 10**xmin
xmax = 10**xmax
return xmin,xmax,c
#################################
def SetMultiLimits(x=None,y=None,z=None,w=None,xmin=None,xmax=None,ymin=None,ymax=None,zmin=None,zmax=None,wmin=None,wmax=None,log=None):
#################################
if x!=None:
xmin,xmax,c = FindLimits(xmin,xmax,x,log,c=None)
if y!=None:
ymin,ymax,c = FindLimits(ymin,ymax,y,log,c=c)
if z!=None:
zmin,zmax,c = FindLimits(zmin,zmax,z,log,c=c)
if w!=None:
wmin,wmax,c = FindLimits(wmin,wmax,z,log,c=c)
return xmin,xmax,ymin,ymax,zmin,zmax,wmin,wmax
#################################
def SetLimitsFromDataPoints(xmin,xmax,ymin,ymax,datas,log=None):
#################################
x = array([],float)
y = array([],float)
for data in datas:
x = concatenate((x,data.x))
y = concatenate((y,data.y))
return SetLimits(xmin,xmax,ymin,ymax,x,y,log=log)
#################################
def SetLimits(xmin,xmax,ymin,ymax,x,y,log=None):
#################################
if log!=None:
if string.find(log,'x')!=-1:
x,y = CleanVectorsForLogX(x,y)
x = log10(x)
if xmin!=None:
xmin=log10(xmin)
if xmax!=None:
xmax=log10(xmax)
#############################
# set x
if xmin == None:
xmin = min(x)
if xmax == None:
xmax = max(x)
if xmin==xmax:
xmin = xmin-0.05*xmin
xmax = xmax+0.05*xmax
else:
xmin = xmin-0.05*(xmax-xmin)
xmax = xmax+0.05*(xmax-xmin)
# cut y values based on x
c = (x>=xmin)*(x<=xmax)
y = compress(c,y)
if log!=None:
if string.find(log,'x')!=-1:
#if log=='x' or log=='xy':
xmin = 10**xmin
xmax = 10**xmax
#############################
# set y
if log!=None:
if string.find(log,'y')!=-1:
x,y = CleanVectorsForLogY(x,y)
y = log10(y)
if ymin!=None:
ymin=log10(ymin)
if ymax!=None:
ymax=log10(ymax)
if ymin == None:
ymin = min(y)
if ymax == None:
ymax = max(y)
if ymin==ymax:
ymin = ymin-0.05*ymin
ymax = ymax+0.05*ymax
else:
ymin = ymin-0.05*(ymax-ymin)
ymax = ymax+0.05*(ymax-ymin)
if log!=None:
if string.find(log,'y')!=-1:
#if log=='y' or log=='xy':
ymin = 10**ymin
ymax = 10**ymax
return xmin,xmax,ymin,ymax
#################################
def SetAxis(xmin,xmax,ymin,ymax,log=None):
#################################
if log!=None:
if string.find(log,'x')!=-1 and string.find(log,'y')!=-1:
loglog()
elif string.find(log,'x')!=-1:
semilogx()
else:
semilogy()
axis([xmin,xmax,ymin,ymax])
# set the
ax = gca()
if log==None:
log = 'None'
if string.find(log,'x')==-1:
ax.xaxis.set_major_locator(AutoLocator())
try:
x_major = ax.xaxis.get_majorticklocs()
dx_minor = (x_major[-1]-x_major[0])/(len(x_major)-1) /5.
ax.xaxis.set_minor_locator(MultipleLocator(dx_minor))
except:
print "SetAxis : you should uptdate matplotlib"
if string.find(log,'y')==-1:
ax.yaxis.set_major_locator(AutoLocator())
try:
y_major = ax.yaxis.get_majorticklocs()
dy_minor = (y_major[-1]-y_major[0])/(len(y_major)-1) /5.
ax.yaxis.set_minor_locator(MultipleLocator(dy_minor))
except:
print "SetAxis : you should uptdate matplotlib"
#################################
def SetColorsForFiles(files,palette,cols=None):
#################################
if files==None:
return None
# set colors
colors = {}
i = 0
for file in files:
if cols!=None:
#colors[file]=cols[i]
print "SetColorsForFiles : not implemented"
else:
f = i*255/len(files)
colors[file] = SetColor(f,palette)
i = i + 1
return colors
#################################
def SetColor(f,palette):
#################################
'''
from a float between 0 and 255, get the correspondig triplet
'''
# set colors
rgb = palette[clip(int(f),0,255)]
return (rgb[0],rgb[1],rgb[2])
#################################
def InitPlot(files=None,opt=None):
#################################
global pcolors
global labelfont
labelfont = 16
try:
figure(figsize=opt.size)
except:
figure(figsize=(8,6))
#ax = gca()
#setp(ax.get_xticklabels(),fontsize=labelfont)
#setp(ax.get_yticklabels(),fontsize=labelfont)
palette = GetPalette()
pcolors = SetColorsForFiles(files,palette)
#################################
def EndPlot(files=None,opt=None):
#################################
if opt==None:
show()
return
if opt.ps!=None:
savefig(opt.ps)
else:
show()
#################################
def LegendFromDataPoints(datas=None,loc=None):
#################################
if datas!=None:
tags = []
for data in datas:
print "-->",data.label
#tags.append(r"%s"%CleanString(data.label)) # remove _ in c_{star} for example...
tags.append("%s"%data.label)
print "<--",tags[-1]
legend( tags,loc=loc )
else:
legend(loc=loc )
####################################################################################################################################
#
# PALETTES FUNCTIONS
#
####################################################################################################################################
def GetPalette(name='rainbow4'):
"""list of available palettes:
aips0 backgr bgyrw blue blulut color green heat idl2 idl4
idl5 idl6 idl11 idl12 idl14 idl15 isophot light manycol
pastel rainbow rainbow1 rainbow2 rainbow3 rainbow4 ramp
random random1 random2 random3 random4 random5 random6
real red smooth smooth1 smooth2 smooth3 staircase stairs8
stairs9 standard
"""
_name = os.path.normpath((os.path.join(__path__[0], 'rgb_tables', name)))
_r = []
_g = []
_b = []
try:
_f = open(_name)
except:
return
_f.readline()
for i in range(255):
line = _f.readline()
line = string.split(line)
_r.append(float(line[0]))
_g.append(float(line[1]))
_b.append(float(line[2]))
del line
_r = array(_r)
_g = array(_g)
_b = array(_b)
_f.close()
return transpose(array([_r,_g,_b]))/255.
# end def palette()
def GetColormap(name='rainbow4',revesed=False):
"""return a matplolib color map from a palette
"""
LUTSIZE = rcParams['image.lut']
palette = GetPalette(name)
red = []
green = []
blue = []
if not revesed:
for i in range(len(palette)):
r,g,b= palette[i][0],palette[i][1],palette[i][2]
x = i/float(len(palette)-1)
red.append( (x, r, r))
green.append( (x, g, g))
blue.append( (x, b, b))
else:
for i in range(len(palette)):
r,g,b= palette[len(palette)-i-1][0],palette[len(palette)-i-1][1],palette[len(palette)-i-1][2]
x = i/float(len(palette)-1)
red.append( (x, r, r))
green.append( (x, g, g))
blue.append( (x, b, b))
cmapdata = {'red':red,'green':green,'blue':blue}
cmap = colors.LinearSegmentedColormap(name, cmapdata, LUTSIZE)
return cmap
####################################################################################################################################
#
# UNITS
#
####################################################################################################################################
#################################
def do_units_options(options):
#################################
'''
This function allow to extract units options from option object
'''
try:
UnitLength_in_cm = options.UnitLength_in_cm
except:
UnitLength_in_cm = None
try:
UnitMass_in_g = options.UnitMass_in_g
except:
UnitMass_in_g = None
try:
UnitVelocity_in_cm_per_s = options.UnitVelocity_in_cm_per_s
except:
UnitVelocity_in_cm_per_s = None
try:
GadgetParameterFile = options.GadgetParameterFile
except:
GadgetParameterFile = None
if GadgetParameterFile != None:
params = io.read_params(GadgetParameterFile)
else:
if UnitLength_in_cm==None or UnitMass_in_g==None or UnitVelocity_in_cm_per_s==None:
params = None
else:
params={}
params['UnitLength_in_cm'] =UnitLength_in_cm
params['UnitMass_in_g'] =UnitMass_in_g
params['UnitVelocity_in_cm_per_s'] =UnitVelocity_in_cm_per_s
return params
####################################################################################################################################
#
# OTHERS
#
####################################################################################################################################
def CleanVectors(x,y):
'''
remove bad values
'''
c = isfinite(x) * isfinite(y)
x = compress(c,x)
y = compress(c,y)
return x.astype(float),y.astype(float)
def CleanVectorsForLogX(x,y):
'''
remove negative values
'''
c = (x>0)
x = compress(c,x)
y = compress(c,y)
return x.astype(float),y.astype(float)
def CleanVectorsForLogY(x,y):
'''
remove negative values
'''
c = (y>0)
x = compress(c,x)
y = compress(c,y)
return x.astype(float),y.astype(float)
####################################################################################################################################
#
#
#
####################################################################################################################################

Event Timeline

c4science · Help