Add bunch_X properties to Beam class

Add bunch_X properties to Beam class
Add a generator over "not empty" bunches to Beam class
Modify Bunch class in order to have mean, std and emit as property and not method

Tracking/beam.py

@@ -8,7 +8,6 @@ Beam and bunch elements
 
 import numpy as np
 import pandas as pd
-import warnings
 
 class Bunch:
     """
@@ -38,10 +37,7 @@ class Bunch:
         Bunch current in [A]
         
     Methods
-    ------
-    stats()
-        Compute the mean and the standard deviation of alive particles for each
-        coordinates.
+    -------
     init_gaussian(cov=None, mean=None, **kwargs)
         Initialize bunch particles with 6D gaussian phase space.
     """
@@ -132,7 +128,8 @@ class Bunch:
     @current.setter
     def current(self, value):
         self.charge_per_mp = value * self.ring.T0 / self.__len__()
-        
+    
+    @property    
     def mean(self):
         """
         Compute the mean position of alive particles for each 
@@ -146,6 +143,7 @@ class Bunch:
         mean = [[self[name].mean()] for name in self]
         return np.array(mean)
     
+    @property
     def std(self):
         """
         Compute the standard deviation of the position of alive 
@@ -158,7 +156,8 @@ class Bunch:
         """
         std = [[self[name].std()] for name in self]
         return np.array(std)
-        
+    
+    @property    
     def emit(self):
         """
         Compute the bunch emittance for each plane [1].
@@ -245,7 +244,7 @@ class Beam:
         Total bunch charge in [C]
     particle_number : int
         Total number of particle in the beam
-    filling_pattern : list of bool
+    filling_pattern : array of bool
         Filling pattern of the beam
         
     Methods
@@ -270,9 +269,8 @@ class Beam:
     def __len__(self):
         """Return the number of (not empty) bunches"""
         length = 0
-        for bunch in self:
-            if bunch.current != 0:
-                length += 1
+        for bunch in self.not_empty:
+            length += 1
         return length        
     
     def __getitem__(self, i):
@@ -284,8 +282,17 @@ class Beam:
         self.bunch_list.__setitem__(i, value)
     
     def __iter__(self):
-        """Iterate over the bunches"""
+        """Iterate over all bunches"""
         return self.bunch_list.__iter__()
+   
+    @property             
+    def not_empty(self):
+        """Return a generator to iterate over not empty bunches"""
+        for bunch in self:
+            if bunch.current == 0:
+                pass
+            else:
+                yield bunch
         
     def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
                   mp_per_bunch=1e3):
@@ -333,41 +340,77 @@ class Beam:
                 
         self.bunch_list = bunch_list
         
-        for bunch in self:
+        for bunch in self.not_empty:
             bunch.init_gaussian()
+    
+    @property
+    def filling_pattern(self):
+        """Return an array with the filling pattern of the beam as bool"""
+        filling_pattern = []
+        for bunch in self:
+            if filling_pattern != 0:
+                filling_pattern.append(True)
+            else:
+                filling_pattern.append(False)
+        return np.array(filling_pattern)
         
+    @property
+    def bunch_current(self):
+        """Return an array with the current in each bunch in [A]"""
+        bunch_current = [bunch.current for bunch in self]
+        return np.array(bunch_current)
+    
+    @property
+    def bunch_charge(self):
+        """Return an array with the charge in each bunch in [C]"""
+        bunch_charge = [bunch.charge for bunch in self]
+        return np.array(bunch_charge)
+    
+    @property
+    def bunch_particle(self):
+        """Return an array with the particle number in each bunch"""
+        bunch_particle = [bunch.particle_number for bunch in self]
+        return np.array(bunch_particle)
+    
     @property
     def current(self):
         """Total beam current in [A]"""
-        current = 0
-        for bunch in self:
-            current += bunch.current
-        return current
+        return np.sum(self.bunch_current)
     
     @property
     def charge(self):
         """Total beam charge in [C]"""
-        charge = 0
-        for bunch in self:
-            charge += bunch.charge
-        return charge
+        return np.sum(self.bunch_charge)
     
     @property
     def particle_number(self):
-        """Total number of particle in the beam"""
-        particle_number = 0
-        for bunch in self:
-            particle_number += bunch.particle_number
-        return particle_number
+        """Total number of particles in the beam"""
+        return np.sum(self.bunch_particle)
     
     @property
-    def filling_pattern(self):
-        """Return an array with the filling pattern of the beam as bool"""
-        filling_pattern = []
-        for bunch in self:
-            if filling_pattern != 0:
-                filling_pattern.append(True)
-            else:
-                filling_pattern.append(False)
-        return np.array(filling_pattern)
-        
+    def bunch_mean(self):
+        """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):
+            bunch_mean[:,index] = np.squeeze(bunch.mean)
+        return bunch_mean
+    
+    @property
+    def bunch_std(self):
+        """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):
+            bunch_std[:,index] = np.squeeze(bunch.std)
+        return bunch_std
+    
+    @property
+    def bunch_emit(self):
+        """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):
+            bunch_emit[:,index] = np.squeeze(bunch.emit)
+        return bunch_emit
+    
\ No newline at end of file