Skip to content
Snippets Groups Projects
TraceWin.py 6.58 KiB
Newer Older
Yngve Levinsen's avatar
Yngve Levinsen committed

class dst:
Yngve Levinsen's avatar
Yngve Levinsen committed
    '''
    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]
    '''
Yngve Levinsen's avatar
Yngve Levinsen committed
    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):
Yngve Levinsen's avatar
Yngve Levinsen committed
        # Thanks Emma!
Yngve Levinsen's avatar
Yngve Levinsen committed

        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)

Yngve Levinsen's avatar
Yngve Levinsen committed
        Footer=numpy.fromfile(fin, dtype=numpy.float64, count=1)
Yngve Levinsen's avatar
Yngve Levinsen committed
        self.mass=Footer[0]

    def __getitem__(self, key):
Yngve Levinsen's avatar
Yngve Levinsen committed
        # makes the class function as a dictionary
        # e.g. dst['x'] returns the x array..
Yngve Levinsen's avatar
Yngve Levinsen committed
        if key in self._columns:
            i=self._columns.index(key)
            return self._data[:,i]
Yngve Levinsen's avatar
Yngve Levinsen committed

class density_file:
    '''
    Simple class to read a TraceWin density file
    into a pythonized object
    '''

    def __init__(self, filename):
        import numpy, sys
Yngve Levinsen's avatar
Yngve Levinsen committed


        self.filename=filename
        self.fin=file(self.filename, 'r')

        # 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
Yngve Levinsen's avatar
Yngve Levinsen committed
        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._max=numpy.zeros((counter,7))
        self._min=numpy.zeros((counter,7))

        self.lost=numpy.zeros((counter,self.Nrun))
        self.powlost=numpy.zeros((counter,self.Nrun))


        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]
        if shift:
            raise ValueError("ERROR, shifted "+str(shift*2)+" bytes")

        self.vlong=numpy.fromfile(self.fin, dtype=numpy.int16, count=1)[0]
        self.Nrun=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]

        self.version=version
        self.year=year

    def _skipAndCount(self):

        import numpy

        self._getHeader()

        numpy.fromfile(self.fin, dtype=numpy.int16, count=15132/2)
Yngve Levinsen's avatar
Yngve Levinsen committed


    def _getFullContent(self):

        import numpy

        #self._getHeader()
        # no need to read the header again:
        # (though only if we are SURE about content!)
Yngve Levinsen's avatar
Yngve Levinsen committed
        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

        moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
        moy2=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:
            rms_emit=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)
            rms_emit2=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)
Yngve Levinsen's avatar
Yngve Levinsen committed

        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]
            lost2=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            Minlost=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            Maxlost=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
            powlost2=numpy.fromfile(self.fin, dtype=numpy.float64, count=1)[0]
            Minpowlost=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
            Maxpowlost=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
Yngve Levinsen's avatar
Yngve Levinsen committed
            if self.vlong==1:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint64, count=n*step)
Yngve Levinsen's avatar
Yngve Levinsen committed
            else:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint32, count=n*step)
Yngve Levinsen's avatar
Yngve Levinsen committed
    
            if self.ib[self.i]>0:
                tabp=numpy.fromfile(self.fin, dtype=numpy.uint32, count=3*step)