Add sample_voltage method to CavityResonator

Bug fix in plot_cavitydata.
New method to sample the voltage seen by a zero charge particle during an RF period.

mbtrack2/tracking/monitors/plotting.py

@@ -1132,7 +1132,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
     units = [" voltage [MV]", " voltage [MV]", " voltage [MV]", " current [A]"]
     units_val = [1e-6, 1e-6, 1e-6, 1]
     
-    if plot_type == "bunch" or "mean":
+    if plot_type == "bunch" or plot_type == "mean":
     
         if plot_type == "bunch":
             data = [cavity_data["cavity_phasor_record"][bunch_number,:], 

mbtrack2/tracking/rf.py

@@ -185,6 +185,8 @@ class CavityResonator():
         Return the generator voltage minus beam loading voltage.
     update_feedback()
         Force feedback update from current CavityResonator parameters.
+    sample_voltage()
+        Sample the voltage seen by a zero charge particle during an RF period.
     
     References
     ----------
@@ -924,6 +926,53 @@ class CavityResonator():
             if isinstance(FB, (ProportionalIntegralLoop, DirectFeedback)):
                 FB.init_Ig2Vg_matrix()
             
+    def sample_voltage(self, n_points=1e4, index=0):
+        """
+        Sample the voltage seen by a zero charge particle during an RF period.
+
+        Parameters
+        ----------
+        n_points : int or float, optional
+            Number of sample points. The default is 1e4.
+        index : int, optional
+            RF bucket number to sample. Be carful if index > 0 as no new beam 
+            loading is not taken into account here.
+            The default is 0.
+
+        Returns
+        -------
+        pos : array of float
+            Array of position from -T1/2 to T1/2.
+        voltage_rec : array of float
+            Recoring of the voltage.
+
+        """
+        # Init
+        n_points = int(n_points)
+        index = 0
+        voltage_rec = np.zeros(n_points)
+        pos = np.linspace(-self.ring.T1/2, self.ring.T1/2, n_points)
+        DeltaT = self.ring.T1/(n_points-1)
+        
+        # From t=0 of first non empty bunch to -T1/2
+        self.phasor_decay(-self.ring.T1/2 + index*self.ring.T1, 
+                          ref_frame="beam")
+
+        # Goes from (-T1/2) to (T1/2 + DeltaT) in n_points steps
+        for i in range(n_points):
+            phase = self.m * self.ring.omega1 * (pos[i] + self.ring.T1* (index + self.ring.h * self.nturn))
+            Vgene = np.real(self.generator_phasor_record[index]*np.exp(1j*phase))
+            Vbeam = np.real(self.beam_phasor)
+            Vtot = Vgene + Vbeam
+            voltage_rec[i] = Vtot
+            self.phasor_decay(DeltaT, ref_frame="beam")
+
+        # Get back to t=0
+        self.phasor_decay(-DeltaT*n_points + self.ring.T1/2 - index*self.ring.T1, 
+                          ref_frame="beam")
+        
+        return pos, voltage_rec
+
 class ProportionalLoop():
     """
     Proportional feedback loop to control a CavityResonator amplitude and phase.