Make new branch for RF related development

mbtrack2/tracking/bk_rf.py

+# -*- coding: utf-8 -*-
+"""
+This module handles radio-frequency (RF) cavitiy elements. 
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+import matplotlib.patches as mpatches
+from matplotlib.legend_handler import HandlerPatch
+from mbtrack2.tracking.element import Element
+
+class RFCavity(Element):
+    """
+    Perfect RF cavity class for main and harmonic RF cavities.
+    Use cosine definition.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    m : int
+        Harmonic number of the cavity
+    Vc : float
+        Amplitude of cavity voltage [V]
+    theta : float
+        Phase of Cavity voltage
+    """
+    def __init__(self, ring, m, Vc, theta):
+        self.ring = ring
+        self.m = m 
+        self.Vc = Vc
+        self.theta = theta
+        
+    @Element.parallel    
+    def track(self,bunch):
+        """
+        Tracking method for the element.
+        No bunch to bunch interaction, so written for Bunch objects and
+        @Element.parallel is used to handle Beam objects.
+        
+        Parameters
+        ----------
+        bunch : Bunch or Beam object
+        """
+        bunch["delta"] += self.Vc / self.ring.E0 * np.cos(
+                self.m * self.ring.omega1 * bunch["tau"] + self.theta )
+        
+    def value(self, val):
+        return self.Vc / self.ring.E0 * np.cos( 
+                self.m * self.ring.omega1 * val + self.theta )
+    
+    
+class CavityResonator():
+    """Cavity resonator class for active or passive RF cavity with beam
+    loading or HOM, based on [1,2].
+    
+    Use cosine definition.
+    
+    If used with mpi, beam.mpi.share_distributions must be called before the 
+    track method call.
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    m : int or float
+        Harmonic number of the cavity.
+    Rs : float
+        Shunt impedance of the cavities in [Ohm], defined as 0.5*Vc*Vc/Pc.
+        If Ncav = 1, used for the total shunt impedance.
+        If Ncav > 1, used for the shunt impedance per cavity.
+    Q : float
+        Quality factor of the cavity.
+    QL : float
+        Loaded quality factor of the cavity.
+    detune : float
+        Detuing of the cavity in [Hz], defined as (fr - m*ring.f1).
+    Ncav : int, optional
+        Number of cavities.
+    Vc : float, optinal
+        Total cavity voltage in [V].
+    theta : float, optional
+        Total cavity phase 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 
+        beam.mpi.share_distributions method.
+        The default is 75.
+        
+    Attributes
+    ----------
+    beam_phasor : complex
+        Beam phasor in [V].
+    beam_phasor_record : array of complex
+        Last beam phasor value of each bunch in [V].
+    generator_phasor : complex
+        Generator phasor in [V].
+    cavity_phasor : complex
+        Cavity phasor in [V].
+    cavity_phasor_record : array of complex
+        Last cavity phasor value of each bunch in [V].
+    cavity_voltage : float
+        Cavity total voltage in [V].
+    cavity_phase : float
+        Cavity total phase in [rad].
+    loss_factor : float
+        Cavity loss factor in [V/C].
+    Rs_per_cavity : float
+        Shunt impedance of a single cavity in [Ohm], defined as 0.5*Vc*Vc/Pc.
+    beta : float
+        Coupling coefficient of the cavity.
+    fr : float
+        Resonance frequency of the cavity in [Hz].
+    wr : float
+        Angular resonance frequency in [Hz.rad].
+    psi : float
+        Tuning angle in [rad].
+    filling_time : float
+        Cavity filling time in [s].
+    Pc : float
+        Power dissipated in the cavity walls in [W].
+    Pg : float
+        Generator power in [W].
+    Vgr : float
+        Generator voltage at resonance in [V].
+    theta_gr : float
+        Generator phase at resonance in [rad].
+    Vg : float
+        Generator voltage in [V].
+    theta_g : float
+        Generator phase in [rad].
+    tracking : bool
+        True if the tracking has been initialized.
+    bunch_index : int
+        Number of the tracked bunch in the current core.
+    distance : array
+        Distance between bunches.
+    valid_bunch_index : array
+    
+    Methods
+    -------
+    Vbr(I0)
+        Return beam voltage at resonance in [V].
+    Vb(I0)
+        Return beam voltage in [V].
+    Pb(I0)
+        Return power transmitted to the beam in [W].
+    Pr(I0)
+        Return power reflected back to the generator in [W].
+    Z(f)
+        Cavity impedance in [Ohm] for a given frequency f in [Hz].
+    set_optimal_coupling(I0)
+        Set coupling to optimal value.
+    set_optimal_detune(I0)
+        Set detuning to optimal conditions.
+    set_generator(I0)
+        Set generator parameters.
+    plot_phasor(I0)
+        Plot phasor diagram.
+    is_DC_Robinson_stable(I0)
+        Check DC Robinson stability.
+    plot_DC_Robinson_stability()
+        Plot DC Robinson stability limit.
+    init_tracking(beam)
+        Initialization of the tracking.
+    track(beam)
+        Tracking method.
+    phasor_decay(time)
+        Compute the beam phasor decay during a given time span.
+    phasor_evol(profile, bin_length, charge_per_mp)
+        Compute the beam phasor evolution during the crossing of a bunch.
+    VRF(z, I0)
+        Return the total RF voltage.
+    dVRF(z, I0)
+        Return derivative of total RF voltage.
+    ddVRF(z, I0)
+        Return the second derivative of total RF voltage.
+    deltaVRF(z, I0)
+        Return the generator voltage minus beam loading voltage.
+    init_FB(gain_A, gain_P, delay)
+        Initialize and switch on amplitude and phase feedback.
+    track_FB()
+        Tracking method for the amplitude and phase feedback.
+    
+    References
+    ----------
+    [1] Wilson, P. B. (1994). Fundamental-mode rf design in e+ e− storage ring 
+    factories. In Frontiers of Particle Beams: Factories with e+ e-Rings 
+    (pp. 293-311). Springer, Berlin, Heidelberg.
+    
+    [2] Yamamoto, Naoto, Alexis Gamelin, and Ryutaro Nagaoka. "Investigation 
+    of Longitudinal Beam Dynamics With Harmonic Cavities by Using the Code 
+    Mbtrack." IPAC’19, Melbourne, Australia, 2019.
+    """
+    def __init__(self, ring, m, Rs, Q, QL, detune, Ncav=1, Vc=0, theta=0, 
+                 n_bin=75):
+        self.ring = ring
+        self.m = m
+        self.Ncav = Ncav
+        if Ncav != 1:
+            self.Rs_per_cavity = Rs
+        else:
+            self.Rs = Rs
+        self.Q = Q
+        self.QL = QL
+        self.detune = detune
+        self.Vc = Vc
+        self.theta = theta
+        self.beam_phasor = np.zeros(1, dtype=np.complex)
+        self.beam_phasor_record = np.zeros((self.ring.h), dtype=np.complex)
+        self.tracking = False
+        self.Vg = 0
+        self.theta_g = 0
+        self.Vgr = 0
+        self.theta_gr = 0
+        self.Pg = 0
+        self.n_bin = int(n_bin)
+        self.FB = False
+        self.igFB = False
+        
+    def init_tracking(self, beam):
+        """
+        Initialization of the tracking.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        """
+        if beam.mpi_switch:
+            self.bunch_index = beam.mpi.bunch_num # Number of the tracked bunch in this processor
+            
+        self.distance = beam.distance_between_bunches
+        self.valid_bunch_index = beam.bunch_index
+        self.tracking = True
+        self.nturn = 0
+    
+    def track(self, beam):
+        """
+        Track a Beam object through the CavityResonator object.
+        
+        Can be used with or without mpi.
+        If used with mpi, beam.mpi.share_distributions must be called before.
+        
+        The beam phasor is given at t=0 (synchronous particle) of the first 
+        non empty bunch.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        """
+        
+        if self.tracking is False:
+            self.init_tracking(beam)
+        
+        for index, bunch in enumerate(beam):
+            
+            if beam.filling_pattern[index]:
+                
+                if beam.mpi_switch:
+                    # get rank of bunch n° index
+                    rank = beam.mpi.bunch_to_rank(index)
+                    # mpi -> get shared bunch profile for current bunch
+                    center = beam.mpi.tau_center[rank]
+                    profile = beam.mpi.tau_profile[rank]
+                    bin_length = center[1]-center[0]
+                    charge_per_mp = beam.mpi.charge_per_mp_all[rank]
+                    if index == self.bunch_index:
+                        sorted_index = beam.mpi.tau_sorted_index
+                else:
+                    # no mpi -> get bunch profile for current bunch
+                    if len(bunch) != 0:
+                        (bins, sorted_index, profile, center) = bunch.binning(n_bin=self.n_bin)
+                        bin_length = center[1]-center[0]
+                        charge_per_mp = bunch.charge_per_mp
+                        self.bunch_index = index
+                    else:
+                        # Update filling pattern
+                        beam.update_filling_pattern()
+                        beam.update_distance_between_bunches()
+                        # save beam phasor value
+                        self.beam_phasor_record[index] = self.beam_phasor
+                        # phasor decay to be at t=0 of the next bunch
+                        self.phasor_decay(self.ring.T1, ref_frame="beam")
+                        continue
+                
+                energy_change = bunch["tau"]*0
+                
+                # remove part of beam phasor decay to be at the start of the binning (=bins[0])
+                self.phasor_decay(center[0] - bin_length/2, ref_frame="beam")
+                
+                if index != self.bunch_index:
+                    self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="beam")
+                else:
+                    # modify beam phasor
+                    for i, center0 in enumerate(center):
+                        mp_per_bin = profile[i]
+                        
+                        if mp_per_bin == 0:
+                            self.phasor_decay(bin_length, ref_frame="beam")
+                            continue
+                        
+                        ind = (sorted_index == i)
+                        phase = self.m * self.ring.omega1 * (center0 + self.ring.T1* (index + self.ring.h * self.nturn))
+                        if self.igFB:
+                            Vgene = (self.generator_phasor_record[index]*np.exp(1j*phase)).real
+                            #Vgene = np.abs(self.generator_phasor_record[index])*np.cos(phase+np.angle(self.generator_phasor_record[index]))
+                        else:
+                            Vgene = self.Vg*np.cos(phase + self.theta_g)
+                        Vbeam = np.real(self.beam_phasor)
+                        Vtot = Vgene + Vbeam - charge_per_mp*self.loss_factor*mp_per_bin
+                        energy_change[ind] = Vtot / self.ring.E0
+    
+                        self.beam_phasor -= 2*charge_per_mp*self.loss_factor*mp_per_bin
+                        self.phasor_decay(bin_length, ref_frame="beam")
+                
+                # phasor decay to be at t=0 of the current bunch (=-1*bins[-1])
+                self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="beam")
+                
+                if index == self.bunch_index:
+                    # apply kick
+                    bunch["delta"] += energy_change
+            
+            # save beam phasor value
+            self.beam_phasor_record[index] = self.beam_phasor
+            
+            if self.FB and (self.FB_index == index):
+                self.track_FB()
+            
+            # phasor decay to be at t=0 of the next bunch
+            self.phasor_decay(self.ring.T1, ref_frame="beam")
+                
+        self.nturn += 1
+        
+    def init_phasor_track(self, beam):
+        """
+        Initialize the beam phasor for a given beam distribution using a
+        tracking like method.
+        
+        Follow the same steps as the track method but in the "rf" reference 
+        frame and without any modifications on the beam.
+
+        Parameters
+        ----------
+        beam : Beam object
+
+        """        
+        if self.tracking is False:
+            self.init_tracking(beam)
+            
+        n_turn = int(self.filling_time/self.ring.T0*10)
+        
+        for i in range(n_turn):
+            for j, bunch in enumerate(beam.not_empty):
+                
+                index = self.valid_bunch_index[j]
+                
+                if beam.mpi_switch:
+                    # get shared bunch profile for current bunch
+                    center = beam.mpi.tau_center[j]
+                    profile = beam.mpi.tau_profile[j]
+                    bin_length = center[1]-center[0]
+                    charge_per_mp = beam.mpi.charge_per_mp_all[j]
+                else:
+                    if i == 0:
+                        # get bunch profile for current bunch
+                        (bins, sorted_index, profile, center) = bunch.binning(n_bin=self.n_bin)
+                        if j == 0:
+                            self.profile_save = np.zeros((len(beam),len(profile),))
+                            self.center_save = np.zeros((len(beam),len(center),))
+                        self.profile_save[j,:] = profile
+                        self.center_save[j,:] = center
+                    else:
+                        profile = self.profile_save[j,:]
+                        center = self.center_save[j,:]
+                        
+                    bin_length = center[1]-center[0]
+                    charge_per_mp = bunch.charge_per_mp
+                
+                self.phasor_decay(center[0] - bin_length/2, ref_frame="rf")
+                self.phasor_evol(profile, bin_length, charge_per_mp, ref_frame="rf")
+                self.phasor_decay(-1 * (center[-1] + bin_length/2), ref_frame="rf")
+                self.phasor_decay( (self.distance[index] * self.ring.T1), ref_frame="rf")
+            
+    def phasor_decay(self, time, ref_frame="beam"):
+        """
+        Compute the beam phasor decay during a given time span, assuming that 
+        no particles are crossing the cavity during the time span.
+
+        Parameters
+        ----------
+        time : float
+            Time span in [s], can be positive or negative.
+        ref_frame : string, optional
+            Reference frame to be used, can be "beam" or "rf".
+
+        """
+        if ref_frame == "beam":
+            delta = self.wr
+        elif ref_frame == "rf":
+            delta = (self.wr - self.m*self.ring.omega1)
+        self.beam_phasor = self.beam_phasor * np.exp((-1/self.filling_time +
+                                  1j*delta)*time)
+        
+    def phasor_evol(self, profile, bin_length, charge_per_mp, ref_frame="beam"):
+        """
+        Compute the beam phasor evolution during the crossing of a bunch using 
+        an analytic formula [1].
+        
+        Assume that the phasor decay happens before the beam loading.
+
+        Parameters
+        ----------
+        profile : array
+            Longitudinal profile of the bunch in [number of macro-particle].
+        bin_length : float
+            Length of a bin in [s].
+        charge_per_mp : float
+            Charge per macro-particle in [C].
+        ref_frame : string, optional
+            Reference frame to be used, can be "beam" or "rf".
+            
+        References
+        ----------
+        [1] mbtrack2 manual.
+            
+        """
+        if ref_frame == "beam":
+            delta = self.wr
+        elif ref_frame == "rf":
+            delta = (self.wr - self.m*self.ring.omega1)
+            
+        n_bin = len(profile)
+        
+        # Phasor decay during crossing time
+        deltaT = n_bin*bin_length
+        self.phasor_decay(deltaT, ref_frame)
+        
+        # Phasor evolution due to induced voltage by marco-particles
+        k = np.arange(0, n_bin)
+        var = np.exp( (-1/self.filling_time + 1j*delta) * 
+                      (n_bin-k) * bin_length )
+        sum_tot = np.sum(profile * var)
+        sum_val = -2 * sum_tot * charge_per_mp * self.loss_factor
+        self.beam_phasor += sum_val
+        
+    def init_phasor(self, beam):
+        """
+        Initialize the beam phasor for a given beam distribution using an
+        analytic formula [1].
+        
+        No modifications on the Beam object.
+
+        Parameters
+        ----------
+        beam : Beam object
+            
+        References
+        ----------
+        [1] mbtrack2 manual.
+
+        """
+        
+        # Initialization
+        if self.tracking is False:
+            self.init_tracking(beam)
+        
+        N = self.n_bin - 1
+        delta = (self.wr - self.m*self.ring.omega1)
+        n_turn = int(self.filling_time/self.ring.T0*10)
+        
+        T = np.ones(self.ring.h)*self.ring.T1
+        bin_length = np.zeros(self.ring.h)
+        charge_per_mp = np.zeros(self.ring.h)
+        profile = np.zeros((N, self.ring.h))
+        center = np.zeros((N, self.ring.h))
+        
+        # Gather beam distribution data
+        for j, bunch in enumerate(beam.not_empty):
+            index = self.valid_bunch_index[j]
+            if beam.mpi_switch:
+                beam.mpi.share_distributions(beam, n_bin=self.n_bin)
+                center[:,index] = beam.mpi.tau_center[j]
+                profile[:,index] = beam.mpi.tau_profile[j]
+                bin_length[index] = center[1, index]-center[0, index]
+                charge_per_mp[index] = beam.mpi.charge_per_mp_all[j]
+            else:
+                (bins, sorted_index, profile[:, index], center[:, index]) = bunch.binning(n_bin=self.n_bin)
+                bin_length[index] = center[1, index]-center[0, index]
+                charge_per_mp[index] = bunch.charge_per_mp
+            T[index] -= (center[-1, index] + bin_length[index]/2)
+            if index != 0:
+                T[index - 1] += (center[0, index] - bin_length[index]/2)
+        T[self.ring.h - 1] += (center[0, 0] - bin_length[0]/2)
+
+        # Compute matrix coefficients
+        k = np.arange(0, N)
+        Tkj = np.zeros((N, self.ring.h))
+        for j in range(self.ring.h):
+            sum_t = np.array([T[n] + N*bin_length[n] for n in range(j+1,self.ring.h)])
+            Tkj[:,j] = (N-k)*bin_length[j] + T[j] + np.sum(sum_t)
+            
+        var = np.exp( (-1/self.filling_time + 1j*delta) * Tkj )
+        sum_tot = np.sum((profile*charge_per_mp) * var)
+        
+        # Use the formula n_turn times
+        for i in range(n_turn):
+            # Phasor decay during one turn
+            self.phasor_decay(self.ring.T0, ref_frame="rf")
+            # Phasor evolution due to induced voltage by marco-particles during one turn
+            sum_val = -2 * sum_tot * self.loss_factor
+            self.beam_phasor += sum_val
+        
+        # Replace phasor at t=0 (synchronous particle) of the first non empty bunch.
+        idx0 = self.valid_bunch_index[0]
+        self.phasor_decay(center[-1,idx0] + bin_length[idx0]/2, ref_frame="rf")
+    
+    @property
+    def generator_phasor(self):
+        """Generator phasor in [V]"""
+        return self.Vg*np.exp(1j*self.theta_g)
+    
+    @property
+    def cavity_phasor(self):
+        """Cavity total phasor in [V]"""
+        return self.generator_phasor + self.beam_phasor
+    
+    @property
+    def cavity_phasor_record(self):
+        """Last cavity phasor value of each bunch in [V]"""
+        if self.igFB:
+            return self.generator_phasor_record + self.beam_phasor_record
+        else:
+            return self.generator_phasor + self.beam_phasor_record
+    
+    @property
+    def cavity_voltage(self):
+        """Cavity total voltage in [V]"""
+        return np.abs(self.cavity_phasor)
+    
+    @property
+    def cavity_phase(self):
+        """Cavity total phase in [rad]"""
+        return np.angle(self.cavity_phasor)
+    
+    @property
+    def beam_voltage(self):
+        """Beam loading voltage in [V]"""
+        return np.abs(self.beam_phasor)
+    
+    @property
+    def beam_phase(self):
+        """Beam loading phase in [rad]"""
+        return np.angle(self.beam_phasor)
+    
+    @property
+    def loss_factor(self):
+        """Cavity loss factor in [V/C]"""
+        return self.wr*self.Rs/(2 * self.Q)
+
+    @property
+    def m(self):
+        """Harmonic number of the cavity"""
+        return self._m
+
+    @m.setter
+    def m(self, value):
+        self._m = value
+        
+    @property
+    def Ncav(self):
+        """Number of cavities"""
+        return self._Ncav
+
+    @Ncav.setter
+    def Ncav(self, value):
+        self._Ncav = value
+        
+    @property
+    def Rs_per_cavity(self):
+        """Shunt impedance of a single cavity in [Ohm], defined as 
+        0.5*Vc*Vc/Pc."""
+        return self._Rs_per_cavity
+
+    @Rs_per_cavity.setter
+    def Rs_per_cavity(self, value):
+        self._Rs_per_cavity = value
+
+    @property
+    def Rs(self):
+        """Shunt impedance [ohm]"""
+        return self.Rs_per_cavity * self.Ncav
+
+    @Rs.setter
+    def Rs(self, value):
+        self.Rs_per_cavity = value / self.Ncav
+
+    @property
+    def Q(self):
+        """Quality factor"""
+        return self._Q
+
+    @Q.setter
+    def Q(self, value):
+        self._Q = value
+
+    @property
+    def QL(self):
+        """Loaded quality factor"""
+        return self._QL
+
+    @QL.setter
+    def QL(self, value):
+        self._QL = value
+        self._beta = self.Q/self.QL - 1
+        self._RL = self.Rs/(1 + self._beta)
+
+    @property
+    def beta(self):
+        """Coupling coefficient"""
+        return self._beta
+
+    @beta.setter
+    def beta(self, value):
+        self.QL = self.Q/(1 + value)
+
+    @property
+    def detune(self):
+        """Cavity detuning [Hz] - defined as (fr - m*f1)"""
+        return self._detune
+
+    @detune.setter
+    def detune(self, value):
+        self._detune = value
+        self._fr = self.detune + self.m*self.ring.f1
+        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))
+
+    @property
+    def fr(self):
+        """Resonance frequency of the cavity in [Hz]"""
+        return self._fr
+
+    @fr.setter
+    def fr(self, value):
+        self.detune = value - self.m*self.ring.f1
+
+    @property
+    def wr(self):
+        """Angular resonance frequency in [Hz.rad]"""
+        return self._wr
+
+    @wr.setter
+    def wr(self, value):
+        self.detune = (value - self.m*self.ring.f1)*2*np.pi
+
+    @property
+    def psi(self):
+        """Tuning angle in [rad]"""
+        return self._psi
+
+    @psi.setter
+    def psi(self, value):
+        delta = (self.ring.f1*self.m*np.tan(value)/self.QL)**2 + 4*(self.ring.f1*self.m)**2
+        fr = (self.ring.f1*self.m*np.tan(value)/self.QL + np.sqrt(delta))/2
+        self.detune = fr - self.m*self.ring.f1
+        
+    @property
+    def filling_time(self):
+        """Cavity filling time in [s]"""
+        return 2*self.QL/self.wr
+    
+    @property
+    def Pc(self):
+        """Power dissipated in the cavity walls in [W]"""
+        return self.Vc**2 / (2 * self.Rs)
+    
+    def Pb(self, I0):
+        """
+        Return power transmitted to the beam in [W] - near Eq. (4.2.3) in [1].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Power transmitted to the beam in [W].
+
+        """
+        return I0 * self.Vc * np.cos(self.theta)
+    
+    def Pr(self, I0):
+        """
+        Power reflected back to the generator in [W].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Power reflected back to the generator in [W].
+
+        """
+        return self.Pg - self.Pb(I0) - self.Pc
+
+    def Vbr(self, I0):
+        """
+        Return beam voltage at resonance in [V].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Beam voltage at resonance in [V].
+
+        """
+        return 2*I0*self.Rs/(1+self.beta)
+    
+    def Vb(self, I0):
+        """
+        Return beam voltage in [V].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        float
+            Beam voltage in [V].
+
+        """
+        return self.Vbr(I0)*np.cos(self.psi)
+    
+    def Z(self, f):
+        """Cavity impedance in [Ohm] for a given frequency f in [Hz]"""
+        return self.Rs/(1 + 1j*self.QL*(self.fr/f - f/self.fr))
+    
+    def set_optimal_detune(self, I0):
+        """
+        Set detuning to optimal conditions - second Eq. (4.2.1) in [1].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        """
+        self.psi = np.arctan(-self.Vbr(I0)/self.Vc*np.sin(self.theta))
+        
+    def set_optimal_coupling(self, I0):
+        """
+        Set coupling to optimal value - Eq. (4.2.3) in [1].
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        """
+        self.beta = 1 + (2 * I0 * self.Rs * np.cos(self.theta) / 
+                         self.Vc)
+                
+    def set_generator(self, I0):
+        """
+        Set generator parameters (Pg, Vgr, theta_gr, Vg and theta_g) for a 
+        given current and set of parameters.
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        """
+        
+        # Generator power [W] - Eq. (4.1.2) [1] corrected with factor (1+beta)**2 instead of (1+beta**2)
+        self.Pg = self.Vc**2*(1+self.beta)**2/(2*self.Rs*4*self.beta*np.cos(self.psi)**2)*(
+            (np.cos(self.theta) + 2*I0*self.Rs/(self.Vc*(1+self.beta))*np.cos(self.psi)**2 )**2 + 
+            (np.sin(self.theta) + 2*I0*self.Rs/(self.Vc*(1+self.beta))*np.cos(self.psi)*np.sin(self.psi) )**2)
+        # Generator voltage at resonance [V] - Eq. (3.2.2) [1]
+        self.Vgr = 2*self.beta**(1/2)/(1+self.beta)*(2*self.Rs*self.Pg)**(1/2)
+        # Generator phase at resonance [rad] - from Eq. (4.1.1)
+        self.theta_gr = np.arctan((self.Vc*np.sin(self.theta) + self.Vbr(I0)*np.cos(self.psi)*np.sin(self.psi))/
+                    (self.Vc*np.cos(self.theta) + self.Vbr(I0)*np.cos(self.psi)**2)) - self.psi
+        # Generator voltage [V]
+        self.Vg = self.Vgr*np.cos(self.psi)
+        # Generator phase [rad]
+        self.theta_g = self.theta_gr + self.psi
+        
+    def plot_phasor(self, I0):
+        """
+        Plot phasor diagram showing the vector addition of generator and beam 
+        loading voltage.
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+            
+        Returns
+        -------
+        Figure.
+
+        """
+
+        def make_legend_arrow(legend, orig_handle,
+                              xdescent, ydescent,
+                              width, height, fontsize):
+            p = mpatches.FancyArrow(0, 0.5*height, width, 0, length_includes_head=True, head_width=0.75*height )
+            return p
+
+        fig = plt.figure()
+        ax= fig.add_subplot(111, polar=True)
+        ax.set_rmax(max([1.2,self.Vb(I0)/self.Vc*1.2,self.Vg/self.Vc*1.2]))
+        arr1 = ax.arrow(self.theta, 0, 0, 1, alpha = 0.5, width = 0.015,
+                         edgecolor = 'black', lw = 2)
+
+        arr2 = ax.arrow(self.psi + np.pi, 0, 0,self.Vb(I0)/self.Vc, alpha = 0.5, width = 0.015,
+                         edgecolor = 'red', lw = 2)
+
+        arr3 = ax.arrow(self.theta_g, 0, 0,self.Vg/self.Vc, alpha = 0.5, width = 0.015,
+                         edgecolor = 'blue', lw = 2)
+
+        ax.set_rticks([])  # less radial ticks
+        plt.legend([arr1,arr2,arr3], ['Vc','Vb','Vg'],handler_map={mpatches.FancyArrow : HandlerPatch(patch_func=make_legend_arrow),})
+        
+        return fig
+        
+    def is_DC_Robinson_stable(self, I0):
+        """
+        Check DC Robinson stability - Eq. (6.1.1) [1]
+
+        Parameters
+        ----------
+        I0 : float
+            Beam current in [A].
+
+        Returns
+        -------
+        bool
+
+        """
+        return 2*self.Vc*np.sin(self.theta) + self.Vbr(I0)*np.sin(2*self.psi) > 0
+    
+    def plot_DC_Robinson_stability(self, detune_range = [-1e5,1e5]):
+        """
+        Plot DC Robinson stability limit.
+
+        Parameters
+        ----------
+        detune_range : list or array, optional
+            Range of tuning to plot in [Hz].
+
+        Returns
+        -------
+        Figure.
+
+        """
+        old_detune = self.psi
+        
+        x = np.linspace(detune_range[0],detune_range[1],1000)
+        y = []
+        for i in range(0,x.size):
+            self.detune = x[i]
+            y.append(-self.Vc*(1+self.beta)/(self.Rs*np.sin(2*self.psi))*np.sin(self.theta)) # droite de stabilité
+            
+        fig = plt.figure()
+        ax = plt.gca()
+        ax.plot(x,y)
+        ax.set_xlabel("Detune [Hz]")
+        ax.set_ylabel("Threshold current [A]")
+        ax.set_title("DC Robinson stability limit")
+        
+        self.psi = old_detune
+        
+        return fig
+        
+    def VRF(self, z, I0, F = 1, PHI = 0):
+        """Total RF voltage taking into account form factor amplitude F and form factor phase PHI"""
+        return self.Vg*np.cos(self.ring.k1*self.m*z + self.theta_g) - self.Vb(I0)*F*np.cos(self.ring.k1*self.m*z + self.psi - PHI)
+    
+    def dVRF(self, z, I0, F = 1, PHI = 0):
+        """Return derivative of total RF voltage taking into account form factor amplitude F and form factor phase PHI"""
+        return -1*self.Vg*self.ring.k1*self.m*np.sin(self.ring.k1*self.m*z + self.theta_g) + self.Vb(I0)*F*self.ring.k1*self.m*np.sin(self.ring.k1*self.m*z + self.psi - PHI)
+    
+    def ddVRF(self, z, I0, F = 1, PHI = 0):
+        """Return the second derivative of total RF 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 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 init_FB(self, gain_A, gain_P, delay, FB_index=0):
+        """
+        Initialize and switch on amplitude and phase feedback.
+        
+        Objective amplitude and phase are self.Vc and self.theta.
+
+        Parameters
+        ----------
+        gain_A : float
+            Amplitude (voltage) gain of the feedback.
+        gain_P : float
+            Phase gain of the feedback.
+        delay : int
+            Feedback delay in unit of turns.
+        FB_index : int, optional
+            Bunch index at which the feedback is applied.
+            Default is 0.
+
+        Returns
+        -------
+        None.
+
+        """
+        self.gain_A = gain_A
+        self.gain_P = gain_P
+        self.delay = int(delay)
+        self.FB = True
+        self.volt_delay = np.ones(self.delay)*self.Vc
+        self.phase_delay = np.ones(self.delay)*self.theta
+        self.FB_index = int(FB_index)
+        
+    def track_FB(self):
+        """
+        Tracking method for the amplitude and phase feedback.
+
+        Returns
+        -------
+        None.
+
+        """
+        diff_A = self.volt_delay[-1] - self.Vc
+        diff_P = self.phase_delay[-1] - self.theta
+        self.Vg -= self.gain_A*diff_A
+        self.theta_g -= self.gain_P*diff_P
+        self.volt_delay = np.roll(self.volt_delay, 1)
+        self.phase_delay = np.roll(self.phase_delay, 1)
+        self.volt_delay[0] = self.cavity_voltage
+        self.phase_delay[0] = self.cavity_phase
+
+    def init_generatorCurrentFB(self,gain,sample_num,every,delay,Vref=None,IQ=True):
+        """
+        Generator current (ig) Feedback,
+        assumption :  delay >> samle_every >> sample_num
+
+        Adjusting Vg parameter to keep the Vc to target(Vref) values.
+        The function "track_directSample_FB" is
+           1) Monitor the Vc phasor
+           Mean Vc value between specified bunch number (sample_num) is monitored
+           with specified interval period (every).
+           2) Changing the ig phasor
+           The ig phasor is changed according the difference of the specified
+           reference values (Vref) with specified gain (gain).
+           By using ig instead of Vg, the cavity response can be taken account.
+           3) ig changes are reflected to Vg after the specifed delay (delay) of the system
+
+        Parameters
+        -------
+        gain : float list
+            Pgain and Igain of the feedback system
+            For IQ feedback, same gain set is used for I and Q.
+        sample_num : int
+            Sample number to monitor Vc 
+            The averaged Vc value in sample_num is monitored.
+            Unit is a bucket number.
+        every : int
+            interval to monitor and change Vc
+            Unit is a bucket number.
+        delay : int
+            Loop delay of the assumed system. 
+            Unit is a bucket number.
+        Vref : float, complex
+            target value of the feedback
+            if None, .Vc and .theta is set as the reference
+        """
+        if Vref is not None:
+            if isinstance(Vref,complex):
+                self.Vref = Vref
+            elif isinstance(Vref,list):
+                self.Vref = Vref[0] * np.exp(1j*Vref[1])
+            else:
+                self.Vref = Vref * np.exp(1j*self.ring.U0/Vref)
+        else: 
+            self.Vref = self.Vc * np.exp(1j*self.theta)
+        self.Vref_theta = np.angle(self.Vref)
+        self.IQ = False
+        if IQ:
+            self.IQ = True
+        #else:
+        #    self.fbRef = self.Vref[0] * np.exp(-1j*self.Vref[1])
+        #    #self.sbRef[0] = self.sbFB_ref_Vc * np.cos(self.sbFB_ref_theta )
+        #    #self.sbRef[1] = self.sbFB_ref_Vc * np.sin(self.sbFB_ref_theta )
+
+        if isinstance(gain,list): # if IQ is true, len(gain) should be 2
+            self.igFB_g = gain
+        else:
+            self.igFB_g = [gain,0]
+
+        if delay > 0:
+            self.igFB_delay = int(delay)
+        else:
+            self.igFB_delay = 1
+        if every > 0:
+            self.igFB_every = int(every)
+        else:
+            self.igFB_every = 1
+        record_size = int(np.ceil(self.igFB_delay/self.igFB_every))
+        if every - sample_num < 0:
+            raise ValueError("Bad parameter set : sample_num or sample_every")
+        if record_size < 1:
+            raise ValueError("Bad parameter set : delay or sample_every")
+        self.igFB_sample = int(sample_num)
+
+        # init lists for FB process
+        self.generator_phasor_record = np.ones(self.ring.h)*self.generator_phasor
+        self.igFB_ig_phasor = np.ones(self.ring.h+1)*self._Vg2Ig(self.generator_phasor)
+        #self.igFB_vc_previous = np.ones(self.igFB_sample)*self.cavity_phasor # 
+        self.igFB_vc_previous = np.ones(self.igFB_sample)*self.Vref # The diffs at 1st turn become almost zero.
+        self.igFB_diff_record = np.zeros(record_size,dtype=complex)
+        self.igFB_I_record = 0+0j
+
+        self.igFB_motRot = np.exp(-1j*np.angle(self.Vref))
+        self.igFB_ref = self.Vref*self.igFB_motRot
+
+        self.igFB_sample_list = range(0,self.ring.h,self.igFB_every)
+        self.igFB = True
+
+    def track_generatorCurrentFB(self):
+        """
+        update self.igFB_generator_phasor
+        This function should be called every turn. 
+
+        1) sampling Vc, recording diff
+        2) feedback (update Ig)
+        3) Ig -> Vg for every bucket
+
+        Vc -->IQ-->(rot)-->(-)-->(V->I,fac)-->PI-->
+                            |
+                            Ref
+        """
+        vc_list = np.concatenate([self.igFB_vc_previous,self.cavity_phasor_record])
+        #print("Debug vc_list",vc_list[-5],vc_list[5],np.mean(vc_list))
+        #print("Debug cavity_phasor", np.mean(self.cavity_phasor),np.mean(self.cavity_phasor_record))
+        #print("Debug",np.min(vc_list),np.max(vc_list))
+        #fb_bucket_list = self.igFB_sample_list
+        for index in self.igFB_sample_list:
+            # 2) updating Ig using last item of the list
+            diff = self.igFB_diff_record[-1]
+            self.igFB_I_record += diff/self.ring.f1
+            fb_ratio = self.igFB_g[0]*diff + self.igFB_g[1]*self.igFB_I_record
+            self.igFB_ig_phasor[index:] += (fb_ratio.real*self.igFB_ig_phasor[index:].real 
+                    + 1j*fb_ratio.imag*self.igFB_ig_phasor[index:].imag)
+            # Shift the record
+            self.igFB_diff_record = np.roll(self.igFB_diff_record,1)
+            # 1) recording diff as a first item of the list
+            mean_vc = np.mean(vc_list[index:self.igFB_sample+index])*self.igFB_motRot
+            self.igFB_diff_record[0] =  self._Vg2Ig(self.igFB_ref - mean_vc)
+            #mean_vc = np.mean(vc_list[index:self.igFB_sample+index])
+            #self.igFB_diff_record[0] = 1 + 1j - (mean_vc.real/self.Vref.real + 1j*mean_vc.imag/self.Vref.imag)
+            #print("Debug1,",self.igFB_diff_record[0],np.mean(vc_list[index:self.igFB_sample+index]),np.abs(np.mean(vc_list)),np.angle(np.mean(vc_list)))
+            #print("Debug2,", self.igFB_genCurrent_phasor[index])
+            #self.igFB_diff_record[0] = 1 - np.mean(vc_list[index-self.igFB_sample:index])/self.Vref
+        # update sample_list for next turn
+        self.igFB_sample_list = range(index+self.igFB_every-self.ring.h,self.ring.h,self.igFB_every)
+        # update vc_previous for next turn
+        self.igFB_vc_previous = self.cavity_phasor_record[- self.igFB_sample:]
+
+        # Ig -> Vg
+        for index in range(self.ring.h):
+            self.generator_phasor_record[index] = (self.generator_phasor_record[index-1]
+                    *np.exp(-1/self.filling_time*(1 - 1j*np.tan(self.psi))*self.ring.T1)
+                    + self.igFB_ig_phasor[index]*self.loss_factor*self.ring.T1)
+        self.Vg=np.mean(np.abs(self.generator_phasor_record))
+        self.theta_g=np.mean(np.angle(self.generator_phasor_record))
+
+    def _Vg2Ig(self,Vg):
+        """
+        Return Ig from Vg
+        assuming constant Vg
+
+        Eq.25 of N.Yamamoto et al., PRAB 21, 012001 (2018)
+        """
+        fac = 1/self.filling_time/self.loss_factor 
+        ig = fac*Vg*( 1 - 1j*np.tan(self.psi) )
+        return ig
+