Binning on Bunch + others

Add T1 = RF period to synchrotron object
Add binning and plot_profile to Bunch object // plot_profile to improve
Modify beam.filling_pattern to store the array instead of computing it every time
Add distance_between_bunches method to Beam object
Modify not_empty method of Beam to use filling_pattern

tracking/particles.py

@@ -10,6 +10,7 @@ import numpy as np
 import pandas as pd
 from tracking.parallel import Mpi
 from scipy.constants import c, m_e, m_p, e
+import matplotlib.pyplot as plt
 
 class Particle:
     """
@@ -271,6 +272,53 @@ class Bunch:
         self.particles["tau"] = values[:,4]
         self.particles["delta"] = values[:,5]
         
+    def binning(self, dimension="tau", n_bin=75):
+        """
+        Bin macro-particles.
+
+        Parameters
+        ----------
+        dimension : str, optional
+            Dimension in which the binning is done. The default is "tau".
+        n_bin : int, optional
+            Number of bins. The default is 75.
+
+        Returns
+        -------
+        bins : IntervalIndex of shape (n_bin,)
+            Bins where the particles are sorted.
+        sorted_index : Series of shape (self.mp_number,)
+            Bin number of each macro-particles.
+        profile : Series of shape (n_bin,)
+            Number of marco-particles in each bin.
+
+        """
+        bins = pd.interval_range(start=min(self[dimension]), 
+                                 end=max(self[dimension]), 
+                                 periods=n_bin,
+                                 name=dimension)
+        sorted_index = pd.cut(self[dimension], bins)
+        sorted_index = sorted_index.reindex(self.alive.index)
+        profile = sorted_index.value_counts().sort_index()
+        return (bins, sorted_index, profile)
+    
+    def plot_profile(self, dimension="tau", n_bin=75):
+        """
+        Plot bunch profile.
+
+        Parameters
+        ----------
+        dimension : str, optional
+            Dimension to plot. The default is "tau".
+        n_bin : int, optional
+            Number of bins. The default is 75.
+
+        """
+        bins, sorted_index, profile = self.binning(dimension, n_bin)
+        fig = plt.figure()
+        ax = fig.gca()
+        ax.plot(bins.mid, profile)
+        
 class Beam:
     """
     Define a Beam object composed of several Bunch objects. 
@@ -355,12 +403,29 @@ class Beam:
    
     @property             
     def not_empty(self):
-        """Return a generator to iterate over not empty bunches"""
-        for bunch in self:
-            if bunch.current == 0:
-                pass
+        """Return a generator to iterate over not empty bunches."""
+        for index, value in enumerate(self.filling_pattern):
+            if value == True:
+                yield self[index]
             else:
-                yield bunch
+                pass
+    
+    @property
+    def distance_between_bunches(self):
+        filling_pattern = self.filling_pattern
+        distance = np.zeros(filling_pattern.shape)
+        for index, value in enumerate(filling_pattern):
+            if value == False:
+                pass
+            elif value == True:
+                count = 1
+                for value2 in filling_pattern[index+1:]:
+                    if value2 == False:
+                        count += 1
+                    elif value2 == True:
+                        break
+                distance[index] = count
+        return distance
         
     def init_beam(self, filling_pattern, current_per_bunch=1e-3, 
                   mp_per_bunch=1e3):
@@ -410,17 +475,23 @@ class Beam:
         
         for bunch in self.not_empty:
             bunch.init_gaussian()
+            
+        self.update_filling_pattern()
     
-    @property
-    def filling_pattern(self):
-        """Return an array with the filling pattern of the beam as bool"""
+    def update_filling_pattern(self):
+        """Update the beam filling pattern."""
         filling_pattern = []
         for bunch in self:
             if bunch.current != 0:
                 filling_pattern.append(True)
             else:
                 filling_pattern.append(False)
-        return np.array(filling_pattern)
+        self._filling_pattern = np.array(filling_pattern)
+    
+    @property
+    def filling_pattern(self):
+        """Return an array with the filling pattern of the beam as bool"""
+        return self._filling_pattern
         
     @property
     def bunch_current(self):

tracking/synchrotron.py

@@ -53,6 +53,8 @@ class Synchrotron:
         Revolution frequency in [Hz].
     omega0 : float
         Angular revolution frequency in [Hz.rad]
+    T1 : flaot
+        Fundamental RF period in [s].
     f1 : float
         Fundamental RF frequency in [Hz].
     omega1 : float
@@ -112,6 +114,7 @@ class Synchrotron:
     def L(self,value):
         self._L = value
         self._T0 = self.L/c
+        self._T1 = self.T0/self.h
         self._f0 = 1/self.T0
         self._omega0 = 2*np.pi*self.f0
         self._f1 = self.h*self.f0
@@ -127,6 +130,15 @@ class Synchrotron:
     def T0(self, value):
         self.L = c*value
         
+    @property
+    def T1(self):
+        """"Fundamental RF period [s]"""
+        return self._T1
+    
+    @T1.setter
+    def T1(self, value):
+        self.L = c*value*self.h
+        
     @property
     def f0(self):
         """Revolution frequency [Hz]"""