Add n_bin optional parameter for CavityResonator

Add n_bin optional parameter for CavityResonator used only if mpi is false.
Various docstring improvements.

mbtrack2/tracking/monitors/monitors.py

@@ -1357,7 +1357,7 @@ class BeamSpectrumMonitor(Monitor):
         
 class CavityMonitor(Monitor):
     """
-    Monitor a CavityResonator object and save attributes (mean, std, emit and current).
+    Monitor a CavityResonator object and save attributes.
     
     Parameters
     ----------
@@ -1453,4 +1453,4 @@ class CavityMonitor(Monitor):
                     pass
             else:                            
                 raise TypeError("cavity should be a CavityResonator object.")
-        self.track_count += 1       
+        self.track_count += 1       
\ No newline at end of file

mbtrack2/tracking/parallel.py

@@ -30,6 +30,10 @@ class Mpi:
         associated bunch number
     bunch_num : int
         Return the bunch number corresponding to the current processor
+    next_bunch : int
+        Return the rank of the next tracked bunch
+    previous_bunch : int
+        Return the rank of the previous tracked bunch
         
     Methods
     -------
@@ -40,6 +44,10 @@ class Mpi:
         Return the bunch number corresponding to rank
     bunch_to_rank(bunch_num)
         Return the rank corresponding to the bunch number bunch_num
+    share_distributions(beam)
+        Compute the bunch profiles and share it between the different bunches.
+    share_means(beam)
+        Compute the bunch means and share it between the different bunches.
         
     References
     ----------

mbtrack2/tracking/rf.py

@@ -51,9 +51,13 @@ class RFCavity(Element):
     
 class CavityResonator():
     """Cavity resonator class for active or passive RF cavity with beam
-    loading or HOM, based on [1].
+    loading or HOM, based on [1,2].
+    
     Use cosine definition.
     
+    If used with mpi, beam.mpi.share_distributions must be called before the 
+    track method call.
+    
     Parameters
     ----------
     ring : Synchrotron object
@@ -75,6 +79,11 @@ class CavityResonator():
         Total cavity voltage in [V].
     theta : float, optional
         Total cavity phase in [rad].
+    n_bin : int, optional
+        Number of bins used for the beam loading computation. 
+        Only used if MPI is not used, otherwise n_bin must be specified in the 
+        beam.mpi.share_distributions method.
+        The default is 75.
         
     Attributes
     ----------
@@ -172,8 +181,13 @@ class CavityResonator():
     [1] Wilson, P. B. (1994). Fundamental-mode rf design in e+ e− storage ring 
     factories. In Frontiers of Particle Beams: Factories with e+ e-Rings 
     (pp. 293-311). Springer, Berlin, Heidelberg.
+    
+    [2] Yamamoto, Naoto, Alexis Gamelin, and Ryutaro Nagaoka. "Investigation 
+    of Longitudinal Beam Dynamics With Harmonic Cavities by Using the Code 
+    Mbtrack." IPAC’19, Melbourne, Australia, 2019.
     """
-    def __init__(self, ring, m, Rs, Q, QL, detune, Ncav=1, Vc=0, theta=0):
+    def __init__(self, ring, m, Rs, Q, QL, detune, Ncav=1, Vc=0, theta=0, 
+                 n_bin=75):
         self.ring = ring
         self.m = m
         self.Ncav = Ncav
@@ -194,6 +208,7 @@ class CavityResonator():
         self.Vgr = 0
         self.theta_gr = 0
         self.Pg = 0
+        self.n_bin = int(n_bin)
         
     def init_tracking(self, beam):
         """
@@ -217,7 +232,7 @@ class CavityResonator():
         Track a Beam object through the CavityResonator object.
         
         Can be used with or without mpi.
-        If used with mpi, beam.mpi_share_distributions must be called before.
+        If used with mpi, beam.mpi.share_distributions must be called before.
         
         The beam phasor is given at t=0 (synchronous particle) of the first 
         non empty bunch.
@@ -248,7 +263,7 @@ class CavityResonator():
                 else:
                     # no mpi -> get bunch profile for current bunch
                     if len(bunch) != 0:
-                        (bins, sorted_index, profile, center) = bunch.binning()
+                        (bins, sorted_index, profile, center) = bunch.binning(n_bin=self.n_bin)
                         bin_length = center[1]-center[0]
                         charge_per_mp = bunch.charge_per_mp
                         self.bunch_index = index
@@ -336,7 +351,7 @@ class CavityResonator():
                 else:
                     if i == 0:
                         # get bunch profile for current bunch
-                        (bins, sorted_index, profile, center) = bunch.binning()
+                        (bins, sorted_index, profile, center) = bunch.binning(n_bin=self.n_bin)
                         if j == 0:
                             self.profile_save = np.zeros((len(beam),len(profile),))
                             self.center_save = np.zeros((len(beam),len(center),))
@@ -416,7 +431,7 @@ class CavityResonator():
         sum_val = -2 * sum_tot * charge_per_mp * self.loss_factor
         self.beam_phasor += sum_val
         
-    def init_phasor(self, beam, n_bin=75):
+    def init_phasor(self, beam):
         """
         Initialize the beam phasor for a given beam distribution using an
         analytic formula [1].
@@ -426,8 +441,6 @@ class CavityResonator():
         Parameters
         ----------
         beam : Beam object
-        n_bin : int, optional
-            Number of bins to use. The default is 75.
             
         References
         ----------
@@ -439,7 +452,7 @@ class CavityResonator():
         if self.tracking is False:
             self.init_tracking(beam)
         
-        N = n_bin-1
+        N = self.n_bin - 1
         delta = (self.wr - self.m*self.ring.omega1)
         n_turn = int(self.filling_time/self.ring.T0*10)
         
@@ -453,13 +466,13 @@ class CavityResonator():
         for j, bunch in enumerate(beam.not_empty):
             index = self.valid_bunch_index[j]
             if beam.mpi_switch:
-                beam.mpi.share_distributions(beam, n_bin=n_bin)
+                beam.mpi.share_distributions(beam, n_bin=self.n_bin)
                 center[:,index] = beam.mpi.tau_center[j]
                 profile[:,index] = beam.mpi.tau_profile[j]
                 bin_length[index] = center[1, index]-center[0, index]
                 charge_per_mp[index] = beam.mpi.charge_per_mp_all[j]
             else:
-                (bins, sorted_index, profile[:, index], center[:, index]) = bunch.binning(n_bin=n_bin)
+                (bins, sorted_index, profile[:, index], center[:, index]) = bunch.binning(n_bin=self.n_bin)
                 bin_length[index] = center[1, index]-center[0, index]
                 charge_per_mp[index] = bunch.charge_per_mp
             T[index] -= (center[-1, index] + bin_length[index]/2)

mbtrack2/utilities/beamloading.py

@@ -10,9 +10,9 @@ from scipy.constants import c
 from scipy.integrate import quad
 
 class BeamLoadingEquilibrium():
-    """Class used to compute beam equilibrium profile and stability for a given
-    storage ring and a list of RF cavities of any harmonic. The class assumes
-    an uniform filling of the storage ring. Based on an extension of [1].
+    """Class used to compute beam equilibrium profile for a given storage ring 
+    and a list of RF cavities of any harmonic. The class assumes an uniform 
+    filling of the storage ring. Based on an extension of [1].
 
     [1] Venturini, M. (2018). Passive higher-harmonic rf cavities with general
     settings and multibunch instabilities in electron storage rings.