-
Yngve Levinsen authoredYngve Levinsen authored
Code owners
Assign users and groups as approvers for specific file changes. Learn more.
TraceWin.py 10.44 KiB
class dst:
'''
Simple class to read in a
TraceWin distribution file
Class afterwards hold the following
dictionary items:
- x [cm]
- xp [rad]
- y [cm]
- yp [rad]
- phi [rad]
- E [MeV]
'''
def __init__(self, filename):
# easy storage..
self.filename=filename
# used to create dict behaviour..
self._columns=['x','xp','y','yp','phi','E']
# read in the file..
self._readBinaryFile()
def _readBinaryFile(self):
# Thanks Emma!
import numpy
fin=file(self.filename,'r')
# dummy, Np, Ib, freq, dummy
Header_type = numpy.dtype([
('dummy12', numpy.int16),
('Np', numpy.int32),
('Ib', numpy.float64),
('freq', numpy.float64),
('dummy3', numpy.int8)
])
Header=numpy.fromfile(fin, dtype=Header_type, count=1)
self.Np=Header['Np'][0]
self.Ib=Header['Ib'][0]
self.freq=Header['freq'][0]
Table=numpy.fromfile(fin, dtype=numpy.float64, count=self.Np*6)
self._data=Table.reshape(self.Np,6)
Footer=numpy.fromfile(fin, dtype=numpy.float64, count=1)
self.mass=Footer[0]
def __getitem__(self, key):
# makes the class function as a dictionary
# e.g. dst['x'] returns the x array..
if key in self._columns:
i=self._columns.index(key)
return self._data[:,i]
class density_file:
'''
Simple class to read a TraceWin density file
into a pythonized object
'''
def __init__(self, filename):
import numpy, sys
self.filename=filename
self.fin=file(self.filename, 'r')
# currently unknown:
self.version=0
# first we simply count how many elements we have:
counter=0
while True:
try:
self._skipAndCount()
counter+=1
except IndexError: # EOF reached..
break
if sys.flags.debug:
print "Number of steps found:", counter
self.fin.seek(0)
# set up the arrays..
self.i=0
# z position [m] :
self.z=numpy.zeros(counter)
# current [mA] :
self.ib=numpy.zeros(counter)
# number of lost particles:
self.Np=numpy.zeros(counter)
self.Xouv=numpy.zeros(counter)
self.Youv=numpy.zeros(counter)
self.moy=numpy.zeros((counter,7))
self.moy2=numpy.zeros((counter,7))
self._max=numpy.zeros((counter,7))
self._min=numpy.zeros((counter,7))
if self.version>=7:
self.rms_emit=numpy.zeros((counter,3))
self.rms_emit2=numpy.zeros((counter,3))
self.lost=numpy.zeros((counter,self.Nrun))
self.powlost=numpy.zeros((counter,self.Nrun))
self.lost2=numpy.zeros(counter)
self.Minlost=numpy.zeros(counter)
self.Maxlost=numpy.zeros(counter)
self.powlost2=numpy.zeros(counter)
self.Minpowlost=numpy.zeros(counter)
self.Maxpowlost=numpy.zeros(counter)
while self.i<counter:
self._getFullContent()
self.i+=1
def _getHeader(self):
import numpy
# header..
version=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
year=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
# there is much more data written, but it is undocumented. Our trick to get back "on line":
shift=0
while year!=2011 or version!=8:
shift+=1
version=year
year=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
self.vlong=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
self.Nrun=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
if shift:
raise ValueError("ERROR, shifted "+str(shift*2)+" bytes")
self.version=version
self.year=year
def _skipAndCount(self):
import numpy
self._getHeader()
if self.version==8:
numpy.fromfile(self.fin, dtype=numpy.int16, count=8344/2)
if self.Nrun>1:
#WARN not 100% sure if this is correct..
numpy.fromfile(self.fin, dtype=numpy.int16, count=((5588+self.Nrun*12)/2))
def _get_7dim_array(array):
return dict(x=array[0],
y=array[1],
phase=array[2],
energy=array[3],
r=array[4],
z=array[5],
dpp=array[6],
)
def _getFullContent(self):
import numpy
#self._getHeader()
# no need to read the header again:
# (though only if we are SURE about content!)
numpy.fromfile(self.fin, dtype=numpy.int16, count=5)
nelp=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
self.ib[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
self.z[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
# Aperture
self.Xouv[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
self.Youv[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
step=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
n=7 # x [m], y[m], Phase [deg], Energy [MeV], R[m], Z[m], dp/p
self.moy[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
self.moy2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
self._max[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
self._min[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
if self.version>=5:
rms=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
rms_size2=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
if self.version>=6:
min_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
max_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
if self.version>=7:
self.rms_emit[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
self.rms_emit2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
if self.version>=8:
e_ouv=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
phase_ouv_pos=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
phase_ouv_neg=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
self.Np[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
if self.Np[self.i]:
powlost=numpy.zeros(self.Nrun)
for i in xrange(self.Nrun):
self.lost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
self.powlost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
self.lost2[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
self.Minlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
self.Maxlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
self.powlost2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float64, count=1)[0]
self.Minpowlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
self.Maxpowlost[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
if self.vlong==1:
tab=numpy.fromfile(self.fin, dtype=numpy.uint64, count=n*step)
else:
tab=numpy.fromfile(self.fin, dtype=numpy.uint32, count=n*step)
if self.ib[self.i]>0:
tabp=numpy.fromfile(self.fin, dtype=numpy.uint32, count=3*step)
class remote_data_merger:
def __init__(self, base='.'):
self._base=base
self._files=[]
def add_file(self,filepath):
import os
if os.path.exists(filepath):
fname=filepath
else:
fullpath=os.path.join(self._base,filepath)
if os.path.exists(fullpath):
fname=fullpath
else:
raise ValueError("Could not find file "+filepath)
if fname not in self._files:
self._files.append(fname)
def generate_partran_out(self,filename=None):
'''
Creates a string to be written to file
each line is a list.
If filename is given, writes directly to output file.
'''
import numpy as np
h1=[]
h2=[]
d1=[]
d2=[]
d3=[]
if self._files:
for f in self._files:
string=file(f,'r').read()
split=string.split('$$$')
if split[9]!='Data_Error':
raise ValueError("Magic problem, please complain to Yngve")
thisdata=split[10].strip().split('\n')
if not h1:
h1=[thisdata[0]+" (std in paranthesis)"]
h2=thisdata[2:10]
d1.append(thisdata[1].split())
d2.append(thisdata[10])
d3.append(thisdata[11])
# fix d1:
for i in xrange(len(d1)):
for j in xrange(len(d1[0])):
d1[i][j]=float(d1[i][j])
d1=np.array(d1)
means=d1.mean(axis=0)
stds=d1.std(axis=0)
d1=[]
for i in xrange(len(stds)):
if stds[i]/means[i]<1e-10:
stds[i]=0.0
for i in xrange(len(stds)):
# some small std are removed..
if stds[i]/means[i]>1e-8:
d1.append('%f(%f)' % (means[i],stds[i]))
else: #error is 0
d1.append(str(means[i]))
d1=[' '.join(d1)]
# create data:
data=h1+d1+h2+d2+d3
if filename:
file(filename,'w').write('\n'.join(data))
return data
class partran(dict):
'''
Read partran1.out files..
'''
def __init__(self,filename):
self.filename=filename
self._readAsciiFile()
def _readAsciiFile(self):
import numpy
stream=file(self.filename,'r')
for i in xrange(10):
line=stream.readline()
self.columns=['NUM']+line.split()[1:]
self.data=numpy.loadtxt(stream)
self._dict={}
for i in xrange(len(self.columns)):
self[self.columns[i]]=self.data[:,i]