TraceWin.py

            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 partran(dict):
    '''
    Read partran1.out files..
    '''

    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')
        l = fin.readline().split()
        self.header = []
        self.start = []
        self.end = []
        numindexes = []
        while len(l) > 1:
            [self.header.append(float(i)) for i in l]
            numindexes.append(int(l[0]) + 1)
            if len(l) == 2:
                self.start.append(0.0)
                self.end.append(float(l[1]))
            else:
                self.start.append(float(l[1]))
                self.end.append(float(l[2]))
            l = fin.readline().split()

        if len(self.start) == 1:
            self.z = numpy.mgrid[self.start[0]:self.end[0]:numindexes[0]*1j]
            print(new_z,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(l[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))