improving of TraceWin.dst class

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()