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,
                 freq=352.21,
                 mass=938.272,
                 Ib=0.0):
        # 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 = Ib
            self.freq = freq
            self._data = numpy.zeros((self.Np, 6))
            self.mass = mass

    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 append_many(self, array):
        '''
        Append a matrix of particle vectors.

        Matrix on form 6xN, where N is number of particles.
        Each row should hold [x,xp,y,yp,phi,E]
        Units m,rad, MeV
        '''
        import numpy

        self._data = numpy.append(self._data, array, 0)
        self.Np += len(array)

    def combine_dst(self, other):
        '''
        Appends the particles from another dst object to this one

        '''

        if self.mass != other.mass:
            raise ValueError("You are trying to add two distributions with differing mass: {} and {}".format( self.mass, other.mass))
        if self.freq != other.freq:
            raise ValueError("You are trying to add two distributions with differing freq: {} and {}".format( self.freq, other.freq))

        self.append_many(other._data)
        self.Ib = ( self.Ib*self.Np + other.Ib*other.Np ) / ( self.Np + other.Np)

    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)

        fout.write(pack('{}d'.format(len(data)),*data))

        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)