TraceWin.py

            self._max = self._fun_merge(o, numpy.maximum, "_max")
            self._min = self._fun_merge(o, numpy.minimum, "_min")

            if self.version >= 5:
                # this looks strange to me, but it is what TraceWin does..
                self.rms_size = self._sum_merge(o, "rms_size")
                self.rms_size2 = self._sum_merge(o, "rms_size2")

            if self.version >= 6:
                self.max_pos_moy = self._fun_merge(o, numpy.maximum, "max_pos_moy")
                self.min_pos_moy = self._fun_merge(o, numpy.minimum, "min_pos_moy")

            if self.version >= 7:
                # this looks strange to me, but it is what TraceWin does..
                self.rms_emit = self._sum_merge(o, "rms_emit")
                self.rms_emit2 = self._sum_merge(o, "rms_emit2")

            if self.version >= 8:
                # Warning: TraceWin does NOT merge these data in any way
                self.energy_accept = self._avg_merge(o, "energy_accept")
                self.phase_ouv_pos = self._avg_merge(o, "phase_ouv_pos")
                self.phase_ouv_neg = self._avg_merge(o, "phase_ouv_neg")

            # Note, we don't get into the problem of differing table sizes
            # particles are lost, because we have written zeroes for
            # the rest of the tables

            self.lost = self._concatenate_merge(o, "lost")
            self.powlost = self._concatenate_merge(o, "powlost")

            self.lost2 = self._sum_merge(o, "lost2")
            self.powlost2 = self._sum_merge(o, "powlost2")

            self.Minlost = self._fun_merge(o, numpy.minimum, "Minlost")
            self.Maxlost = self._fun_merge(o, numpy.maximum, "Maxlost")
            self.Minpowlost = self._fun_merge(o, numpy.minimum, "Minpowlost")
            self.Maxpowlost = self._fun_merge(o, numpy.maximum, "Maxpowlost")

            # Note: We are ignoring tab/tabp data...

            # merge final info (make sure to do this last!)
            self.Np = self._sum_merge(o, "Np")
            self.Nrun += o.Nrun

    def savetohdf(self, filename="Density.h5", group="TraceWin", force=False):
        """
        Saves data to HDF5
        """
        import h5py
        import sys

        fout = h5py.File(filename, "a")
        if group in fout:
            if force:
                del fout[group]
            else:
                if sys.flags.debug:
                    print("Group {} already exist in {}".format(group, filename))
                return

        group = fout.create_group(group)

        # header attributes..
        group.attrs["version"] = self.version
        group.attrs["year"] = self.year
        group.attrs["Nrun"] = self.Nrun
        group.attrs["vlong"] = self.vlong

        length = len(self.z)

        partran = sum(self.Np) > 0

        # one number per location
        arrays = ["z", "nelp", "ib", "Np", "Xouv", "Youv"]
        array_units = ["m", "", "mA", "", "m", "m"]
        if self.version >= 8:
            arrays += ["energy_accept", "phase_ouv_pos", "phase_ouv_neg"]
            array_units += ["eV", "deg", "deg"]
        if partran:
            arrays += [
                "lost2",
                "Minlost",
                "Maxlost",
                "powlost2",
                "Minpowlost",
                "Maxpowlost",
            ]
            array_units += ["", "", "", "W*w", "W", "W"]

        # 7 numbers per location..
        coordinates = ["moy", "moy2", "_max", "_min"]
        coordinate_units = ["m", "m*m", "m", "m"]
        if self.version >= 5 and partran:
            coordinates += ["rms_size", "rms_size2"]
            coordinate_units += ["m", "m*m"]
        if self.version >= 6 and partran:
            coordinates += ["min_pos_moy", "max_pos_moy"]
            coordinate_units += ["m", "m"]

        for val, unit in zip(arrays, array_units):
            data_set = group.create_dataset(val, (length,), dtype="f")
            data_set[...] = getattr(self, val)
            if unit:
                data_set.attrs["unit"] = unit

        for val, unit in zip(coordinates, coordinate_units):
            data_set = group.create_dataset(val, (length, 7), dtype="f")
            data_set[...] = getattr(self, val)
            if unit:
                data_set.attrs["unit"] = unit

        if self.version >= 7 and partran:
            # 3 numbers per location..
            emit_data = ["rms_emit", "rms_emit2"]
            emit_units = ["m*rad", "m*m*rad*rad"]
            for val, unit in zip(emit_data, emit_units):
                data_set = group.create_dataset(val, (length, 3), dtype="f")
                data_set[...] = getattr(self, val)
                if unit:
                    data_set.attrs["unit"] = unit
        if partran:
            # 1 numbers per location and per run..
            data = ["lost", "powlost"]
            units = ["", "W"]
            for val, unit in zip(data, units):
                data_set = group.create_dataset(val, (length, self.Nrun), dtype="f")
                data_set[...] = getattr(self, val)
                if unit:
                    data_set.attrs["unit"] = unit

        fout.close()


