Refactor and small improvements on CavityResonator

Add comments.
Remove usage of bins, use only center.

tracking/rf.py

@@ -169,14 +169,13 @@ class CavityResonator():
         
     def init_tracking(self, beam):
         """
-        Initialization of the tracking
+        Initialization of the tracking.
 
         Parameters
         ----------
         beam : Beam object
 
         """
-        
         if beam.mpi_switch:
             self.bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
             
@@ -185,10 +184,19 @@ class CavityResonator():
         self.tracking = True
         self.nturn = 0
     
-    def track(self, beam, ref_frame="beam"):
-        """version type mbtrack + formule analytique
+    def track(self, beam):
+        """
+        Track a Beam object through the CavityResonator object.
         
-        phasor is given at t=0 synchronous particle of the bunch 0
+        Can be used with or without mpi.
+        
+        The beam phasor is given at t=0 (synchronous particle) of the first 
+        non empty bunch.
+
+        Parameters
+        ----------
+        beam : Beam object
+
         """
         
         if self.tracking is False:
@@ -200,7 +208,6 @@ class CavityResonator():
             
             if beam.mpi_switch:
                 # get shared bunch profile for current bunch
-                bins = beam.mpi.bins_all[j]
                 center = beam.mpi.center_all[j]
                 profile = beam.mpi.profile_all[j]
                 bin_length = center[1]-center[0]
@@ -216,19 +223,18 @@ class CavityResonator():
             
             energy_change = bunch["tau"]*0
             
-            # phasor is given at t=0 synchronous particle of the bunch
-            # remove part of beam phasor decay to be at the start of the binning
-            self.phasor_decay(bins[0], ref_frame=ref_frame)
+            # remove part of beam phasor decay to be at the start of the binning (=bins[0])
+            self.phasor_decay(center[0] - bin_length/2, ref_frame="beam")
             
             if index != self.bunch_index:
-                self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame=ref_frame)
+                self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="beam")
             else:
                 # modify beam phasor
                 for i, center0 in enumerate(center):
                     mp_per_bin = profile[i]
                     
                     if mp_per_bin == 0:
-                        self.phasor_decay(bin_length, ref_frame=ref_frame)
+                        self.phasor_decay(bin_length, ref_frame="beam")
                         continue
                     
                     ind = (sorted_index == i)
@@ -239,22 +245,32 @@ class CavityResonator():
                     energy_change[ind] = Vtot / self.ring.E0
 
                     self.beam_phasor -= 2*charge_per_mp*self.loss_factor*mp_per_bin
-                    self.phasor_decay(bin_length, ref_frame=ref_frame)
-                
-            self.phasor_decay(-1*bins[-1], ref_frame=ref_frame)
-            self.phasor_decay(self.distance[index] * self.ring.T1, ref_frame=ref_frame)
+                    self.phasor_decay(bin_length, ref_frame="beam")
+            
+            # phasor decay to be at t=0 of the current bunch (=-1*bins[-1])
+            self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="beam")
+            # phasor decay to be at t=0 of the next bunch
+            self.phasor_decay( (self.distance[index] * self.ring.T1), ref_frame="beam")
             
             if index == self.bunch_index:
                 # apply kick
                 bunch["delta"] += energy_change
-                bunch.energy_change = energy_change
                 
             self.nturn += 1
                 
         
     def init_phasor(self, beam):
-        """version type mbtrack + formule analytique"""
+        """
+        Initialize the beam phasor for a given beam distribution.
         
+        Follow the same steps as the track method but in the "rf" reference 
+        frame and without any modifications on the beam.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        """        
         if self.tracking is False:
             self.init_tracking(beam)
             
@@ -267,7 +283,6 @@ class CavityResonator():
                 
                 if beam.mpi_switch:
                     # get shared bunch profile for current bunch
-                    bins = beam.mpi.bins_all[j]
                     center = beam.mpi.center_all[j]
                     profile = beam.mpi.profile_all[j]
                     bin_length = center[1]-center[0]
@@ -278,11 +293,10 @@ class CavityResonator():
                     bin_length = center[1]-center[0]
                     charge_per_mp = bunch.charge_per_mp
                 
-                # phasor is given at t=0 synchronous particle of the bunch
-                # remove part of beam phasor decay to be at the start of the binning
-                self.phasor_decay(bins[0], ref_frame="rf")
+                self.phasor_decay(center[0] - bin_length/2, ref_frame="rf")
                 self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="rf")
-                self.phasor_decay( (self.distance[index] * self.ring.T1) - bins[-1], ref_frame="rf")
+                self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="rf")
+                self.phasor_decay( (self.distance[index] * self.ring.T1), ref_frame="rf")
             
     def phasor_decay(self, time, ref_frame="beam"):
         """
@@ -293,6 +307,8 @@ class CavityResonator():
         ----------
         time : float
             Time span in [s], can be positive or negative.
+        ref_frame : string, optional
+            Reference frame to be used, can be "beam" or "rf".
 
         """
         if ref_frame == "beam":
@@ -305,6 +321,8 @@ class CavityResonator():
     def phasor_evol(self, profile, bin_length, charge_per_mp, ref_frame="beam"):
         """
         Compute the beam phasor evolution during the crossing of a bunch.
+        
+        Assume that the phasor decay happens before the beam loading.
 
         Parameters
         ----------
@@ -314,6 +332,9 @@ class CavityResonator():
             Length of a bin in [s].
         charge_per_mp : float
             Charge per macro-particle in [C].
+        ref_frame : string, optional
+            Reference frame to be used, can be "beam" or "rf".
+            
         """
         if ref_frame == "beam":
             delta = self.wr