Skip to content
Snippets Groups Projects
TraceWin.py 6.84 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


        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
        print 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._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=8344/2)
        numpy.fromfile(self.fin, dtype=numpy.int16, count=6788/2)


    def _getFullContent(self):

        import numpy

        #self._getHeader()
        # no need to read the header again:
        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=n)
            rms_emit2=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)

        self.Np[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
        
        #print self.i, self.Np[self.i] # Nrun, nelp, ib, self.z[self.i], Np

        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]
    
            if self.vlong==1:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint64, count=n)
            else:
                tab=numpy.fromfile(self.fin, dtype=numpy.uint32, count=n)
    
            if self.ib[self.i]>0:
                tabp=[]
                for j in xrange(n):
                    tabp.append(numpy.fromfile(self.fin, dtype=numpy.uint32, count=step))
    

        
        # skip jibberish
        if self.Np[self.i]:
            numpy.fromfile(self.fin, dtype=numpy.int16, count=2352/2)
            if self.Nrun==100:
                numpy.fromfile(self.fin, dtype=numpy.int16, count=5572/2)
        else:
            numpy.fromfile(self.fin, dtype=numpy.int16, count=8032/2)
            if self.Nrun==100:
                numpy.fromfile(self.fin, dtype=numpy.int16, count=6788/2)