class remote_data_merger:
    def __init__(self, base="."):
        self._base = base
        self._files = []

    def add_file(self, filepath):
        import os

        if os.path.exists(filepath):
            fname = filepath
        else:
            fullpath = os.path.join(self._base, filepath)
            if os.path.exists(fullpath):
                fname = fullpath
            else:
                raise ValueError("Could not find file " + filepath)
        if fname not in self._files:
            self._files.append(fname)

    def generate_partran_out(self, filename=None):
        """
        Creates a string to be written to file
        each line is a list.

        If filename is given, writes directly to output file.

        """

        import numpy as np

        h1 = []
        h2 = []

        d1 = []
        d2 = []
        d3 = []

        if self._files:
            for f in self._files:
                string = open(f, "r").read()
                split = string.split("$$$")
                if split[9] != "Data_Error":
                    raise ValueError("Magic problem, please complain to Yngve")

                thisdata = split[10].strip().split("\n")

                if not h1:
                    h1 = [thisdata[0] + " (std in paranthesis)"]
                    h2 = thisdata[2:10]
                d1.append(thisdata[1].split())
                d2.append(thisdata[10])
                d3.append(thisdata[11])

            # fix d1:
            for i in range(len(d1)):
                for j in range(len(d1[0])):
                    d1[i][j] = float(d1[i][j])
            d1 = np.array(d1)
            means = d1.mean(axis=0)
            stds = d1.std(axis=0)
            d1 = []
            for i in range(len(stds)):
                if stds[i] / means[i] < 1e-10:
                    stds[i] = 0.0
            for i in range(len(stds)):
                # some small std are removed..
                if stds[i] / means[i] > 1e-8:
                    d1.append("%f(%f)" % (means[i], stds[i]))
                else:  # error is 0
                    d1.append(str(means[i]))
            d1 = [" ".join(d1)]

            # create data:
            data = h1 + d1 + h2 + d2 + d3

            if filename:
                open(filename, "w").write("\n".join(data))

            return data


