diff --git a/ess/TraceWin.py b/ess/TraceWin.py index f35ded288853a9fcd2f022e0a707eee5a34376b2..e49e9900a95d66f602a9861c8797eda494a8b4ac 100644 --- a/ess/TraceWin.py +++ b/ess/TraceWin.py @@ -14,13 +14,22 @@ class dst: - phi [rad] - E [MeV] (kinetic energy) ''' - def __init__(self, filename): + def __init__(self, filename=None): # easy storage.. self.filename=filename # used to create dict behaviour.. self._columns=['x','xp','y','yp','phi','E'] - # read in the file.. - self._readBinaryFile() + if filename: + # read in the file.. + self._readBinaryFile() + else: + import numpy + self.Np=0 + self.Ib=0.0 + self.freq=352.21 + self._data=numpy.zeros((self.Np,6)) + self.mass=0.0 + def append(self,x=0.0,xp=0.0,y=0.0,yp=0.0,E=0.0,phi=0.0): ''' @@ -51,7 +60,7 @@ class dst: self.Np-=1 def _readBinaryFile(self): - # Thanks Emma! + # Thanks Ema! import numpy @@ -119,12 +128,12 @@ class dst: import numpy 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) + print(pack('b',125),end="",file=fout) + print(pack('b',100),end="",file=fout) + print(pack('i',self.Np),end="",file=fout) + print(pack('d',self.Ib),end="",file=fout) + print(pack('d',self.freq),end="",file=fout) + print(pack('b',125),end="",file=fout) data=self._data.copy() @@ -143,10 +152,9 @@ class dst: data=data.reshape(self.Np*6,1) for x in data: - out+=pack('d',x) + print(pack('d',x),end="",file=fout) - out+=pack('d',self.mass) - print >>fout, out + print(pack('d',self.mass),end="",file=fout) fout.close() def subplot(self, index, x, y=None, nb=100, mask=None): @@ -261,7 +269,6 @@ class plt: # Thanks Emma! import numpy - from builtins import range fin=open(self.filename,'r') @@ -297,7 +304,7 @@ class plt: while i<self.Ne: SubHeader=numpy.fromfile(fin, dtype=SubHeader_type, count=1) # unfinished files need this fix (simulation still running) - if not SubHeader['Nelp']: + if len(SubHeader['Nelp'])==0: break i=SubHeader['Nelp'][0] @@ -374,7 +381,6 @@ class plt: (NOT necessarily reference) ''' import numpy - from builtins import range if not hasattr(self,'avg'): self.calc_avg() @@ -426,7 +432,6 @@ class plt: ''' import numpy - from builtins import range if not hasattr(self,'avg'): self.calc_avg() @@ -449,7 +454,6 @@ class plt: ''' import numpy - from builtins import range if not hasattr(self,'sigma'): self.calc_sigma() @@ -469,7 +473,6 @@ class plt: ''' import numpy - from builtins import range if not hasattr(self,'sigma'): self.calc_sigma() @@ -497,6 +500,25 @@ class plt: self.twiss_alpha = numpy.array(self.twiss_alpha) + def save_dst(self, index,filename=''): + import numpy + + dset=self[index] + + _dst=dst() + _dst.freq=self.freq + _dst.Ib=self.Ib + _dst.Np=len(dset['x']) + _dst._data=numpy.array([dset['x'], + dset['xp'], + dset['y'], + dset['yp'], + dset['phi'], + dset['E']]).transpose() + if filename: + _dst.save(filename) + return _dst + class density_file: @@ -651,7 +673,6 @@ class density_file: def _getFullContent(self): import numpy - from builtins import range #self._getHeader() # no need to read the header again: @@ -962,7 +983,6 @@ class remote_data_merger: ''' import numpy as np - from builtins import range h1=[] h2=[] @@ -1018,8 +1038,6 @@ class partran(dict): ''' Read partran1.out files.. ''' - # python 3 forward compatibility: - from builtins import range def __init__(self,filename): self.filename=filename