fixed TraceWin.density_file.merge

if merging files with differing length of arrays, assume
that the shorter one lost all beam

ess/TraceWin.py

@@ -577,10 +577,71 @@ class density_file:
         '''
         returns the average of the parameter
         weighted by how many Nruns in self and other object
+
+        This allows for different lengths of the two arrays..
         '''
         mine=getattr(self,param)
         new=getattr(other,param)
-        return (mine*self.Nrun+new*other.Nrun)/(self.Nrun+other.Nrun)
+        if len(mine)>len(new):
+            ret = mine.copy()
+            ret[:len(new)] = (mine[:len(new)]*self.Nrun+new*other.Nrun)/(self.Nrun+other.Nrun)
+        elif len(mine)<len(new):
+            ret = new.copy()
+            ret[:len(mine)] = (mine*self.Nrun+new[:len(mine)]*other.Nrun)/(self.Nrun+other.Nrun)
+        else:
+            ret = (mine*self.Nrun+new*other.Nrun)/(self.Nrun+other.Nrun)
+        return ret
+
+    def _sum_merge(self,other,param):
+        '''
+        returns the sum of the parameter
+
+        This allows for different lengths of the two arrays..
+        '''
+        mine=getattr(self,param)
+        new=getattr(other,param)
+        if len(mine)>len(new):
+            ret = mine.copy()
+            ret[:len(new)] += new
+        elif len(mine)<len(new):
+            ret = new.copy()
+            ret[:len(mine)] += mine
+        else:
+            ret = mine+new
+        return ret
+
+    def _concatenate_merge(self,other,param):
+        '''
+        returns the concatenation of the two matrices
+
+        This allows for different lengths of the two arrays/matrices..
+        '''
+        import numpy
+
+        mine=getattr(self,param)
+        new=getattr(other,param)
+        ret=numpy.zeros(( max([len(mine),len(new)]), len(mine[0])+len(new[0]) ))
+        ret[:len(mine),:len(mine[0])]=mine
+        ret[:len(new),len(mine[0]):]=new
+        return ret
+
+    def _fun_merge(self,other,function,param):
+        '''
+        returns the function applied on the parameter
+
+        This allows for different lengths of the two arrays..
+        '''
+        mine=getattr(self,param)
+        new=getattr(other,param)
+        if len(mine)>len(new):
+            ret = mine.copy()
+            ret[:len(new)] = function(mine[:len(new)],new)
+        elif len(mine)<len(new):
+            ret = new.copy()
+            ret[:len(mine)] = function(mine,new[:len(mine)])
+        else:
+            ret = function(mine,new)
+        return ret
 
     def merge(self,objects):
         '''
@@ -598,28 +659,31 @@ class density_file:
 
         # merge info..
         for o in objects:
+            if len(self.ib)<len(o.ib):
+                raise ValueError("Sorry, not implemented yet. Complain to Yngve")
+
             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.moy=self._sum_merge(o,'moy')
+            self.moy2=self._sum_merge(o,'moy')
 
-            self._max=numpy.maximum(self._max,o._max)
-            self._min=numpy.minimum(self._min,o._min)
+            self._max=self._fun_merge(o,numpy.maximum,'_max')
+            self._min=self._fun_merge(o,numpy.minimum,'_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
+                self.rms_size=self._sum_merge(o,'rms_size')
+                self.rms_size2=self._sum_merge(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)
+                self.max_pos_moy=self._fun_merge(o,numpy.maximum,'max_pos_moy')
+                self.min_pos_moy=self._fun_merge(o,numpy.minimum,'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
+                self.rms_emit=self._sum_merge(o,'rms_emit')
+                self.rms_emit2=self._sum_merge(o,'rms_emit2')
 
             if self.version>=8:
                 # Warning: TraceWin does NOT merge these data in any way
@@ -631,22 +695,21 @@ class density_file:
             # 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.lost=self._concatenate_merge(o,'lost')
+            self.powlost=self._concatenate_merge(o,'powlost')
 
-            self.lost2+=o.lost2
-            self.powlost2+=o.powlost2
+            self.lost2=self._sum_merge(o,'lost2')
+            self.powlost2=self._sum_merge(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)
+            self.Minlost=self._fun_merge(o,numpy.minimum,'Minlost')
+            self.Maxlost=self._fun_merge(o,numpy.maximum,'Maxlost')
+            self.Minpowlost=self._fun_merge(o,numpy.minimum,'Minpowlost')
+            self.Maxpowlost=self._fun_merge(o,numpy.maximum,'Maxpowlost')
 
             # Note: We are ignoring tab/tabp data...
 
             # merge final info (make sure to do this last!)
-            self.Np+=o.Np
+            self.Np=self._sum_merge(o,'Np')
             self.Nrun+=o.Nrun
 
     def savetohdf(self,filename='Density.h5',group='TraceWin'):