class envDiag:
    """
    Read ENV_diag1.dat file

    This contains e.g. the absolute phase at each diag

    For now we do not read in all info from the file,
    so feel free to request or add anything else you would like.

    """

    def __init__(self, filename):
        self.filename = filename
        self.elements = {}
        # Needed to get an ordered dictionary:
        self._elementList = []
        self._readAsciiFile()
        self.units = {}
        self._setUnits()

    def _readAsciiFile(self):
        """
        Read the file
        """
        current = None
        for line in open(self.filename, "r"):
            lsp = line.split()
            if lsp[0] == "DIAG":
                self.elements[int(lsp[2])] = {}
                self._elementList.append(int(lsp[2]))
                current = self.elements[int(lsp[2])]
                current["loc"] = float(lsp[4])
            elif lsp[0] == "Ibeam:":
                current["current"] = float(lsp[1])
            elif lsp[0] == "Positions":
                current["phase"] = float(lsp[5])
                current["energy"] = float(lsp[6])
            elif lsp[0] == "RMS":
                current["x_rms"] = float(lsp[4]) * 0.01
                current["y_rms"] = float(lsp[5]) * 0.01
                current["phase_rms"] = float(lsp[6])
                current["energy_rms"] = float(lsp[7])
            elif lsp[0] == "Emittances":
                current["emit_x"] = float(lsp[3])
                current["emit_y"] = float(lsp[4])
                current["emit_z"] = float(lsp[5])
            elif lsp[0] == "Emittances99":
                current["emit99_x"] = float(lsp[3])
                current["emit99_y"] = float(lsp[4])
                current["emit99_z"] = float(lsp[5])
            elif lsp[0] == "Twiss":
                if lsp[1] == "Alpha" and lsp[3] == "(XXp,":
                    current["alpha_x"] = float(lsp[6])
                    current["alpha_y"] = float(lsp[7])
                elif lsp[1] == "Alpha" and lsp[3] == "(ZDp/p)":
                    current["alpha_z"] = float(lsp[5])
                elif lsp[1] == "Beta":
                    current["beta_x"] = float(lsp[5])
                    current["beta_y"] = float(lsp[6])
                    current["beta_z"] = float(lsp[7])

    def _setUnits(self):
        """
        Set the units for each element in the element dictionary
        (empty string for all undefined)
        """
        for key in ["loc", "x_rms", "y_rms"]:
            self.units[key] = "m"
        for key in ["emit_x", "emit_y", "emit_z", "emit99_x", "emit99_y", "emit99_z"]:
            self.units[key] = "Pi.mm.mrad"
        for key in ["current"]:
            self.units[key] = "mA"
        for key in ["energy", "energy_rms"]:
            self.units[key] = "MeV"
        for key in ["phase", "phase_rms"]:
            self.units[key] = "deg"
        for key in ["beta_x", "beta_y", "beta_z"]:
            self.units[key] = "mm/mrad"

        for element in self.elements:
            for key in self.elements[element]:
                if key not in self.units:
                    self.units[key] = ""

    def printTable(self):
        """
        Make a pretty print of the content
        """
        first = True
        rjust = 12
        for ekey in self._elementList:
            element = self.elements[ekey]

            # Print header if this is the first element..
            if first:
                keys = [key for key in element]
                print("#", end=" ")
                print("NUM".rjust(rjust), end=" ")
                for key in keys:
                    print(key.rjust(rjust), end=" ")
                print()
                print("#", end=" ")
                print("".rjust(rjust), end=" ")
                for key in keys:
                    print(self.units[key].rjust(rjust), end=" ")
                print()
                first = False

            print("  " + str(ekey).rjust(rjust), end=" ")
            for key in keys:
                num = element[key]
                if isinstance(num, float):
                    strnum = "{:.5e}".format(num)
                else:
                    strnum = str(element[key])
                print(strnum.rjust(rjust), end=" ")
            print()

    def getElement(self, elementId):
        """
        Returns the element dictionary for the given ID
        """
        return self.elements[elementId]

    def getElementAtLoc(self, location):
        """
        Returns a list of elements at the location requested
        """
        ret = []
        for key in self.elements:
            if abs(self.elements[key]["loc"] - location) < 1e-6:
                ret.append(self.elements[key])
        return ret

    def getParameterFromAll(self, parameter):
        """
        Returns a list containing the given parameter from all DIAGS,
        ordered by the location of the DIAGs
        """
        if parameter == "NUM":
            return self._elementList[:]
        ret = []
        for key in self._elementList:
            ret.append(self.elements[key][parameter])

        return ret


class partran(dict):
    """
    Read partran1.out files..

    This class can also read tracewin.out (same format)
    """

    def __init__(self, filename):
        self.filename = filename
        self._readAsciiFile()

    def _readAsciiFile(self):

        import numpy

        stream = open(self.filename, "r")
        for i in range(10):
            line = stream.readline()
            if line.strip()[0] == "#":
                break
        self.columns = ["NUM"] + line.split()[1:]
        self.data = numpy.loadtxt(stream)

        self._dict = {}
        for i in range(len(self.columns)):
            self[self.columns[i]] = self.data[:, i]


