Skip to content
Snippets Groups Projects
TraceWin.py 33.6 KiB
Newer Older

            if filename:
                file(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=file(self.filename,'r')
        for i in xrange(10):
            line=stream.readline()
Yngve Levinsen's avatar
Yngve Levinsen committed
            if line.strip()[0]=='#':
                break
        self.columns=['NUM']+line.split()[1:]
        self.data=numpy.loadtxt(stream)

        self._dict={}
        for i in xrange(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,numpy

        if not os.path.isfile(filename):
            raise ValueError("Cannot find file {}".format(filename))
        fin = file(filename,'r')
        l=fin.readline().split()
        self.start=[]
        self.end=[]
        while len(l)>1:
            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()

        self.z=numpy.arange(self.start[0],self.end[0],(self.end[0]-self.start[0])/numindexes[0])
        if len(self.start)>1:
            self.y=numpy.arange(self.start[1],self.end[1],(self.end[1]-self.start[1])/numindexes[1])
        if len(self.start)>2:
            self.x=numpy.arange(self.start[2],self.end[2],(self.end[2]-self.start[2])/(numindexes[2]))

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

    def savemap(self, filename):
        fout=file(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))
        l=1
        for i in self.map.shape:
            l*=i
        data=self.map.reshape(l)
        for j in data:
            fout.write('{}\n'.format(j))