diff --git a/ess/TraceWin.py b/ess/TraceWin.py index d3bd2cac31f9a2d180e97ed69a5ef5f9b84be568..fa66906f8d8397f983d063506f4679daec26145d 100644 --- a/ess/TraceWin.py +++ b/ess/TraceWin.py @@ -15,7 +15,11 @@ class dst: - phi [rad] - E [MeV] (kinetic energy) ''' - def __init__(self, filename=None): + def __init__(self, + filename=None, + freq=352.21, + mass=938.272, + Ib=0.0): # easy storage.. self.filename = filename # used to create dict behaviour.. @@ -26,10 +30,10 @@ class dst: else: import numpy self.Np = 0 - self.Ib = 0.0 - self.freq = 352.21 + self.Ib = Ib + self.freq = freq self._data = numpy.zeros((self.Np, 6)) - self.mass = 0.0 + self.mass = mass def append(self, x=0.0, xp=0.0, y=0.0, yp=0.0, E=0.0, phi=0.0): ''' @@ -45,6 +49,19 @@ class dst: self._data = numpy.append(self._data, [[x, xp, y, yp, phi, E]], 0) self.Np += 1 + def append_many(self, array): + ''' + Append a matrix of particle vectors + + Matrix on form 6xN, where N is number of particles. + Each row should hold [x,xp,y,yp,phi,E] + Units m,rad, MeV + ''' + import numpy + + self._data = numpy.append(self._data, array, 0) + self.Np += len(array) + def remove(self, i=None): ''' Removes all particles from the distribution, or the line specified by i @@ -150,8 +167,7 @@ class dst: else: data = data.reshape(self.Np*6, 1) - for x in data: - fout.write(pack('d', x)) + fout.write(pack('{}d'.format(len(data)),*data)) fout.write(pack('d', self.mass)) fout.close()