Merge branch 'RF-FB2' into develop

mbtrack2/tracking/__init__.py

@@ -6,7 +6,11 @@ from mbtrack2.tracking.particles import (Electron,
                                          Particle)
 from mbtrack2.tracking.synchrotron import Synchrotron
 from mbtrack2.tracking.rf import (RFCavity, 
-                                  CavityResonator)
+                                  CavityResonator,
+                                  ProportionalLoop,
+                                  TunerLoop,
+                                  ProportionalIntegralLoop,
+                                  DirectFeedback)
 from mbtrack2.tracking.parallel import Mpi
 from mbtrack2.tracking.element import (Element, 
                                        LongitudinalMap, 

mbtrack2/tracking/monitors/monitors.py

@@ -1398,6 +1398,8 @@ 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,),
+                       "ig_phasor_record":(ring.h, buffer_size,),
                        "detune":(buffer_size,),
                        "psi":(buffer_size,),
                        "Vg":(buffer_size,),
@@ -1408,6 +1410,8 @@ 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,),
+                     "ig_phasor_record":(ring.h, total_size,),
                      "detune":(total_size,),
                      "psi":(total_size,),
                      "Vg":(total_size,),
@@ -1418,6 +1422,8 @@ class CavityMonitor(Monitor):
                      "QL":(total_size,)}
         dict_dtype = {"cavity_phasor_record":complex,
                       "beam_phasor_record":complex,
+                      "generator_phasor_record":complex,
+                      "ig_phasor_record":complex,
                       "detune":float,
                       "psi":float,
                       "Vg":float,

mbtrack2/tracking/monitors/plotting.py

@@ -1087,15 +1087,16 @@ 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", "generator" or 
+        "ig".
     plot_type : str, optional
         Type of plot:
             - "bunch" plots the phasor voltage and angle versus time for a 
             given bunch.
             - "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.
+            - "streak_amplitude" plots the phasor amplitude versus bunch index 
+            and time.
             - "streak_angle" plots the phasor angle versus bunch index and 
             time.
             - "detune" or "psi" plots the detuning or tuning angle versus time.
@@ -1114,26 +1115,29 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         Figure object with the plot on it.
 
     """
-    
     file = hp.File(filename, "r")
     cavity_data = file[cavity_name]
     
     time = np.array(cavity_data["time"])
     
-    ph = {"cavity":0, "beam":1}
-    labels = ["Cavity", "Beam"]
+    ph = {"cavity":0, "beam":1, "generator":2, "ig":3}
+    labels = ["Cavity", "Beam", "Generator", "Generator"]
+    units = [" voltage [MV]", " voltage [MV]", " voltage [MV]", " current [A]"]
+    units_val = [1e-6, 1e-6, 1e-6, 1]
     
     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,:],
+                cavity_data["ig_phasor_record"][bunch_number,:]]
 
-        ylabel1 = labels[ph[phasor]] + " voltage [MV]"
+        ylabel1 = labels[ph[phasor]] + units[ph[phasor]]
         ylabel2 = labels[ph[phasor]] + " phase [rad]"
         
         fig, ax = plt.subplots()
         twin = ax.twinx()
-        p1, = ax.plot(time, np.abs(data[ph[phasor]])*1e-6, color="r",label=ylabel1)
+        p1, = ax.plot(time, np.abs(data[ph[phasor]])*units_val[ph[phasor]], color="r",label=ylabel1)
         p2, = twin.plot(time, np.angle(data[ph[phasor]]), color="b", label=ylabel2)
         ax.set_xlabel("Turn number")
         ax.set_ylabel(ylabel1)
@@ -1150,20 +1154,20 @@ 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],
+                cavity_data["ig_phasor_record"][:,index]]
         
         h=len(data[0])
         x=np.arange(h)
 
-        ylabel1 = labels[ph[phasor]] + " voltage [MV]"
+        ylabel1 = labels[ph[phasor]] + units[ph[phasor]]
         ylabel2 = labels[ph[phasor]] + " phase [rad]"
         
         fig, ax = plt.subplots()
         twin = ax.twinx()
-        p1, = ax.plot(x, np.abs(data[ph[phasor]])*1e-6, color="r",label=ylabel1)
+        p1, = ax.plot(x, np.abs(data[ph[phasor]])*units_val[ph[phasor]], color="r",label=ylabel1)
         p2, = twin.plot(x, np.angle(data[ph[phasor]]), color="b", label=ylabel2)
         ax.set_xlabel("Bunch index")
         ax.set_ylabel(ylabel1)
@@ -1175,14 +1179,19 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         ax.yaxis.label.set_color("r")
         twin.yaxis.label.set_color("b")
         
-    if plot_type == "streak_volt" or plot_type == "streak_phase":
+    if plot_type == "streak_amplitude" or plot_type == "streak_phase":
+        
+        data = [cavity_data["cavity_phasor_record"][:,:], 
+                cavity_data["beam_phasor_record"][:,:],
+                cavity_data["generator_phasor_record"][:,:],
+                cavity_data["ig_phasor_record"][:,:]]
         
-        if plot_type == "streak_volt":
-            data = np.transpose(np.abs(cavity_data["cavity_phasor_record"][:,:])*1e-6)
-            ylabel = labels[ph[phasor]] + " voltage [MV]"
+        if plot_type == "streak_amplitude":
+            data = np.transpose(np.abs(data[ph[phasor]])*units_val[ph[phasor]])
+            ylabel = labels[ph[phasor]] + units[ph[phasor]]
             cmap = mpl.cm.coolwarm # diverging
         elif plot_type == "streak_phase":
-            data = np.transpose(np.angle(cavity_data["cavity_phasor_record"][:,:]))
+            data = np.transpose(np.angle(data[ph[phasor]]))
             ylabel = labels[ph[phasor]] + " phase [rad]"
             cmap = mpl.cm.coolwarm # diverging