new function TraceWin.density_file.merge

ess/TraceWin.py

@@ -413,6 +413,8 @@ class density_file:
         self.i=0
         # z position [m] :
         self.z=numpy.zeros(counter)
+        # element index number
+        self.nelp=numpy.zeros(counter)
         # current [mA] :
         self.ib=numpy.zeros(counter)
         # number of lost particles:
@@ -428,10 +430,23 @@ class density_file:
         self._max=numpy.zeros((counter,7))
         self._min=numpy.zeros((counter,7))
 
+        if self.version>=5:
+            self.rms_size=numpy.zeros((counter,7))
+            self.rms_size2=numpy.zeros((counter,7))
+
+        if self.version>=6:
+            self.min_pos_moy=numpy.zeros((counter,7))
+            self.max_pos_moy=numpy.zeros((counter,7))
+
         if self.version>=7:
             self.rms_emit=numpy.zeros((counter,3))
             self.rms_emit2=numpy.zeros((counter,3))
 
+        if self.version>=8:
+            self.energy_accept=numpy.zeros(counter)
+            self.phase_ouv_pos=numpy.zeros(counter)
+            self.phase_ouv_neg=numpy.zeros(counter)
+
         self.lost=numpy.zeros((counter,self.Nrun))
         self.powlost=numpy.zeros((counter,self.Nrun))
 
@@ -486,6 +501,9 @@ class density_file:
             numpy.fromfile(self.fin, dtype=numpy.int16, count=((5588+self.Nrun*12)/2))
 
     def _get_7dim_array(array):
+        '''
+        Unused?
+        '''
         return dict(x=array[0],
                     y=array[1],
                     phase=array[2],
@@ -504,7 +522,7 @@ class density_file:
         # (though only if we are SURE about content!)
         numpy.fromfile(self.fin, dtype=numpy.int16, count=5)
 
-        nelp=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
+        self.nelp[self.i]=numpy.fromfile(self.fin, dtype=numpy.int32, count=1)[0]
         self.ib[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
         self.z[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
         # Aperture
@@ -521,23 +539,22 @@ class density_file:
         self._min[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)[:]
 
         if self.version>=5:
-            rms=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
-            rms_size2=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
+            self.rms_size[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
+            self.rms_size2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
         if self.version>=6:
-            min_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
-            max_pos_moy=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
+            self.min_pos_moy[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
+            self.max_pos_moy[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=n)
         if self.version>=7:
             self.rms_emit[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
             self.rms_emit2[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=3)[:]
         if self.version>=8:
-            e_ouv=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
-            phase_ouv_pos=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
-            phase_ouv_neg=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
+            self.energy_accept[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
+            self.phase_ouv_pos[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
+            self.phase_ouv_neg[self.i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)
 
         self.Np[self.i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
 
         if self.Np[self.i]:
-            powlost=numpy.zeros(self.Nrun)
             for i in xrange(self.Nrun):
                 self.lost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.int64, count=1)[0]
                 self.powlost[self.i,i]=numpy.fromfile(self.fin, dtype=numpy.float32, count=1)[0]
@@ -556,6 +573,83 @@ class density_file:
             if self.ib[self.i]>0:
                 tabp=numpy.fromfile(self.fin, dtype=numpy.uint32, count=3*step)
 
+    def _avg_merge(self,other,param):
+        '''
+        returns the average of the parameter
+        weighted by how many Nruns in self and other object
+        '''
+        mine=getattr(self,param)
+        new=getattr(other,param)
+        return (mine*self.Nrun+new*other.Nrun)/(self.Nrun+other.Nrun)
+
+    def merge(self,objects):
+        '''
+        Merge with list of objects
+        '''
+        import numpy
+
+        if not isinstance(objects,list):
+            raise TypeError("You tried to merge a non-list")
+
+        # for now we only allow objects with same version..
+        for o in objects:
+            if self.version != o.version:
+                raise ValueError("Cannot merge files with differing version")
+
+        # merge info..
+        for o in objects:
+            self.ib=self._avg_merge(o,'ib')
+
+            # this looks strange to me, but it is what TraceWin does..
+            self.moy+=o.moy
+            self.moy2+=o.moy2
+
+            self._max=numpy.maximum(self._max,o._max)
+            self._min=numpy.minimum(self._min,o._min)
+
+            if self.version>=5:
+                # this looks strange to me, but it is what TraceWin does..
+                self.rms_size+=o.rms_size
+                self.rms_size2+=o.rms_size2
+
+            if self.version>=6:
+                self.max_pos_moy=numpy.maximum(self.max_pos_moy,o.max_pos_moy)
+                self.min_pos_moy=numpy.minimum(self.min_pos_moy,o.min_pos_moy)
+
+            if self.version>=7:
+                # this looks strange to me, but it is what TraceWin does..
+                self.rms_emit+=o.rms_emit
+                self.rms_emit2+=o.rms_emit2
+
+            if self.version>=8:
+                # Warning: TraceWin does NOT merge these data in any way
+                self.energy_accept=self._avg_merge(o,'energy_accept')
+                self.phase_ouv_pos=self._avg_merge(o,'phase_ouv_pos')
+                self.phase_ouv_neg=self._avg_merge(o,'phase_ouv_neg')
+
+            # Note, we don't get into the problem of differing table sizes
+            # particles are lost, because we have written zeroes for
+            # the rest of the tables
+
+            # numpy.c_ == column stack objects..
+            self.lost=numpy.c_[self.lost,o.lost]
+            self.powlost=numpy.c_[self.powlost,o.powlost]
+
+            self.lost2+=o.lost2
+            self.powlost2+=o.powlost2
+
+            self.Minlost=numpy.minimum(self.Minlost,o.Minlost)
+            self.Maxlost=numpy.maximum(self.Maxlost,o.Maxlost)
+            self.Minpowlost=numpy.minimum(self.Minpowlost,o.Minpowlost)
+            self.Maxpowlost=numpy.maximum(self.Maxpowlost,o.Maxpowlost)
+
+            # Note: We are ignoring tab/tabp data...
+
+            # merge final info (make sure to do this last!)
+            self.Np+=o.Np
+            self.Nrun+=o.Nrun
+
+
 class remote_data_merger:
     def __init__(self, base='.'):
         self._base=base