Page MenuHomec4science
Files F74489798

Mkgmov.py
No OneTemporary
Actions

File Metadata

Created
Sun, Jul 28, 02:31

Mkgmov.py
View Options

import os
from pNbody import *
from pNbody.param import Params
import copy
import types
import gzip
import Mtools
import Mtools as mt
#######################################################################
# some usefull functions
#######################################################################
def ReadNbodyParameters(paramname):
'''
read param from a parameter Nbody file
'''
gparams = Params(paramname,None)
param = {}
# create new params
for p in gparams.params:
param[p[0]]=p[3]
return param
def gzip_compress(file):
f = open(file,'r')
content = f.read()
f.close()
f = gzip.open(file+'.gz', 'wb')
f.write(content)
f.close()
#######################################################################
#
# C L A S S D E F I N I T I O N
#
#######################################################################
class Movie():
def __init__(self,parameterfile='filmparam.py',format=None,imdir=None,timesteps=None,pio=False,compress=True):
self.DEFAULT_TIMESTEPS = None
self.DEFAULT_FORMAT = "fits"
self.DEFAULT_IMDIR = "fits"
self.DEFAULT_PIO = False
self.DEFAULT_COMPRESS = True
self.DEFAULT_SCALE = "log"
self.DEFAULT_CD = 0.
self.DEFAULT_MN = 0.
self.DEFAULT_MX = 0.
self.DEFAULT_PALETTE = "light"
self.parameterfile=parameterfile
# read the parameter file
self.read_parameterfile()
# use options
if format != None:
self.film['format'] = format
if imdir != None:
self.film['imdir'] = imdir
if timesteps != None:
self.film['timesteps'] = timesteps
if pio != None:
self.film['pio'] = pio
if compress != None:
self.film['compress'] = compress
self.imdir = self.film['imdir']
self.pio = self.film['pio']
self.compress = self.film['compress']
# deal with timesteps
self.set_timesteps()
self.ifile=0
if mpi.mpi_IsMaster():
if self.pio:
self.pio = "yes"
else:
self.pio = "no"
if self.parameterfile==None:
print "you must specify a parameter file"
sys.exit()
if not os.path.exists(self.parameterfile):
print "file %s does not exists"%self.parameterfile
sys.exit()
if not os.path.exists(self.imdir):
os.mkdir(self.imdir)
def info(self):
print "INFO INFO INFO"
#print self.film
print self.parameterfile
print self.getftype()
print "INFO INFO INFO"
def read_parameterfile(self):
if not os.path.isfile(self.parameterfile):
raise IOError(915,'file %s not found ! Pease check the file name.'%(self.parameterfile))
filmparam = __import__(os.path.splitext(self.parameterfile)[0],globals(), locals(), [], -1)
self.film = filmparam.film
# set some defaults
if not self.film.has_key('timesteps'):
self.film['timesteps'] = self.DEFAULT_TIMESTEPS
if not self.film.has_key('imdir'):
self.film['imdir'] = self.DEFAULT_IMDIR
if not self.film.has_key('format'):
self.film['format'] = self.DEFAULT_FORMAT
if not self.film.has_key('pio'):
self.film['pio'] = self.DEFAULT_PIO
if not self.film.has_key('compress'):
self.film['compress'] = self.DEFAULT_COMPRESS
self.setftype(self.film['ftype'])
# post process
for i,frame in enumerate(self.film['frames']):
frame['id'] = i
# check
#for frame in self.film['frames']:
# print frame['id']
# for component in frame['components']:
# print " ",component['id']
##################################
# time steps stuffs
##################################
def set_timesteps(self):
"""
define self.times (which is a list) based
on the value contained in self.film['timesteps']
"""
# self.times
if self.film['timesteps']=='every':
self.times="every"
elif type(self.film['timesteps']) == types.StringType:
fname = self.film['timesteps']
if not os.path.isfile(fname):
raise IOError(916,'file %s not found ! Pease check the file name.'%(fname))
times = io.read_ascii(fname,[0])[0]
times = take( times,len(times)-1-arange(len(times))) # invert order
times = times.tolist()
self.times = times
elif type(self.film['timesteps']) == types.ListType:
self.times = self.film['timesteps']
elif type(self.film['timesteps']) == types.TupleType:
t0 = self.film['timesteps'][0]
t1 = self.film['timesteps'][1]
dt = self.film['timesteps'][2]
times = arange(t0,t1,dt)
times = take( times,len(times)-1-arange(len(times))) # invert order
times = times.tolist()
self.times = times
else:
self.times=[]
def set_next_time(self):
if self.times!="every":
if len(self.times)>0:
self.times.pop()
def get_next_time(self):
if self.times=="every":
return 0.0
if len(self.times)==0:
return None
else:
return self.times[-1]
def getftype(self):
return self.ftype
def setftype(self,ftype):
self.ftype = ftype
def ApplyFilmParam(self,nb,film):
# set time reference for this file
exec("nb.tnow = %s"%film['time'])
# exec1
if film.has_key('exec'):
if film['exec'] != None:
exec(film['exec'])
# macro
if film.has_key('macro'):
if film['macro'] != None:
execfile(film['macro'])
return nb
def ApplyFrameParam(self,nb,frame):
nbf = nb
# exec
if frame.has_key('exec'):
if frame['exec'] != None:
exec(frame['exec'])
# macro
if frame.has_key('macro'):
if frame['macro'] != None:
execfile(frame['macro'])
return nbf
def ApplyComponentParam(self,nbf,component):
if component['id'][0] == '@':
# here, all tasks must have an object containing all particles
# ok, but not in the right order !!!
nbfc = Nbody(componentid,self.getftype())
nbfc = nbfc.SendAllToAll()
nbfc.sort()
nbfc.componentid=component['id']#[1:]
elif component['id'][0] == '#':
nbfc = nbf
nbfc.componentid = component['id']#[1:]
else:
nbfc = nbf.select(component['id'])
nbfc.componentid = component['id']
# exec
if component['exec'] != None:
exec(component['exec'])
# macro
if component['macro'] != None:
execfile(component['macro'])
#print "------------------------"
#print min(nbfc.u),max(nbfc.u)
#print min(nbfc.rho),max(nbfc.rho)
#print min(nbfc.tpe),max(nbfc.tpe)
#print "temperature",min(nbfc.T()),max(nbfc.T())
#print nbfc.nbody
#print min(nbfc.rsp),max(nbfc.rsp)
#print "------------------------"
return nbfc
def dump(self,dict):
# exctract dict
atime = dict['atime']
pos = dict['pos']
# create nbody object
nb = Nbody(pos=pos,ftype='gadget')
nb.atime = atime
# add other arrays
if dict.has_key('vel'):
nb.vel = dict['vel']
if dict.has_key('num'):
nb.num = dict['num']
if dict.has_key('mass'):
nb.mass = dict['mass']
if dict.has_key('tpe'):
nb.tpe = dict['tpe']
if dict.has_key('u'):
nb.u = dict['u']
if dict.has_key('rho'):
nb.rho = dict['rho']
if dict.has_key('rsp'):
nb.rsp = dict['rsp']
if dict.has_key('metals'):
nb.metals = dict['metals']
#!!!
nb.flag_chimie_extraheader=1
nb.ChimieNelements = 5
nb.flag_metals = 5
nb.ChimieSolarAbundances={}
nb.ChimieSolarAbundances['Fe']=0.00176604
nb.init()
#print "################################"
#print "writing qq.dat"
#print "################################"
#nb.rename('qq.dat')
#nb.write()
self.dumpimage(nb=nb)
def dumpimage(self,nb=None,file=None):
# for each frame one can create an img
imgs = []
if file!=None:
#print "(task=%04d) reading "%(mpi.ThisTask),file,self.getftype(),self.pio
nb = Nbody(file,ftype=self.getftype(),pio=self.pio)
else:
if nb==None:
raise "you must specify at least a file or give an nbody object"
film = self.film
nb = self.ApplyFilmParam(nb,film)
for frame in film['frames']:
nbf = self.ApplyFrameParam(nb,frame)
if frame.has_key('ext_cmd'):
if len(frame['ext_cmd'])>0:
output="/tmp/tmp.png"
for cmd in frame['ext_cmd']:
exec(cmd)
img = Image.open(output)
imgs.append(img)
continue
# composition parameters
if frame.has_key('cargs'):
if len(frame['cargs'])!=0:
frame['compose']=True
datas = []
else:
frame['compose']=False
else:
frame['cargs']=[]
frame['compose']=False
for component in frame['components']:
print "------------------------"
print "component"
print "------------------------"
nbfc = self.ApplyComponentParam(nbf,component)
# find the observer position
# 1) from params
# 2) from pfile
# 3) from tdir
# and transform into parameter
if frame['tdir']!=None:
tfiles = glob.glob(os.path.join(frame['tdir'],"*"))
tfiles.sort()
bname = os.path.basename(file)
tfiles_for_this_file = []
for j in xrange(len(tfiles)):
tfile = "%s.%05d"%(os.path.basename(file),j)
tmp_tfile = os.path.join(frame['tdir'],tfile)
if os.path.exists(tmp_tfile):
tfiles_for_this_file.append(tmp_tfile)
elif frame['pfile']!=None:
if not os.path.isfile(frame['pfile']):
print "parameter file %s does not exists(1)..."%(frame['pfile'])
# read from pfile defined in frame
param = ReadNbodyParameters(frame['pfile'])
tfiles_for_this_file = [None]
else:
# take frame as parameter
param = copy.copy(frame)
tfiles_for_this_file = [None]
# loop over different oberver positions for this file
for iobs,tfile in enumerate(tfiles_for_this_file):
if tfile!=None:
param = ReadNbodyParameters(tfile)
# add parameters defined by user in the parameter file
for key in component.keys():
param[key] = component[key]
# set image shape using frame
param['shape'] = (frame['width'],frame['height'])
# compute map
mat = nbfc.CombiMap(param)
if mpi.mpi_IsMaster():
if frame['compose']:
datas.append(mat)
# save output
if self.film["format"]=="fits":
output = '%04d_%04d-%s-%06d.fits'%(self.ifile,frame['id'],component['id'],iobs)
output = os.path.join(self.imdir,output)
print nb.atime,output
if os.path.exists(output):
os.remove(output)
header = [('TIME',nb.tnow,'snapshot time')]
io.WriteFits(transpose(mat), output, extraHeader = header)
# compress
if self.compress:
gzip_compress(output)
os.remove(output)
elif self.film["format"]=="png":
output = '%04d_%04d-%s-%06d.png'%(self.ifile,frame['id'],nbfc.componentid,iobs)
output = os.path.join(self.imdir,output)
print nb.atime,output
# here, we should use component['scale'] ... not frame['scale'], no ?
if not frame.has_key('scale'):
frame['scale']=self.DEFAULT_SCALE
if not frame.has_key('cd'):
frame['cd']=self.DEFAULT_CD
if not frame.has_key('mn'):
frame['mn']=self.DEFAULT_MN
if not frame.has_key('mx'):
frame['mx']=self.DEFAULT_MX
if not frame.has_key('palette'):
frame['palette']=self.DEFAULT_PALETTE
matint,mn_opt,mx_opt,cd_opt = set_ranges(mat,scale=frame['scale'],cd=frame['cd'],mn=frame['mn'],mx=frame['mx'])
frame['mn'] = mn_opt
frame['mx'] = mx_opt
frame['cd'] = cd_opt
img = get_image(matint,palette_name=frame['palette'])
img.save(output)
print frame['mn'],frame['mx'],frame['cd']
# need to create an img
if component.has_key('to_img'):
if component['to_img']:
# get params
if not component.has_key('scale'):
component['scale']=self.DEFAULT_SCALE
if not component.has_key('cd'):
component['cd']=self.DEFAULT_CD
if not component.has_key('mn'):
component['mn']=self.DEFAULT_MN
if not component.has_key('mx'):
component['mx']=self.DEFAULT_MX
if not component.has_key('palette'):
component['palette']=self.DEFAULT_PALETTE
matint,mn_opt,mx_opt,cd_opt = set_ranges(mat,scale=component['scale'],cd=component['cd'],mn=component['mn'],mx=component['mx'])
img = get_image(matint,palette_name=component['palette'])
# here we can add img commands....
if component.has_key('img_cmd'):
if len(component['img_cmd'])>0:
for cmd in component['img_cmd']:
exec(cmd)
# append img to list
img.atime = nb.atime
imgs.append(img)
#######################
# compose components
#######################
if frame['compose']:
if mpi.mpi_IsMaster():
img = Mtools.fits_compose_colors_img(datas,frame['cargs'])
# save
#output = '%04d_%04d.png'%(self.ifile,frame['id'])
#output = os.path.join(self.imdir,output)
#img.save(output)
# append img to list
img.atime = nb.atime
imgs.append(img)
#######################
# compose frames
#######################
if film.has_key('img_cmd'):
if len(film['img_cmd'])>0:
for cmd in film['img_cmd']:
exec(cmd)
output = '%04d.png'%(self.ifile)
output = os.path.join(self.imdir,output)
img.save(output)
img.i=0
# increment counter
self.ifile+=1

Event Timeline

c4science · Help