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.