diff --git a/ess/TraceWin.py b/ess/TraceWin.py index 150e5e41d829eee2e0bd5035cf553696e581b3f4..5cfefe95a0c92fa156a5a034cd5b53f429d96955 100644 --- a/ess/TraceWin.py +++ b/ess/TraceWin.py @@ -499,8 +499,10 @@ class plt: self.twiss_alpha = [[-self.sigma[j][i][i+1]/self.twiss_eps[j,i/2] for i in (0,2,4)] for j in range(len(self.Nelp))] self.twiss_alpha = numpy.array(self.twiss_alpha) - - def save_dst(self, index,filename=''): + def get_dst(self, index): + ''' + Returns the dst corresponding to the given index + ''' import numpy dset=self[index] @@ -516,8 +518,16 @@ class plt: dset['yp'], dset['phi'], dset['E']]).transpose() - if filename: - _dst.save(filename) + return _dst + + def save_dst(self, index,filename): + ''' + Saves the dst at the specified index to file + + Returns the same dst object. + ''' + _dst=get_dst(index) + _dst.save(filename) return _dst