Merge of Vlasov code with mbtrack2

Add CavityResonator class
Add BeamLoadingVlasov class, to do tested and developped further.
From BeamLoadingVlasov, only the beam_equilibrium method is well tested and benchmarked.
The cannonical_transform and solveB methods have to be tested and benchmarked further.

tracking/__init__.py

@@ -8,7 +8,7 @@ Created on Tue Jan 14 18:11:33 2020
 from mbtrack2.tracking.particles import (Electron, Proton, Bunch, Beam, 
                                          Particle)
 from mbtrack2.tracking.synchrotron import Synchrotron
-from mbtrack2.tracking.rf import RFCavity
+from mbtrack2.tracking.rf import RFCavity, CavityResonator
 from mbtrack2.tracking.parallel import Mpi
 from mbtrack2.tracking.optics import Optics, PhyisicalModel
 from mbtrack2.tracking.element import (Element, LongitudinalMap, 

tracking/rf.py

@@ -7,6 +7,9 @@ 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):
@@ -50,4 +53,316 @@ class RFCavity(Element):
         
     def value(self, val):
         return self.Vc / self.ring.E0 * np.cos( 
-                self.m * self.ring.omega1 * val + self.theta )
+                self.m * self.ring.omega1 * val + self.theta )
\ No newline at end of file
+    
+    
+class CavityResonator():
+    """Cavity resonator class for active or passive RF cavity with beam
+    loading or HOM, based on [1].
+    
+    Parameters
+    ----------
+    ring : Synchrotron object
+    m : int or float
+        Harmonic number of the cavity.
+    Rs : float
+        Shunt impedance of the cavity in [Ohm].
+    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).
+    Vc : float, optinal
+        Total cavity voltage in [V].
+    theta : float, optional
+        Total cavity phase in [rad].
+    
+    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.
+    """
+    def __init__(self, ring, m, Rs, Q, QL, detune, Vc=0, theta=0):
+        self.ring = ring
+        self.m = m
+        self.Rs = Rs
+        self.Q = Q
+        self.QL = QL
+        self.detune = detune
+        self.Vc = Vc
+        self.theta = theta
+
+    @property
+    def m(self):
+        """Harmonic number of the cavity"""
+        return self._m
+
+    @m.setter
+    def m(self, value):
+        self._m = value
+
+    @property
+    def Rs(self):
+        """Shunt impedance [ohm]"""
+        return self._Rs
+
+    @Rs.setter
+    def Rs(self, value):
+        self._Rs = value
+
+    @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
+
+    @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
+        
+    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_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)
\ No newline at end of file

vlasov/__init__.py

+# -*- coding: utf-8 -*-
+"""
+Created on Tue Jan 14 18:11:33 2020
+
+@author: gamelina
+"""
+
+from mbtrack2.vlasov.beamloading import BeamLoadingVlasov
\ No newline at end of file