diff --git a/ess/TraceWin.py b/ess/TraceWin.py index fdfe8e6dd2e516e4edddf87747f7fea780788521..096e05b260b1ec8b228659ab36f1f106e9f73af3 100644 --- a/ess/TraceWin.py +++ b/ess/TraceWin.py @@ -76,15 +76,19 @@ class dst: except: raise ValueError("Available keys: "+str(self._columns)) - def save(self, filename): + def save(self, filename, toutatis=False): ''' Save the distribution file so it can be read by TraceWin again + :param filename: Name of file + :param toutatis: Include 7th column of zeros + Stolen from Ryoichi's func.py (with permission) ''' from struct import pack + import numpy fout=open(filename,'w') out =pack('b',125) @@ -96,18 +100,25 @@ class dst: data=self._data.copy() + if toutatis and data.shape[1]==6: + data=numpy.append(data,numpy.zeros((len(data),1)),1) + elif not toutatis and data.shape[1]==7: + data=data[:,:-1] + # convert x,y from m to cm: data[:,0]*=1e2 data[:,2]*=1e2 - data=self._data.reshape(self.Np*6,1) + if toutatis: + data=data.reshape(self.Np*7,1) + else: + data=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()