From bc5e46cf4a333bd709d897dd3cb61d5bc1a2310e Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Sat, 9 Sep 2023 08:53:07 +0200
Subject: [PATCH] Add monitoring to ig_phasor_record

Add new ig_phasor_record attribute to CavityResonator, non zero only when specific FB type are used.
Monitor ig_phasor_record with CavityMonitor.
Modify plot_cavitydata to allowing plotting of ig_phasor_record.
---
 mbtrack2/tracking/monitors/monitors.py |  3 ++
 mbtrack2/tracking/monitors/plotting.py | 43 ++++++++++++++++----------
 mbtrack2/tracking/rf.py                | 11 +++++++
 3 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 9bf28e4..e745696 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/mbtrack2/tracking/monitors/monitors.py
@@ -1399,6 +1399,7 @@ class CavityMonitor(Monitor):
         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,),
@@ -1410,6 +1411,7 @@ class CavityMonitor(Monitor):
         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,),
@@ -1421,6 +1423,7 @@ class CavityMonitor(Monitor):
         dict_dtype = {"cavity_phasor_record":complex,
                       "beam_phasor_record":complex,
                       "generator_phasor_record":complex,
+                      "ig_phasor_record":complex,
                       "detune":float,
                       "psi":float,
                       "Vg":float,
diff --git a/mbtrack2/tracking/monitors/plotting.py b/mbtrack2/tracking/monitors/plotting.py
index 3c34309..3ac8e82 100644
--- a/mbtrack2/tracking/monitors/plotting.py
+++ b/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", "cavity" or "generator".
+        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,27 +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, "generator":2}
-    labels = ["Cavity", "Beam", "Generator"]
+    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["generator_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)
@@ -1153,17 +1156,18 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
             raise ValueError("Turn is not valid.")
         data = [cavity_data["cavity_phasor_record"][:,index], 
                 cavity_data["beam_phasor_record"][:,index],
-                cavity_data["generator_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
             
diff --git a/mbtrack2/tracking/rf.py b/mbtrack2/tracking/rf.py
index 62a8b36..23c2de6 100644
--- a/mbtrack2/tracking/rf.py
+++ b/mbtrack2/tracking/rf.py
@@ -102,6 +102,9 @@ class CavityResonator():
         Cavity phasor in [V].
     cavity_phasor_record : array of complex
         Last cavity phasor value of each bunch in [V].
+    ig_phasor_record : array of complex
+        Last current generator phasor of each bunch in [A].
+        Only used for some feedback types.
     cavity_voltage : float
         Cavity total voltage in [V].
     cavity_phase : float
@@ -527,6 +530,14 @@ class CavityResonator():
         """Last cavity phasor value of each bunch in [V]"""
         return self.generator_phasor_record + self.beam_phasor_record
     
+    @property
+    def ig_phasor_record(self):
+        """Last current generator phasor of each bunch in [A]"""
+        for FB in self.feedback:
+            if isinstance(FB, (ProportionalIntegralLoop, DirectFeedback)):
+                return FB.ig_phasor_record
+        return np.zeros(self.ring.h)
+    
     @property
     def cavity_voltage(self):
         """Cavity total voltage in [V]"""
-- 
GitLab