class field_map:
    """
    Class to read in the field map structures

    WARNING: Work in progress!!
    """

    def __init__(self, filename):
        self._filename = filename
        self._load_data(filename)

    def _load_data(self, filename):
        import os
        import numpy

        if not os.path.isfile(filename):
            raise ValueError("Cannot find file {}".format(filename))
        fin = open(filename, "r")
        line = fin.readline().split()
        self.header = []
        self.start = []
        self.end = []
        numindexes = []
        while len(line) > 1:
            [self.header.append(float(i)) for i in line]
            numindexes.append(int(line[0]) + 1)
            if len(line) == 2:
                self.start.append(0.0)
                self.end.append(float(line[1]))
            else:
                self.start.append(float(line[1]))
                self.end.append(float(line[2]))
            line = fin.readline().split()

        if len(self.start) == 1:
            self.z = numpy.mgrid[self.start[0] : self.end[0] : numindexes[0] * 1j]
            print(self.z)
        elif len(self.start) == 2:
            self.z, self.x = numpy.mgrid[
                self.start[0] : self.end[0] : numindexes[0] * 1j,
                self.start[1] : self.end[1] : numindexes[1] * 1j,
            ]
        elif len(self.start) == 3:
            self.z, self.x, self.y = numpy.mgrid[
                self.start[0] : self.end[0] : numindexes[0] * 1j,
                self.start[1] : self.end[1] : numindexes[1] * 1j,
                self.start[2] : self.end[2] : numindexes[2] * 1j,
            ]

        self.norm = float(line[0])
        self.header.append(self.norm)
        self.map = numpy.loadtxt(fin).reshape(numindexes)

    def get_flat_fieldmap(self):
        totmapshape = 1
        for i in self.map.shape:
            totmapshape *= i
        return self.map.reshape(totmapshape)

    def interpolate(self, npoints: tuple, method="cubic"):
        """
        Interpolate the map into a new mesh
        Each value should be an integer with the number of mesh points in each dimension
        intervals should be tuple-like with same number of elements
        as the map dimension, e.g. [0.8,0.8] for 2D
        Can also be a float if you want same interpolation factor in all planes

        method can be 'linear', 'nearest' or 'cubic'
        """
        import numpy
        from scipy.interpolate import griddata

        values = self.map.flatten()

        if len(self.start) == 1:
            points = self.z[:]
            self.z = numpy.mgrid[self.start[0] : self.end[0] : npoints[0] * 1j]
            self.map = griddata(points, values, self.z)
        if len(self.start) == 2:
            points = numpy.array([self.z.flatten(), self.x.flatten()]).transpose()
            self.z, self.x = numpy.mgrid[
                self.start[0] : self.end[0] : npoints[0] * 1j,
                self.start[1] : self.end[1] : npoints[1] * 1j,
            ]
            self.map = griddata(points, values, (self.z, self.x))
        if len(self.start) == 3:
            points = numpy.array(
                [self.z.flatten(), self.x.flatten(), self.y.flatten()]
            ).transpose()
            self.z, self.x, self.y = numpy.mgrid[
                self.start[0] : self.end[0] : npoints[0] * 1j,
                self.start[1] : self.end[1] : npoints[1] * 1j,
                self.start[2] : self.end[2] : npoints[2] * 1j,
            ]
            self.map = griddata(points, values, (self.z, self.x, self.y))
            self.header[0] = npoints[0] - 1
            self.header[2] = npoints[1] - 1
            self.header[5] = npoints[2] - 1

    def savemap(self, filename):
        fout = open(filename, "w")
        for n, s in zip(self.map.shape, self.size):
            fout.write("{} {}\n".format(n - 1, s))
        fout.write("{}\n".format(self.norm))
        totmapshape = 1
        for i in self.map.shape:
            totmapshape *= i
        data = self.map.reshape(totmapshape)
        for j in data:
            fout.write("{}\n".format(j))