Add a bunch_index property to Beam class

bunch_index returns the array with the positions (index) of the non-empty bunches, to be used with Beam.not_empty iterator.
bunch_mean, bunch_std and bunch_emit now iterate only over non empty bunches.

tracking/particles.py

@@ -427,6 +427,8 @@ class Beam:
     mpi_switch : bool
         Status of MPI parallelisation, should not be changed directly but with
         mpi_init() and mpi_close()
+    bunch_index : array of shape (len(self,))
+        Return an array with the positions of the non-empty bunches
         
     Methods
     ------
@@ -608,6 +610,11 @@ class Beam:
     def filling_pattern(self):
         """Return an array with the filling pattern of the beam as bool"""
         return self._filling_pattern
+    
+    @property
+    def bunch_index(self):
+        """Return an array with the positions of the non-empty bunches."""
+        return np.where(self.filling_pattern == True)[0]
         
     @property
     def bunch_current(self):
@@ -647,7 +654,8 @@ class Beam:
         """Return an array with the mean position of alive particles for each
         bunches"""
         bunch_mean = np.zeros((6,self.ring.h))
-        for index, bunch in enumerate(self):
+        for idx, bunch in enumerate(self.not_empty):
+            index = self.bunch_index[idx]
             bunch_mean[:,index] = bunch.mean
         return bunch_mean
     
@@ -656,7 +664,8 @@ class Beam:
         """Return an array with the standard deviation of the position of alive 
         particles for each bunches"""
         bunch_std = np.zeros((6,self.ring.h))
-        for index, bunch in enumerate(self):
+        for idx, bunch in enumerate(self.not_empty):
+            index = self.bunch_index[idx]
             bunch_std[:,index] = bunch.std
         return bunch_std
     
@@ -665,7 +674,8 @@ class Beam:
         """Return an array with the bunch emittance of alive particles for each
         bunches and each plane"""
         bunch_emit = np.zeros((3,self.ring.h))
-        for index, bunch in enumerate(self):
+        for idx, bunch in enumerate(self.not_empty):
+            index = self.bunch_index[idx]
             bunch_emit[:,index] = bunch.emit
         return bunch_emit
     

tracking/rf.py

@@ -187,7 +187,7 @@ class CavityResonator():
             self.bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
             
         self.distance = beam.distance_between_bunches
-        self.valid_bunch_index = np.where(beam.filling_pattern == True)[0]
+        self.valid_bunch_index = beam.bunch_index
         self.tracking = True
         self.nturn = 0