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