From 7d37efa062a6ed39bb1c2a1ac41fd8b72db4303b Mon Sep 17 00:00:00 2001
From: AntoninSauret <antonin.sauret@synchrotron-soleil.fr>
Date: Mon, 27 Jun 2022 11:47:37 +0200
Subject: [PATCH] Feedback class vf

---
 mbtrack2/tracking/FeedBack.py       | 354 +++++++++++++++++++---------
 mbtrack2/tracking/SimpleFeedBack.py |  19 +-
 2 files changed, 251 insertions(+), 122 deletions(-)

diff --git a/mbtrack2/tracking/FeedBack.py b/mbtrack2/tracking/FeedBack.py
index 14a41e0..76cd20e 100644
--- a/mbtrack2/tracking/FeedBack.py
+++ b/mbtrack2/tracking/FeedBack.py
@@ -1,28 +1,106 @@
 # -*- coding: utf-8 -*-
 """
-Created on Tue May  3 09:36:24 2022
-
-@author: sauret
+This module defines both simple and FIR based feedback for tracking.
 """
 from mbtrack2.tracking.element import Element
-import numpy as np
+import numpy as np 
+import matplotlib.pyplot as plt
 
-class FeedBackTransverse(Element) : 
+class SimpleFB(Element) : 
     """ 
-    Transverse feedback system to counter transverse instabilities.
+   Simple feedback system proceeding bunch by bunch correction to counter beam 
+   instabilities
+   
+   Parameters
+   ----------
+   ring : Synchrotron object
+   damping_time : array of shape (3,)
+       Contains the desired damping time in seconds for each plane. Element 0 of
+       the array defines the horizontal plane as 1 and 2 define the vertical 
+       and longitudinal plane respectively. 
+   phase_diff : array of shape (3,)
+       Contains the desired phase in degree for each plane. Element 0 of the 
+       array defines the horizontal plane as 1 and 2 define the vertical and 
+       longitudinal plane respectively.
+   plane : bool array of shape (3,), optional
+       Allow to choose on which plane the feedback is active. State True to
+       turn on the feedback or False otherwise. Element 0 of the array defines
+       the horizontal plane as 1 and 2 define the vertical and longitudinal plane
+       respectively.
+    """
+    
+    def __init__(self, ring, damping_time, phase_diff, plane = np.zeros((3,), dtype= bool)):
+        self.ring = ring
+        self.damping_time = damping_time
+        self.phase_diff = phase_diff
+        self.plane = plane 
+        
+    @Element.parallel 
+    def track(self, bunch):
+        """
+        Tracking method for the feedback system
+        No bunch to bunch interaction, so written for Bunch object and
+        @Element.parallel is used to handle Beam object.
+        
+        Parameters
+        ----------
+        bunch : Bunch or Beam object
+        """
+        if(self.plane[0] == True):
+            bunch["xp"] -= (2*self.ring.T0/self.damping_time[0])*np.sin(self.phase_diff[0])*bunch.mean[1]
+            
+        if(self.plane[1] == True):
+            bunch["yp"] -= (2*self.ring.T0/self.damping_time[1])*np.sin(self.phase_diff[1])*bunch.mean[3]
+            
+        if(self.plane[2] == True):
+            bunch["delta"] -= (2*self.ring.T0/self.damping_time[2])*np.sin(self.phase_diff[1])*bunch.mean[5]
+
+class FIR_Feedback(Element) : 
+    """ 
+    Feedback system based on FIR filters, proceeding bunch by bunch correction 
+    to counter beam instabilities.
     
     Parameters
     ----------
     ring : Synchrotron object
-    turn_delay : int, optional
-        Number of turn delay before applying kick.
-    tap_number : int, optional
-        Number of tap for the FIR filters.
-    ampl : float, optional
-        Gain of the FIR filters.
-    phase : float, optional
-        Phase of the FIR filters in deg.
-            
+    tune : float array of shape (3,)
+        Betatron tune for horizontal and vertical plane in the first and second
+        elements of the array. Synchrotron tune for the third element.
+    turn_delay : int array of shape (3,)
+        Number of turn delay before applying kick. Element 0 of the array 
+        defines the horizontal plane as 1 and 2 define the vertical and 
+        longitudinal plane respectively.
+    tap_number : int array of shape (3,)
+        Number of tap for the FIR filters. Element 0 of the array defines
+        the horizontal plane as 1 and 2 define the vertical and longitudinal
+        plane respectively.
+    gain : float array of shape (3,)
+        Gain of the FIR filters. Element 0 of the array defines
+        the horizontal plane as 1 and 2 define the vertical and longitudinal
+        plane respectively.
+    phase : float array of shape (3,)
+        Phase of the FIR filters in degree. Element 0 of the array defines
+        the horizontal plane as 1 and 2 define the vertical and longitudinal
+        plane respectively.
+    bpm_error : float array of shape (3,), optional
+        Stores the BPM error [m] for horizontal and vertical, [s] for longitudinal. 
+        First, second and third element corespond to the horizontal, vertical 
+        and longitudinal plane respectively.
+    max_kick : float array of shape (3,), optional
+        Stores the maximum kick limitation. First, second and third 
+        element correspond to the horizontal, vertical and longitudinal plane 
+        respectively.
+    max_kick_state : bool array of shape (3,), optional
+        Allow to choose on which plane the max kick limitation is active. State 
+        True to turn on the max kick limitation or False otherwise. Element 0 
+        of the array defines the horizontal plane as 1 and 2 define the vertical 
+        and longitudinal plane respectively.
+    plane : bool array of shape (3,), optional
+        Allow to choose on which plane the feedback is active. State True to
+        turn on the feedback or False otherwise. Element 0 of the array defines
+        the horizontal plane as 1 and 2 define the vertical and longitudinal
+        plane respectively.
+        
     Attributes
     ----------
      x_pos : array
@@ -31,112 +109,135 @@ class FeedBackTransverse(Element) :
      y_pos : array
          Retrieves and stores the position of the bunch, measured by the BPM, 
          in vertical plane.
+     tau_pos : array
+         Retrieves and stores the position of the bunch, measured by the BPM, 
+         in longitudinal plane.
      kick_x : array
          Stores the numerical values of the kicks for the horizontal plane.
      kick_y : array
          Stores the numerical values of the kicks for the vertical plane.
+     kick_tau : array
+         Stores the numerical values of the kicks for the longitudinal plane.
      coef_x : array
          Stores the coefficients of the FIR filter for the horizontal plane.
      coef_y : array
-         Stores the coefficients of the FIR filter for the vertical plane. 
+         Stores the coefficients of the FIR filter for the vertical plane.
+     coef_tau : array
+         Stores the coefficients of the FIR filter for the longitudinal plane.
+         
+    Methods
+    -------
+    get_fir(tap_number, tune, phase, turn_delay, gain)
+         Initialize the FIR filter for the desired plane and return an array
+         containing the FIR coefficients.
+    plot_fir(plane)
+         Plot the gain and the phase of the FIR filter selected.
+      
+    References
+    ----------
+    [1] T.Nakamura, S.Daté, K. Kobayashi, T. Ohshima. Proceedings of EPAC 
+    2004. Transverse bunch by bunch feedback system for the Spring-8 
+    storage ring.
     """
     
-    def __init__(self, ring, turn_delay=1, tap_number=5, ampl =1e-1, phase = -90):
+    def __init__(self, ring, tune, turn_delay, tap_number, gain, phase, bpm_error = np.zeros((3,)), 
+                 max_kick = np.zeros((3,)),plane = np.zeros((3,), dtype = bool)):
         
         self.ring = ring
+        self.tune = tune
         self.turn_delay = turn_delay
         self.tap_number = tap_number
-        self.ampl = ampl
+        self.gain = gain
         self.phase = phase
-        self.x_pos = np.zeros((self.tap_number,))
-        self.y_pos = np.zeros((self.tap_number,))
-        self.kick_x = np.zeros((self.turn_delay+1,))
-        self.kick_y = np.zeros((self.turn_delay+1,))
-        self.coef_x = self.get_fir_x()
-        self.coef_y = self.get_fir_y()
+        self.bpm_error = bpm_error
+        self.max_kick = max_kick
+        self.plane = plane
+        self.x_pos = np.zeros((self.tap_number[0],))
+        self.y_pos = np.zeros((self.tap_number[1],))
+        self.tau_pos = np.zeros((self.tap_number[2],))
+        self.kick_x = np.zeros((self.turn_delay[0]+1,))
+        self.kick_y = np.zeros((self.turn_delay[1]+1,))
+        self.kick_tau = np.zeros((self.turn_delay[2]+1,))
+        self.coef_x = self.get_fir(self.tap_number[0], self.tune[0], self.phase[0], self.turn_delay[0], self.gain[0])
+        self.coef_y = self.get_fir(self.tap_number[1], self.tune[1], self.phase[1], self.turn_delay[1], self.gain[1])
+        self.coef_tau = self.get_fir(self.tap_number[2], self.tune[2], self.phase[2], self.turn_delay[2], self.gain[2])
         
-    def get_fir_x(self):        
+    def get_fir(self, tap_number, tune, phase, turn_delay, gain):        
         """
-        Compute the FIR coefficients for the horizontal plane.
+        Compute the FIR coefficients for the selected plane(s).
+        This method is based on the FIR filter design algorithm developped by
+        T.Nakamura.
         
         Returns
         -------
-        coef_x : array
-            Array containing the FIR coefficients for the horizontal plane.
+        FIR_coef : array
+            Array containing the FIR coefficients for the selected plane(s).
         """
-        phi_x = 2*np.pi*self.ring.tune[0]
-        zeta = (self.phase*2*np.pi)/360
-        coef_x = np.zeros((self.tap_number,))
+        it = np.zeros((tap_number,))
+        CC = np.zeros((5, tap_number,))
+        zeta = (phase*2*np.pi)/360
+        for k in range(tap_number):
+            it[k] = (-k - turn_delay)
         
-        it = np.zeros((self.tap_number,))
-        for k in range(self.tap_number):
-            it[k] = (-k - self.turn_delay)
-            
-        cs_x = np.zeros((self.tap_number,))
-        sn_x = np.zeros((self.tap_number,))
-        for i in range(self.tap_number):
-            cs_x[i] = (np.cos(phi_x*it[i]))
-            sn_x[i] = (np.sin(phi_x*it[i]))
-            
-        CC_x = np.zeros((5, self.tap_number))
-        for j in range(self.tap_number):
-            CC_x[0][j] = 1
-            CC_x[1][j] = cs_x[j]
-            CC_x[2][j] = sn_x[j]
-            CC_x[3][j] = it[j]*sn_x[j]
-            CC_x[4][j] = it[j]*cs_x[j]
-            
-        TCC_x = np.transpose(CC_x)        
-        W_x = np.linalg.inv(CC_x.dot(TCC_x))
-        D_x = W_x.dot(CC_x)
-       
-        for n in range(self.tap_number):
-            a_x = self.ampl*(D_x[1][n]*np.cos(zeta) + D_x[2][n])
-            coef_x[n] = (a_x)
-            
-        return coef_x
+        phi = 2*np.pi*tune
+        cs = np.cos(phi*it)
+        sn = np.sin(phi*it)
+        
+        CC[0][:] = 1
+        CC[1][:] = cs
+        CC[2][:] = sn
+        CC[3][:] = it*sn
+        CC[4][:] = it*cs
+        
+        TCC = np.transpose(CC)
+        W = np.linalg.inv(CC.dot(TCC))
+        D = W.dot(CC)
+        
+        FIR_coef = gain*(D[1][:]*np.cos(zeta) + D[2][:]*np.sin(zeta))
+        return FIR_coef
     
-    def get_fir_y(self):
+    def plot_fir(self, plane):
         """
-        Compute the FIR coefficients for the vertical plane.
-        
+        Plot the gain and the phase of the FIR filter for the desired plane
+
+        Parameters
+        ----------
+        plane : string
+            States "x" for the horizontal plane, "y" for the vertical and
+            "long" for the longitudinal one.
+
         Returns
         -------
-        coef_y : array
-            Array containing the FIR coefficients for the vertical plane.
+        None.
         """
-        phi_y = 2*np.pi*self.ring.tune[1]
-        zeta = (self.phase*2*np.pi)/360
-        coef_y = np.zeros((self.tap_number,))
-        
-        it = np.zeros((self.tap_number,))
-        for k in range(0,self.tap_number):
-            it[k] = (-k - self.turn_delay)
-                
-        cs_y = np.zeros((self.tap_number,))
-        sn_y = np.zeros((self.tap_number,))
-        for i in range(self.tap_number):
-            cs_y[i] = (np.cos(phi_y*it[i]))
-            sn_y[i] = (np.sin(phi_y*it[i]))
-            
-        CC_y = np.zeros((5, self.tap_number))
-        for j in range(self.tap_number):
-            CC_y[0][j] = 1
-            CC_y[1][j] = cs_y[j]
-            CC_y[2][j] = sn_y[j]
-            CC_y[3][j] = it[j]*sn_y[j]
-            CC_y[4][j] = it[j]*cs_y[j]
-            
-        TCC_y = np.transpose(CC_y)
-        W_y = np.linalg.inv(CC_y.dot(TCC_y))        
-        D_y = W_y.dot(CC_y)
-
-        for n in range(self.tap_number):
-            a_y = self.ampl*(D_y[1][n]*np.cos(zeta) + D_y[2][n]*np.sin(zeta))
-            coef_y[n] = (a_y)    
-            
-        return coef_y
-    
+        tune = np.arange(0,1,0.0001)
+        plane_dict = {"x":0 , "y":1 , "long":2}
+        index_plane = plane_dict[plane]
+        t = np.array([len(self.coef_x), len(self.coef_y), len(self.coef_tau)])
+        latency = np.exp(-1j*2*np.pi*tune*self.turn_delay[index_plane])
+        H_FIR = 0
+        liste = []
+        if(plane == "x"):
+            liste = self.coef_x
+        elif(plane == "y"):
+            liste = self.coef_y
+        elif(plane == "long"):
+            liste = self.coef_tau
+        for k in range(t[index_plane]):
+            H_FIR += liste[k]*np.exp(-1j*2*np.pi*(k)*tune)
+        H_tot = H_FIR * latency
+        gain = np.abs(H_tot)
+        phase = np.angle(H_tot, deg = True)
+        plt.figure("Gain"+str(index_plane))
+        plt.plot(tune, gain)
+        plt.title("Gain" + str(index_plane))
+        plt.xlabel("Tune")
+        plt.figure("Phase"+str(index_plane))
+        plt.plot(tune, phase)
+        plt.title("Phase in degree" + str(index_plane))
+        plt.xlabel("Tune")
+        plt.ylabel("Degree")
+         
     @Element.parallel
     def track(self, bunch):
         """
@@ -147,24 +248,51 @@ class FeedBackTransverse(Element) :
         Parameters
         ----------
         bunch : Bunch or Beam object
-        """
-        self.x_pos[0] = bunch.mean[0] 
-        self.y_pos[0] = bunch.mean[2]
-        
-        x_kick = 0
-        y_kick = 0     
+        """             
+        if(self.plane[0] == True):
+            self.x_pos[0] = bunch.mean[0] + np.random.normal(0, self.bpm_error[0])
+            x_kick = 0
+            
+            for k in range(self.tap_number[0]):
+                x_kick += self.coef_x[k]*self.x_pos[k]
+            if(self.max_kick[0] != 0 and x_kick > self.max_kick[0]):
+                x_kick = self.max_kick[0]
+            if(self.max_kick[0] != 0 and x_kick < -1*self.max_kick[0]):
+                x_kick = -1*self.max_kick[0]
+                
+            self.kick_x[-1] = x_kick
+            bunch["xp"] += self.kick_x[0]
+            self.x_pos = np.roll(self.x_pos,1)
+            self.kick_x = np.roll(self.kick_x, -1)
             
-        for k in range(self.tap_number):
+        if(self.plane[1] == True):
+            self.y_pos[0] = bunch.mean[2] + np.random.normal(0, self.bpm_error[1])
+            y_kick = 0
             
-            x_kick += self.coef_x[k]*self.x_pos[k]
-            y_kick += self.coef_y[k]*self.y_pos[k]
+            for k in range(self.tap_number[1]):
+                y_kick += self.coef_y[k]*self.y_pos[k]
+            if(self.max_kick[1] != 0 and y_kick > self.max_kick[1]):
+                y_kick = self.max_kick[1]
+            if(self.max_kick[1] != 0 and y_kick < -1*self.max_kick[1]):
+                y_kick = -1*self.max_kick[1]
+                
+            self.kick_y[-1] = y_kick
+            bunch["yp"] += self.kick_y[0]
+            self.y_pos = np.roll(self.y_pos,1)
+            self.kick_y = np.roll(self.kick_y, -1)
+            
+        if(self.plane[2] == True):
+            self.tau_pos[0] = bunch.mean[4] + np.random.normal(0, self.bpm_error[2])
+            tau_kick = 0
             
-        self.kick_x[-1] = x_kick
-        self.kick_y[-1] = y_kick
-       
-        bunch["xp"] -= self.kick_x[0]
-        bunch["yp"] -= self.kick_y[0]
-        self.x_pos = np.roll(self.x_pos,1)
-        self.y_pos = np.roll(self.y_pos,1)
-        self.kick_x = np.roll(self.kick_x, -1)
-        self.kick_y = np.roll(self.kick_y, -1)
\ No newline at end of file
+            for k in range(self.tap_number[2]):
+                tau_kick += self.coef_tau[k]*self.tau_pos[k]
+            if(self.max_kick[2] != 0 and tau_kick > self.max_kick[2]):
+                tau_kick = self.max_kick[2]
+            if(self.max_kick[2] != 0 and tau_kick < -1*self.max_kick[2]):
+                tau_kick = -1*self.max_kick[2]
+                
+            self.kick_tau[-1] = tau_kick
+            bunch["delta"] += self.kick_tau[0]
+            self.tau_pos = np.roll(self.tau_pos,1)
+            self.kick_tau = np.roll(self.kick_tau, -1)
\ No newline at end of file
diff --git a/mbtrack2/tracking/SimpleFeedBack.py b/mbtrack2/tracking/SimpleFeedBack.py
index 3dba6c7..ddb65ad 100644
--- a/mbtrack2/tracking/SimpleFeedBack.py
+++ b/mbtrack2/tracking/SimpleFeedBack.py
@@ -15,11 +15,12 @@ class SimpleFB(Element) :
    Parameters
    ----------
    ring : Synchrotron object
-   damping_time : array of shape (3,), optional
+   damping_time : array of shape (3,)
        Contains the desired damping time in seconds for each plane in the same 
        order sepcified in plane.
-   phase_diff : float, optional
-       Phase of the damper kick in degrees.
+   phase_diff : array of shape (3,)
+       Contains the desired phase in degree for each plane in the same 
+       order sepcified in plane.
    plane : bool array of shape (3,), optional
        Allow to choose on which plane the feedback is active. State True to
        turn on the feedback or False otherwise. Element 0 of the array defines
@@ -27,7 +28,7 @@ class SimpleFB(Element) :
        respectively.
     """
     
