TraceWin.py

from __future__ import print_function


class dst:
    '''
    Simple class to read in a
    TraceWin distribution file

    Class afterwards hold the following
    dictionary items:
      - x [m]
      - xp [rad]
      - y [m]
      - yp [rad]
      - phi [rad]
      - E [MeV] (kinetic energy)
    '''
    def __init__(self, filename=None):
        # easy storage..
        self.filename = filename
        # used to create dict behaviour..
        self._columns = ['x', 'xp', 'y', 'yp', 'phi', 'E']
        if filename:
            # read in the file..
            self._readBinaryFile()
        else:
            import numpy
            self.Np = 0
            self.Ib = 0.0
            self.freq = 352.21
            self._data = numpy.zeros((self.Np, 6))
            self.mass = 0.0

    def append(self, x=0.0, xp=0.0, y=0.0, yp=0.0, E=0.0, phi=0.0):
        '''
        Append one particle to the distribution

        - Kinetic Energy in MeV
        - x,y in m
        - xp,yp in rad
        - phi in rad
        '''
        import numpy

        self._data = numpy.append(self._data, [[x, xp, y, yp, phi, E]], 0)
        self.Np += 1

    def remove(self, i=None):
        '''
        Removes all particles from the distribution, or the line specified by i
        '''
        import numpy

        if i is None:
            self._data=numpy.delete(self._data,numpy.s_[:], 0)
            self.Np=0
        else:
            self._data=numpy.delete(self._data, i, 0)
            self.Np-=1

    def _readBinaryFile(self):
        # Thanks Ema!

        import numpy

        fin=open(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]

        # Some toutatis distributions has an undocumented 7th line of 0's
        Table = numpy.fromfile(fin, dtype=numpy.float64, count=self.Np*7+1)
        if len(Table)==self.Np*7+1:
            self._data = Table[:-1].reshape(self.Np, 7)
        elif len(Table)==self.Np*6+1: # this is true in most cases
            self._data = Table[:-1].reshape(self.Np, 6)
        else:
            raise ValueError("Incorrect table dimensions found:", len(Table))

        # convert x,y from cm to m:
        self._data[:,0] *= 1e-2
        self._data[:,2] *= 1e-2

        self.mass = Table[-1]

    def keys(self):
        return self._columns[:]

    def __getitem__(self, key):
        # makes the class function as a dictionary
        # e.g. dst['x'] returns the x array..
        try:
            i=self._columns.index(key)
            return self._data[:, i]
        except:
            raise ValueError("Available keys: "+str(self._columns))

    def __setitem__(self, key, value):
        try:
            i=self._columns.index(key)
            self._data[:,i] = value
        except:
            raise ValueError("Available keys: "+str(self._columns))
    
    def save(self, filename, toutatis=False):
        '''
        Save the distribution file
        so it can be read by TraceWin again

        :param filename: Name of file
        :param toutatis: Include 7th column of zeros

        Stolen from Ryoichi's func.py (with permission)
        '''

        from struct import pack
        import numpy

        fout=open(filename, 'wb')
        fout.write(pack('b', 125))
        fout.write(pack('b', 100))
        fout.write(pack('i', self.Np))
        fout.write(pack('d', self.Ib))
        fout.write(pack('d', self.freq))
        fout.write(pack('b', 125))

        data = self._data.copy()

        if toutatis and data.shape[1]==6:
            data = numpy.append(data,numpy.zeros((len(data),1)),1)
        elif not toutatis and data.shape[1]==7:
            data = data[:,:-1]

        # convert x,y from m to cm:
        data[:,0] *= 1e2
        data[:,2] *= 1e2

        if toutatis:
            data = data.reshape(self.Np*7, 1)
        else:
            data = data.reshape(self.Np*6, 1)

        for x in data:
            fout.write(pack('d', x))

        fout.write(pack('d', self.mass))
        fout.close()

    def subplot(self, index, x, y=None, nb=100, mask=None):
        '''
        Create a subplot histogram similar to TraceWin.

        Example::
            import numpy as np
            from ess import TraceWin
            from matplotlib import pyplot as plt
            data=TraceWin.dst('part_dtl1.dst')
            m=np.where(data['E']>3.5)
            data.subplot(221,'x','xp',mask=m)
            data.subplot(222,'y','yp',mask=m)
            data.subplot(223,'phi','E',mask=m)
            data.subplot(224,'x','y',mask=m)
            plt.show()

        '''
        from matplotlib.colors import LogNorm
        import matplotlib.pyplot as plt
        import numpy as np

        units={ 'x': 'mm', 'y': 'mm',
                'xp': 'mrad', 'yp': 'mrad',
                'E': 'MeV', 'phi': 'deg'
              }
        # get X and Y data
        dx=np.array(self[x])
        if mask!=None:
            dx = dx[mask]
        if y!=None:
            dy = np.array(self[y])
            if mask!=None:
                dy=dy[mask]

        if x in ['x','y','xp','yp']:
            dx*=1e3
        if y in ['x','y','xp','yp']:
            dy*=1e3
        if x in ['phi']:
            dx-=np.average(dx)
            dx*=180/np.pi
        if y in ['phi']: