From 8aa00b8a55f7fcaaa539210d9d839b5ab3d5f94e Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Fri, 29 Sep 2023 13:52:42 +0200
Subject: [PATCH] Various improvements related to CavityResonator

Add update_feedback method to CavityResonator to force FB update when needed (when QL or tuning is changed).
The feedback reference targets (setpoints) for ProportionalIntegralLoop is now CavityResonator Vc and theta as for ProportionalLoop.
Add an init_FFconst method to ProportionalIntegralLoop which is called at initialization only (not any more after each setpoint change) and needs to be called manually if we want to change the feedfoward constant during tracking.
Some simplifications on the expressions in ProportionalIntegralLoop.
CavityMonitor now records Vc and theta.
Add option "mean" to plot_cavitydata to plot the mean phasor voltage and angle versus time for non empty bunches (now the default behaviour).
Add optinal parameter show_objective to cavitydata, if True shows the Cavity voltage and phase objective value in plots.
---
 mbtrack2/tracking/monitors/monitors.py | 12 +++-
 mbtrack2/tracking/monitors/plotting.py | 39 ++++++++---
 mbtrack2/tracking/rf.py                | 90 ++++++++++++--------------
 3 files changed, 80 insertions(+), 61 deletions(-)

diff --git a/mbtrack2/tracking/monitors/monitors.py b/mbtrack2/tracking/monitors/monitors.py
index 5d55db8..f8641fe 100644
--- a/mbtrack2/tracking/monitors/monitors.py
+++ b/mbtrack2/tracking/monitors/monitors.py
@@ -1407,7 +1407,9 @@ class CavityMonitor(Monitor):
                        "Pg":(buffer_size,),
                        "Rs":(buffer_size,),
                        "Q":(buffer_size,),
-                       "QL":(buffer_size,)}
+                       "QL":(buffer_size,),
+                       "Vc":(buffer_size,),
+                       "theta":(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,),
@@ -1419,7 +1421,9 @@ class CavityMonitor(Monitor):
                      "Pg":(total_size,),
                      "Rs":(total_size,),
                      "Q":(total_size,),
-                     "QL":(total_size,)}
+                     "QL":(total_size,),
+                     "Vc":(total_size,),
+                     "theta":(total_size,)}
         dict_dtype = {"cavity_phasor_record":complex,
                       "beam_phasor_record":complex,
                       "generator_phasor_record":complex,
@@ -1431,7 +1435,9 @@ class CavityMonitor(Monitor):
                       "Pg":float,
                       "Rs":float,
                       "Q":float,
-                      "QL":float}
+                      "QL":float,
+                      "Vc":float,
+                      "theta":float}
         
         self.monitor_init(group_name, save_every, buffer_size, total_size,
                           dict_buffer, dict_file, file_name, mpi_mode, 
diff --git a/mbtrack2/tracking/monitors/plotting.py b/mbtrack2/tracking/monitors/plotting.py
index 8afe578..8adbae4 100644
--- a/mbtrack2/tracking/monitors/plotting.py
+++ b/mbtrack2/tracking/monitors/plotting.py
@@ -1076,7 +1076,8 @@ def streak_beamspectrum(filename, dim="tau", f0=None, log_scale=True, fmin=None,
     return fig
 
 def plot_cavitydata(filename, cavity_name, phasor="cavity", 
-                    plot_type="bunch", bunch_number=0, turn=None, cm_lim=None):
+                    plot_type="mean", bunch_number=0, turn=None, cm_lim=None,
+                    show_objective=False):
     """
     Plot data recorded by CavityMonitor.
 
@@ -1091,6 +1092,8 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         "ig".
     plot_type : str, optional
         Type of plot:
+            - "mean" plots the mean phasor voltage and angle versus time for 
+            non empty bunches.
             - "bunch" plots the phasor voltage and angle versus time for a 
             given bunch.
             - "turn" plots the phasor voltage and ange versus bunch index for
@@ -1108,6 +1111,10 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         Turn to plot. The default is None.
     cm_lim : list [vmin, vmax], optional
         Colormap scale for the "streak" plots.
+    show_objective : bool, optional
+        If True, shows the Cavity voltage and phase objective value (Vc and 
+        theta) in plot type "mean" and "bunch" for "cavity" phasor.
+        Default is False.
 
     Returns
     -------
@@ -1125,13 +1132,23 @@ 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":
+    if plot_type == "bunch" or "mean":
     
-        data = [cavity_data["cavity_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,:]]
-
+        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["ig_phasor_record"][bunch_number,:]]
+        elif plot_type == "mean":
+            try:
+                bunch_index = (file["Beam"]["current"][:,0] != 0).nonzero()[0]
+            except:
+                print("Beam monitor is needed to show mean voltage.")
+            data = [np.mean(cavity_data["cavity_phasor_record"][bunch_index,:],0), 
+                    np.mean(cavity_data["beam_phasor_record"][bunch_index,:],0),
+                    np.mean(cavity_data["generator_phasor_record"][bunch_index,:],0),
+                    np.mean(cavity_data["ig_phasor_record"][bunch_index,:],0)]
+            
         ylabel1 = labels[ph[phasor]] + units[ph[phasor]]
         ylabel2 = labels[ph[phasor]] + " phase [rad]"
         
@@ -1139,11 +1156,17 @@ def plot_cavitydata(filename, cavity_name, phasor="cavity",
         twin = ax.twinx()
         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)
+        if show_objective:
+            o1, = ax.plot(time, np.array(cavity_data["Vc"])*1e-6, "r--", label="Cavity voltage objective value [MV]")
+            o2, = twin.plot(time, np.array(cavity_data["theta"]), "b--", label="Cavity phase objective value [rad]")
         ax.set_xlabel("Turn number")
         ax.set_ylabel(ylabel1)
         twin.set_ylabel(ylabel2)
         
-        plots = [p1, p2]
+        if show_objective:
+            plots = [p1, o1, p2, o2]
+        else:
+            plots = [p1, p2]
         ax.legend(handles=plots, loc="best")
         
         ax.yaxis.label.set_color("r")
diff --git a/mbtrack2/tracking/rf.py b/mbtrack2/tracking/rf.py
index dbf93a6..7afa2dc 100644
--- a/mbtrack2/tracking/rf.py
+++ b/mbtrack2/tracking/rf.py
@@ -79,9 +79,9 @@ class CavityResonator():
     Ncav : int, optional
         Number of cavities.
     Vc : float, optinal
-        Total cavity voltage in [V].
+        Total cavity voltage objective value in [V].
     theta : float, optional
-        Total cavity phase in [rad].
+        Total cavity phase objective value in [rad].
     n_bin : int, optional
         Number of bins used for the beam loading computation. 
         Only used if MPI is not used, otherwise n_bin must be specified in the 
@@ -183,6 +183,8 @@ class CavityResonator():
         Return the second derivative of total RF voltage.
     deltaVRF(z, I0)
         Return the generator voltage minus beam loading voltage.
+    update_feedback()
+        Force feedback update from current CavityResonator parameters.
     
     References
     ----------
@@ -197,6 +199,7 @@ class CavityResonator():
     def __init__(self, ring, m, Rs, Q, QL, detune, Ncav=1, Vc=0, theta=0, 
                  n_bin=75):
         self.ring = ring
+        self.feedback = []
         self.m = m
         self.Ncav = Ncav
         if Ncav != 1:
@@ -218,7 +221,6 @@ class CavityResonator():
         self.theta_gr = 0
         self.Pg = 0
         self.n_bin = int(n_bin)
-        self.feedback = []
         
     def init_tracking(self, beam):
         """
@@ -623,6 +625,7 @@ class CavityResonator():
     def QL(self, value):
         self._QL = value
         self._beta = self.Q/self.QL - 1
+        self.update_feedback()
 
     @property
     def beta(self):
@@ -645,6 +648,7 @@ class CavityResonator():
         self._wr = self.fr*2*np.pi
         self._psi = np.arctan(self.QL*(self.fr/(self.m*self.ring.f1) -
                                        (self.m*self.ring.f1)/self.fr))
+        self.update_feedback()
 
     @property
     def fr(self):
@@ -913,6 +917,12 @@ class CavityResonator():
     def deltaVRF(self, z, I0, F = 1, PHI = 0):
         """Return the generator voltage minus beam loading voltage taking into account form factor amplitude F and form factor phase PHI"""
         return -1*self.Vg*(self.ring.k1*self.m)**2*np.cos(self.ring.k1*self.m*z + self.theta_g) - self.Vb(I0)*F*(self.ring.k1*self.m)**2*np.cos(self.ring.k1*self.m*z + self.psi - PHI)
+
+    def update_feedback(self):
+        """Force feedback update from current CavityResonator parameters."""
+        for FB in self.feedback:
+            if isinstance(FB, (ProportionalIntegralLoop, DirectFeedback)):
+                FB.init_Ig2Vg_matrix()
     
 class ProportionalLoop():
     """
@@ -1030,21 +1040,24 @@ class ProportionalIntegralLoop():
     Proportional Integral (PI) loop to control a CavityResonator amplitude and 
     phase via generator current (ig) [1,2].
     
+    Feedback reference targets (setpoints) are cav_res.Vc and cav_res.theta.
+    
     The ProportionalIntegralLoop should be initialized only after generator 
     parameters are set.
-    The method init_Ig2Vg_matrix must be called after each parameter change of 
-    cav_res.
     
     The loop must be added to the CavityResonator object using:
         cav_res.feedback.append(loop)
+        
+    When the reference target is changed, it might be needed to re-initialize 
+    the feedforward constant by calling init_FFconst().
 
     Parameters
     ----------
     ring : Synchrotron object
     cav_res : CavityResonator object
         CavityResonator in which the loop will be added. 
-    gain : float or float list
-        Pgain and Igain of the feedback system
+    gain : float list
+        Proportional gain (Pgain) and integral gain (Igain) of the feedback system.
         For IQ feedback, same gain set is used for I and Q.
         In case of normal conducting cavity (QL~1e4),
             the Pgain of ~1.0 and Igain of ~1e4(5) are usually used.
@@ -1062,13 +1075,10 @@ class ProportionalIntegralLoop():
     delay : int
         Loop delay of the PI feedback system. 
         Units are in bucket numbers.
-    Vref : array or complex, optional
-        Target (set point) value of the feedback.
-        If None, cav_res.Vc and cav_res.theta are set as the reference.
-        Default is None.
     FF : bool, optional
+        Boolean switch to use feedforward constant.
         True is recommended to prevent a Vc drop in the beginning of the tracking.
-        In case of small Pgain (QL ~ 1e4), Vc drop may cause baem loss due to static Robinson.
+        In case of small Pgain (QL ~ 1e4), Vc drop may cause beam loss due to static Robinson.
         Default is True.
         
     Methods
@@ -1077,6 +1087,8 @@ class ProportionalIntegralLoop():
         Tracking method for the Cavity PI control feedback.
     init_Ig2Vg_matrix()
         Initialize matrix for Ig2Vg_matrix.
+    init_FFconst()
+        Initialize feedforward constant.
     Ig2Vg_matrix()
         Return Vg from Ig using matrix formalism.
     Ig2Vg()
@@ -1118,15 +1130,12 @@ class ProportionalIntegralLoop():
     of synchrotron light sources. PRAB, 21(1), 012001.
     
     """
-    def __init__(self, ring, cav_res, gain, sample_num, every, delay, 
-                 Vref=None, FF=True):
+    def __init__(self, ring, cav_res, gain, sample_num, every, delay, FF=True):
         self.ring = ring
         self.cav_res = cav_res
         self.Ig2Vg_mat = np.zeros((self.ring.h, self.ring.h), dtype=complex)
         self.gain = gain
         self.FF = FF
-        if not self.FF:
-            self.FFconst = 0
 
         if delay > 0:
             self.delay = int(delay)
@@ -1148,12 +1157,9 @@ class ProportionalIntegralLoop():
         self.diff_record = np.zeros(record_size, dtype=complex)
         self.I_record = 0+0j
 
-        if Vref is None:
-            self.Vref = self.cav_res.Vc*np.exp(1j*self.cav_res.theta)
-        else: 
-            self.Vref = Vref
-
         self.sample_list = range(0, self.ring.h, self.every)
+        
+        self.init_FFconst()
 
         # Pre caclulation for Ig2Vg
         self.init_Ig2Vg_matrix()
@@ -1179,8 +1185,8 @@ class ProportionalIntegralLoop():
             # Shift the record
             self.diff_record = np.roll(self.diff_record, 1)
             # 1) recording diff as a first item of the list
-            mean_vc = np.mean(vc_list[index:self.sample_num+index])*self.Rot
-            self.diff_record[0] = self.VrefRot - mean_vc
+            mean_vc = np.mean(vc_list[index:self.sample_num+index])*np.exp(-1j*self.cav_res.theta)
+            self.diff_record[0] = self.cav_res.Vc - mean_vc
         # update sample_list for next turn
         self.sample_list = range(index+self.every-self.ring.h,self.ring.h,self.every)
         # update vc_previous for next turn
@@ -1190,30 +1196,6 @@ class ProportionalIntegralLoop():
         
         if apply_changes:
             self.Ig2Vg()
-        
-    @property
-    def Vref(self):
-        """Return the target (set point) value of the feedback as a complex."""
-        return self._Vref
-
-    @Vref.setter
-    def Vref(self, value):
-        """
-        Set the target (set point) value of the feedback.
-
-        Vref can be specified as an IQ-complex style or array-like [Amp, Phase].
-        If not specified, cav_res.Vc and cav_res.theta are used.
-
-        """
-        if isinstance(value,complex):
-            self._Vref = value
-        else:
-            self._Vref = value[0]*np.exp(1j*value[1])
-        self.Rot = np.exp(-1j*np.angle(self._Vref))
-        self.VrefRot = self._Vref*self.Rot
-        if self.FF:
-            self.FFconst = np.mean(self.ig_phasor)
-        self.init_Ig2Vg_matrix()
             
     def init_Ig2Vg_matrix(self):
         """
@@ -1228,6 +1210,13 @@ class ProportionalIntegralLoop():
         for idx in np.arange(self.ring.h):
             self.Ig2Vg_mat[idx:,idx] = tempV[:self.ring.h-idx]
             
+    def init_FFconst(self):
+        """Initialize feedforward constant."""
+        if self.FF:
+            self.FFconst = np.mean(self.ig_phasor)
+        else:
+            self.FFconst = 0
+            
     def Ig2Vg_matrix(self):
         """
         Return Vg from Ig using matrix formalism.
@@ -1255,9 +1244,7 @@ class ProportionalIntegralLoop():
 
         Eq.25 of ref [2] assuming the dVg/dt = 0.
         """
-        fac = 1/self.cav_res.filling_time/self.cav_res.loss_factor # 1/R_L
-        ig = fac*Vg*( 1 - 1j*np.tan(self.cav_res.psi) )
-        return ig
+        return Vg * ( 1 - 1j*np.tan(self.cav_res.psi) ) / self.cav_res.RL
 
 class DirectFeedback(ProportionalIntegralLoop):
     """
@@ -1271,6 +1258,9 @@ class DirectFeedback(ProportionalIntegralLoop):
     the CavityResonator generator parameters should be set before 
     initialization.
     
+    The loop must be added to the CavityResonator object using:
+        cav_res.feedback.append(loop)
+    
     Parameters
     ----------
     DFB_gain : float
-- 
GitLab