Improve plot_cavitydata

Allow to plot detuning, tuning angle and powers from tracking.

tracking/monitors/plotting.py

@@ -1072,9 +1072,12 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
             - "turn" plots the phasor voltage and ange versus bunch index for
             a given turn.
             - "streak_volt" plots the phasor voltage versus bunch index and 
-            time
+            time.
             - "streak_angle" plots the phasor angle versus bunch index and 
-            time
+            time.
+            - "detune" or "psi" plots the detuning or tuning angle versus time.
+            - "power" plots the generator, cavity, beam and reflected power
+            versus time.
     bunch_number : int, optional
         Bunch number to select. The default is 0.
     turn : int, optional
@@ -1166,7 +1169,45 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         ax.set_xlabel("Bunch index")
         ax.set_ylabel("Number of turns")
         cbar = fig.colorbar(c, ax=ax)
-        cbar.set_label(ylabel) 
-    
+        cbar.set_label(ylabel)
+        
+    if plot_type == "detune" or plot_type == "psi":
+        
+        fig, ax = plt.subplots()
+        if plot_type == "detune":
+            data = np.array(cavity_data["detune"])*1e-3
+            ylabel = r"Detuning $\Delta f$ [kHz]"
+        elif plot_type == "psi":
+            data = np.array(cavity_data["psi"])
+            ylabel = r"Tuning angle $\psi$"
+            
+        ax.plot(time, data)
+        ax.set_xlabel("Number of turns")
+        ax.set_ylabel(ylabel)
+        
+    if plot_type == "power":
+        Vc = np.mean(np.abs(cavity_data["cavity_phasor_record"]),0)
+        theta = np.mean(np.angle(cavity_data["cavity_phasor_record"]),0)
+        try:
+            bunch_index = (file["Beam"]["current"][:,0] != 0).nonzero()[0]
+            I0 = np.nansum(file["Beam"]["current"][bunch_index,:],0)
+        except:
+            print("Beam monitor is needed to compute power.")
+            
+        Rs = np.array(cavity_data["Rs"])
+        Pc = Vc**2 / (2 * Rs)
+        Pb = I0 * Vc * np.cos(theta)
+        Pg = np.array(cavity_data["Pg"])
+        Pr = Pg - Pb - Pc
+        
+        fig, ax = plt.subplots()
+        ax.plot(time, Pg*1e-3, label="Generator power $P_g$ [kW]")
+        ax.plot(time, Pb*1e-3, label="Beam power $P_b$ [kW]")
+        ax.plot(time, Pc*1e-3, label="Dissipated cavity power $P_c$ [kW]")
+        ax.plot(time, Pr*1e-3, label="Reflected power $P_r$ [kW]")
+        ax.set_xlabel("Number of turns")
+        ax.set_ylabel("Power [kW]")
+        plt.legend()
+        
     file.close()
     return fig