-    def __init__(self, ring,phase_diff = 90, damping_time = np.array([2e-3,2e-5,2e-5]),plane = np.ones((3,), dtype= bool)):
+    def __init__(self, ring,phase_diff, damping_time, plane = np.zeros((3,), dtype= bool)):
         self.ring = ring
         self.damping_time = damping_time
         self.phase_diff = phase_diff
@@ -45,10 +46,10 @@ class SimpleFB(Element) :
         bunch : Bunch or Beam object
         """
         if(self.plane[0] == True):
-            bunch["delta"] -= (2*self.ring.T0/self.damping_time[0])*np.sin(self.phase_diff)*bunch.mean[5]
-        
+            bunch["xp"] -= (2*self.ring.T0/self.damping_time[0])*np.sin(self.phase_diff[0])*bunch.mean[1]
+            
         if(self.plane[1] == True):
-            bunch["xp"] -= (2*self.ring.T0/self.damping_time[1])*np.sin(self.phase_diff)*bunch.mean[1]
-        
+            bunch["yp"] -= (2*self.ring.T0/self.damping_time[1])*np.sin(self.phase_diff[1])*bunch.mean[3]
+            
         if(self.plane[2] == True):
-            bunch["yp"] -= (2*self.ring.T0/self.damping_time[2])*np.sin(self.phase_diff)*bunch.mean[3]
\ No newline at end of file
+            bunch["delta"] -= (2*self.ring.T0/self.damping_time[2])*np.sin(self.phase_diff[1])*bunch.mean[5]
\ No newline at end of file
-- 
GitLab