From 2d994a8da6bd30be2e6663c162fa2bec00652c1e Mon Sep 17 00:00:00 2001
From: Yngve Inntjore Levinsen <Yngve.Levinsen@esss.se>
Date: Fri, 20 Mar 2015 16:59:05 +0100
Subject: [PATCH] in dst class: - added more stuff for dictionary type
behaviour - added save functionality (Thanks Ryoichi!) in plt class: - more
goodies added
---
ess/TraceWin.py | 86 +++++++++++++++++++++++++++++++++++++++++--------
1 file changed, 72 insertions(+), 14 deletions(-)
diff --git a/ess/TraceWin.py b/ess/TraceWin.py
index 2b34e42..e8df760 100644
--- a/ess/TraceWin.py
+++ b/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:
--
GitLab