added TraceWin.plt

ess/TraceWin.py

@@ -55,6 +55,106 @@ class dst:
             i=self._columns.index(key)
             return self._data[:,i]
 
+class plt:
+    '''
+    Simple class to read in a 
+    TraceWin plot file
+
+    Class afterwards hold the following
+    dictionary items:
+      - Ne (number of locations)
+      - Np (number of particles)
+      - Ib [A] (beam current)
+      - freq [MHz]
+      - mc2  [MeV] 
+      - Nelp [m] (locations)
+
+    each plt[i], where i is element number, holds:
+      - Zgen [cm] (location)
+      - phase0 [deg] (ref phase)
+      - wgen [MeV] (ref energy)
+      - x [array, cm]
+      - xp [array, rad]
+      - y [array, cm]
+      - yp [array, rad]
+      - phi [array, rad]
+      - E [array, MeV]
+      - l [array] (is lost)
+
+      Example::
+        plt=ess.TraceWin.plt('calc/dtl1.plt')
+        for i in [97,98]:
+          data=plt[i]$
+          if data:
+            print data['x']
+    '''
+    def __init__(self, filename):
+        # easy storage..
+        self.filename=filename
+        # used to create dict behaviour..
+        self._columns=['x','xp','y','yp','phi','E', 'l']
+        # read in the file..
+        self._readBinaryFile()
+
+    def _readBinaryFile(self):
+        # Thanks Emma!
+
+        import numpy
+
+        fin=file(self.filename,'r')
+
+        # dummy, Np, Ib, freq, dummy
+        Header_type = numpy.dtype([
+            ('dummy12', numpy.int16),
+            ('Ne',      numpy.int32),
+            ('Np',      numpy.int32),
+            ('Ib',      numpy.float64),
+            ('freq',    numpy.float64),
+            ('mc2',     numpy.float64),
+            ])
+        SubHeader_type = numpy.dtype([
+             ('dummy12', numpy.int8),
+             ('Nelp',      numpy.int32),
+             ('Zgen',     numpy.float64),
+             ('phase0',     numpy.float64),
+             ('wgen',     numpy.float64),
+             ])
+
+
+        Header=numpy.fromfile(fin, dtype=Header_type, count=1)
+        self.Np=Header['Np'][0]
+        self.Ne=Header['Ne'][0]
+        self.Ib=Header['Ib'][0]
+        self.freq=Header['freq'][0]
+        self.mc2=Header['mc2'][0]
+
+        self._data=[]
+        self.Nelp=[]
+
+        i=0
+        while i<self.Ne:
+            SubHeader=numpy.fromfile(fin, dtype=SubHeader_type, count=1)
+            i=SubHeader['Nelp'][0]
+
+            self.Nelp.append(i)
+
+            Table=numpy.fromfile(fin, dtype=numpy.float32, count=self.Np*7)
+            Table=Table.reshape(self.Np,7)
+            data={}
+            for key in ['Zgen','phase0','wgen']:
+                data[key]=SubHeader[key][0]
+            for j in xrange(7):
+                    c=self._columns[j]
+                    data[c]=Table[:,j]
+            self._data.append(data)
+
+    def __getitem__(self, key):
+        if key in self.Nelp:
+            i=self.Nelp.index(key)
+            return self._data[i]
+        else:
+            print "No data to plot at element",key
+
 class density_file:
     '''
     Simple class to read a TraceWin density file