Change how empty bunches are dealed with

Empty bunches are now not None any more but represented by Bunch class with 1 mp set to dead
Change init_beam to reflect the change

Tracking/beam.py

@@ -21,7 +21,8 @@ class Bunch:
         Macro-particle number
     current : float, optional
         Bunch current in [A]
-
+    alive : bool, optional
+        If False, the bunch is defined as empty 
     Attributes
     ----------
     mp_number : int
@@ -44,19 +45,26 @@ class Bunch:
         Initialize bunch particles with 6D gaussian phase space.
     """
     
-    def __init__(self, ring, mp_number=1e3, current=1e-3):
-        particles = {"x":np.zeros((int(mp_number),)),
-                     "xp":np.zeros((int(mp_number),)),
-                     "y":np.zeros((int(mp_number),)),
-                     "yp":np.zeros((int(mp_number),)),
-                     "tau":np.zeros((int(mp_number),)),
-                     "delta":np.zeros((int(mp_number),)),
-                     }
-        self.particles = pd.DataFrame(particles)
-        self.alive = pd.Series(np.ones((int(mp_number),),dtype=bool))
+    def __init__(self, ring, mp_number=1e3, current=1e-3, alive=True):
+        
         self.ring = ring
+        if not alive:
+            mp_number = 1
+            current = 0
         self._mp_number = int(mp_number)
+        
+        particles = {"x":np.zeros((self.mp_number,)),
+                     "xp":np.zeros((self.mp_number,)),
+                     "y":np.zeros((self.mp_number,)),
+                     "yp":np.zeros((self.mp_number,)),
+                     "tau":np.zeros((self.mp_number,)),
+                     "delta":np.zeros((self.mp_number,)),
+                     }
+        self.particles = pd.DataFrame(particles)
+        self.alive = pd.Series(np.ones((self.mp_number,),dtype=bool))
         self.current = current
+        if not alive:
+            self.alive = pd.Series(np.zeros((self.mp_number,),dtype=bool))            
         
     def __len__(self):
         """Return the number of alive particles"""
@@ -256,10 +264,10 @@ class Beam:
             self.bunch_list = bunch_list
             
     def __len__(self):
-        """Return the number of bunches"""
+        """Return the number of (not empty) bunches"""
         length = 0
         for bunch in self:
-            if bunch is not None:
+            if bunch.current != 0:
                 length += 1
         return length        
     
@@ -278,15 +286,22 @@ class Beam:
     def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
                   mp_per_bunch=1e3):
         """
-        Initialize beam with a given filling pattern but with uniform current
-        per bunch and marco-particle number per bunch.
+        Initialize beam with a given filling pattern and marco-particle number 
+        per bunch. Then initialize the different bunches with a 6D gaussian
+        phase space.
+        
+        If the filling pattern is an array of bool then the current per bunch 
+        is uniform, else the filling pattern can be an array with the current
+        in each bunch.
         
         Parameters
         ----------
         filling_pattern : numpy array or list of length ring.h
-            Filling pattern of the beam,
-        charge_per_bunch : float, optional
-            Charge per bunch in [C]
+            Filling pattern of the beam, can be a list or an array of bool, 
+            then current_per_bunch is used. Or can be an array with the current
+            in each bunch.
+        current_per_bunch : float, optional
+            Current per bunch in [A]
         mp_per_bunch : float, optional
             Macro-particle number per bunch
         """
@@ -295,24 +310,34 @@ class Beam:
             raise ValueError(("The length of filling pattern is {} ".format(len(filling_pattern)) + 
                               "but should be {}".format(self.ring.h)))
         
+        filling_pattern = np.array(filling_pattern)
         bunch_list = []
-        for value in filling_pattern:
-            if value == True:
-                bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
-            elif value == False:
-                bunch_list.append(None)
-            else:
-                raise ValueError("{} should be True or False".format(value))
+        if filling_pattern.dtype == np.dtype("bool"):
+            for value in filling_pattern:
+                if value == True:
+                    bunch_list.append(Bunch(self.ring, mp_per_bunch, current_per_bunch))
+                elif value == False:
+                    bunch_list.append(Bunch(self.ring, alive=False))
+        elif filling_pattern.dtype == np.dtype("float64"):
+            for current in filling_pattern:
+                if current != 0:
+                    bunch_list.append(Bunch(self.ring, mp_per_bunch, current))
+                elif current == 0:
+                    bunch_list.append(Bunch(self.ring, alive=False))
+        else:
+            raise TypeError("{} should be bool or float64".format(filling_pattern.dtype))
+                
         self.bunch_list = bunch_list
-        # init bunches => gauss
+        
+        for bunch in self:
+            bunch.init_gaussian()
         
     @property
     def current(self):
         """Total beam current in [A]"""
         current = 0
         for bunch in self:
-            if bunch is not None:
-                current += bunch.current
+            current += bunch.current
         return current
     
     @property
@@ -320,8 +345,7 @@ class Beam:
         """Total beam charge in [C]"""
         charge = 0
         for bunch in self:
-            if bunch is not None:
-                charge += bunch.charge
+            charge += bunch.charge
         return charge
     
     @property
@@ -329,18 +353,17 @@ class Beam:
         """Total number of particle in the beam"""
         particle_number = 0
         for bunch in self:
-            if bunch is not None:
-                particle_number += bunch.particle_number
+            particle_number += bunch.particle_number
         return particle_number
     
     @property
     def filling_pattern(self):
-        """Filling pattern of the beam"""
+        """Return an array with the filling pattern of the beam as bool"""
         filling_pattern = []
         for bunch in self:
-            if isinstance(bunch, Bunch):
+            if filling_pattern != 0:
                 filling_pattern.append(True)
             else:
                 filling_pattern.append(False)
-        return filling_pattern
-        
+        return np.array(filling_pattern)
+        
\ No newline at end of file