in dst class:

 - added more stuff for dictionary type behaviour
 - added save functionality (Thanks Ryoichi!)
in plt class:
 - more goodies added

ess/TraceWin.py

@@ -48,12 +48,50 @@ class dst:
         Footer=numpy.fromfile(fin, dtype=numpy.float64, count=1)
         self.mass=Footer[0]
 
+    def keys(self):
+        return self._columns[:]
+
     def __getitem__(self, key):
         # makes the class function as a dictionary
         # e.g. dst['x'] returns the x array..
-        if key in self._columns:
+        try:
             i=self._columns.index(key)
             return self._data[:,i]
+        except:
+            raise ValueError("Available keys: "+str(self._columns))
+
+    def __setitem__(self, key, value):
+        try:
+            i=self._columns.index(key)
+            self._data[:,i]=value
+        except:
+            raise ValueError("Available keys: "+str(self._columns))
+    
+    def save(self, filename):
+        '''
+        Save the distribution file
+        so it can be read by TraceWin again
+
+        Stolen from Ryoichi's func.py (with permission)
+        '''
+
+        from struct import pack
+
+        fout=open(filename,'w')
+        out =pack('b',125)
+        out+=pack('b',100)
+        out+=pack('i',self.Np)
+        out+=pack('d',self.Ib)
+        out+=pack('d',self.freq)
+        out+=pack('b',125)
+        data=self._data.reshape(self.Np*6,1)
+        for x in data:
+            out+=pack('d',x)
+        out+=pack('d',self.mass)
+        print >>fout, out
+        #data.tofile(fout)
+        fout.close()
+
 
 class plt:
     '''
@@ -199,7 +237,7 @@ class plt:
 
         self.avg=dict(x=[], xp=[], y=[], yp=[], E=[], phi=[])
 
-        vals=['x', 'xp', 'y', 'yp', 'E', 'phi']
+        vals=self._columns
 
         for i in self.Nelp:
             data=self[i]
@@ -232,27 +270,47 @@ class plt:
             self.min[v]=numpy.array(self.min[v])
             self.max[v]=numpy.array(self.max[v])
 
-    def calc_std(self):
+    def calc_sigma(self):
+        '''
+        Calculates the sigma matrix
+
+        Creates self.sigma such that self.sigma[i,j]
+        returns the sigma matrix for value i,j.
+
+        The numbering is:
+        0: x
+        1: xp
+        2: y
+        3: yp
+        4: E
+        5: phi
         '''
-        Calculates the sigma of each coordinate
-        such that e.g. self.sigma['x'] is the sigma_x
 
-        Note: SigmaX etc from partran1.out are in mm, not cm
+        import numpy
 
-        Units: cm
+        if not hasattr(self,'avg'):
+               self.calc_avg()
+
+        vals=self._columns
+
+        self.sigma=numpy.array([[numpy.mean( (self[n]-self.avg[n]) * (self[m] - self.avg[m]) ) for n in vals] for m in vals])
+        print self.sigma.shape
+
+    def calc_twiss(self):
+        '''
+        Calculates emittance, beta, alfa, gamma
+        for each plane, x-xp, y-yp, and E-phi
         '''
+
         import numpy
 
+        if not hasattr(self,'sigma'):
+               self.calc_sigma()
 
-        self.sigma=dict(x=[], xp=[], y=[], yp=[])
+        self.eps = [numpy.sqrt(numpy.det(self.sigma[i:i+2][:,i:i+2])) for i in (0,2,4)]
+        self.beta = [self.sigma[i][j]]
 
-        for i in self.Nelp:
-            data=self[i]
-            for key in self.sigma.keys():
-                self.sigma[key].append(data[key].std())
 
-        for key in self.sigma.keys():
-            self.sigma[key]=numpy.array(self.sigma[key])
 
 
 class density_file: