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

class dst:
    def __init__(self, filename):
        '''
        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]
        '''
        # 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='f64', 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]