new class ess.TraceWin.envDiag to read ENV_diag1.dat files

ess/TraceWin.py

@@ -142,7 +142,7 @@ class dst:
             self._data[:,i] = value
         except:
             raise ValueError("Available keys: "+str(self._columns))
-    
+
     def save(self, filename, toutatis=False):
         '''
         Save the distribution file
@@ -236,7 +236,7 @@ class dst:
         if y in ['E'] and max(dy)<0.1:
             dy*=1e3
             units['E']='keV'
-    
+
 
         plt.subplot(index)
         if y!=None:
@@ -255,7 +255,7 @@ class dst:
 
 class plt:
     '''
-    Simple class to read in a 
+    Simple class to read in a
     TraceWin plot file
 
     Class afterwards hold the following
@@ -264,7 +264,7 @@ class plt:
     - Np (number of particles)
     - Ib [A] (beam current)
     - freq [MHz]
-    - mc2  [MeV] 
+    - mc2  [MeV]
     - Nelp [m] (locations)
 
     each plt[i], where i is element number, holds:
@@ -406,7 +406,7 @@ class plt:
     def calc_rel(self):
         '''
         Calculates relativistic gamma/beta
-        at each position, based on 
+        at each position, based on
         AVERAGE beam energy
         (NOT necessarily reference)
         '''
@@ -1109,10 +1109,156 @@ class remote_data_merger:
 
             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
+        '''
+        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):
@@ -1177,8 +1323,8 @@ class field_map:
             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.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])