Add generator_phasor_record to CavityResonator

Also change CavityMonitor and plot_cavitydata to record and display generator_phasor_record data.

mbtrack2/tracking/monitors/monitors.py

@@ -1398,6 +1398,7 @@ class CavityMonitor(Monitor):
         group_name = cavity_name
         dict_buffer = {"cavity_phasor_record":(ring.h, buffer_size,),
                        "beam_phasor_record":(ring.h, buffer_size,),
+                       "generator_phasor_record":(ring.h, buffer_size,),
                        "detune":(buffer_size,),
                        "psi":(buffer_size,),
                        "Vg":(buffer_size,),
@@ -1408,6 +1409,7 @@ class CavityMonitor(Monitor):
                        "QL":(buffer_size,)}
         dict_file = {"cavity_phasor_record":(ring.h, total_size,),
                      "beam_phasor_record":(ring.h, total_size,),
+                     "generator_phasor_record":(ring.h, total_size,),
                      "detune":(total_size,),
                      "psi":(total_size,),
                      "Vg":(total_size,),
@@ -1418,6 +1420,7 @@ class CavityMonitor(Monitor):
                      "QL":(total_size,)}
         dict_dtype = {"cavity_phasor_record":complex,
                       "beam_phasor_record":complex,
+                      "generator_phasor_record":complex,
                       "detune":float,
                       "psi":float,
                       "Vg":float,

mbtrack2/tracking/monitors/plotting.py

@@ -1087,7 +1087,7 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
     cavity_name : str
         Name of the CavityResonator object.
     phasor : str, optional
-        Type of the phasor to plot. Can be "beam" or "cavity".
+        Type of the phasor to plot. Can be "beam", "cavity" or "generator".
     plot_type : str, optional
         Type of plot:
             - "bunch" plots the phasor voltage and angle versus time for a 
@@ -1120,13 +1120,14 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
     
     time = np.array(cavity_data["time"])
     
-    ph = {"cavity":0, "beam":1}
-    labels = ["Cavity", "Beam"]
+    ph = {"cavity":0, "beam":1, "generator":2}
+    labels = ["Cavity", "Beam", "Generator"]
     
     if plot_type == "bunch":
     
         data = [cavity_data["cavity_phasor_record"][bunch_number,:], 
-                cavity_data["beam_phasor_record"][bunch_number,:]]
+                cavity_data["beam_phasor_record"][bunch_number,:],
+                cavity_data["generator_phasor_record"][bunch_number,:]]
 
         ylabel1 = labels[ph[phasor]] + " voltage [MV]"
         ylabel2 = labels[ph[phasor]] + " phase [rad]"
@@ -1150,10 +1151,9 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         index = np.array(time) == turn
         if (index.size == 0):
             raise ValueError("Turn is not valid.")
-        ph = {"cavity":0, "beam":1}
-        data = [np.array(cavity_data["cavity_phasor_record"])[:,index], 
-                np.array(cavity_data["beam_phasor_record"])[:,index]]
-        labels = ["Cavity", "Beam"]
+        data = [cavity_data["cavity_phasor_record"][:,index], 
+                cavity_data["beam_phasor_record"][:,index],
+                cavity_data["generator_phasor_record"][:,index]]
         
         h=len(data[0])
         x=np.arange(h)

mbtrack2/tracking/rf.py

@@ -93,6 +93,8 @@ class CavityResonator():
         Last beam phasor value of each bunch in [V].
     generator_phasor : complex
         Generator phasor in [V].
+    generator_phasor_record : array of complex
+        Last generator phasor value of each bunch in [V].
     cavity_phasor : complex
         Cavity phasor in [V].
     cavity_phasor_record : array of complex
@@ -182,7 +184,7 @@ class CavityResonator():
     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 
+    [2] Naoto Yamamoto, Alexis Gamelin, and Ryutaro Nagaoka. "Investigation 
     of Longitudinal Beam Dynamics With Harmonic Cavities by Using the Code 
     Mbtrack." IPAC’19, Melbourne, Australia, 2019.
     """
@@ -202,6 +204,7 @@ class CavityResonator():
         self.theta = theta
         self.beam_phasor = np.zeros(1, dtype=complex)
         self.beam_phasor_record = np.zeros((self.ring.h), dtype=complex)
+        self.generator_phasor_record = np.zeros((self.ring.h), dtype=complex)
         self.tracking = False
         self.Vg = 0
         self.theta_g = 0
@@ -295,7 +298,7 @@ class CavityResonator():
                         
                         ind = (sorted_index == i)
                         phase = self.m * self.ring.omega1 * (center0 + self.ring.T1* (index + self.ring.h * self.nturn))
-                        Vgene = self.Vg*np.cos(phase + self.theta_g)
+                        Vgene = np.real(self.generator_phasor_record[index]*np.exp(1j*phase))
                         Vbeam = np.real(self.beam_phasor)
                         Vtot = Vgene + Vbeam - charge_per_mp*self.loss_factor*mp_per_bin
                         energy_change[ind] = Vtot / self.ring.E0
@@ -515,7 +518,7 @@ class CavityResonator():
     @property
     def cavity_phasor_record(self):
         """Last cavity phasor value of each bunch in [V]"""
-        return self.generator_phasor + self.beam_phasor_record
+        return self.generator_phasor_record + self.beam_phasor_record
     
     @property
     def cavity_voltage(self):
@@ -786,6 +789,8 @@ class CavityResonator():
         self.Vg = self.Vgr*np.cos(self.psi)
         # Generator phase [rad]
         self.theta_g = self.theta_gr + self.psi
+        # Set generator_phasor_record
+        self.generator_phasor_record = np.ones(self.ring.h)*self.generator_phasor
         
     def plot_phasor(self, I0):
         """