Small additions and bugfix

Add ABI threshold calculation
Bugfix synchrotron_tune computation
Add zoom option to ImpedanceModel.plot_area

collective_effects/__init__.py

@@ -5,7 +5,7 @@ Created on Tue Jan 14 12:25:30 2020
 @author: gamelina
 """
 
-from mbtrack2.collective_effects.instabilities import MBI_threshold
+from mbtrack2.collective_effects.instabilities import mbi_threshold, cbi_threshold
 from mbtrack2.collective_effects.resistive_wall import skin_depth, CircularResistiveWall
 from mbtrack2.collective_effects.tapers import StupakovRectangularTaper, StupakovCircularTaper
 from mbtrack2.collective_effects.wakefield import ComplexData, Impedance, WakeFunction, ImpedanceModel, Wakefield

collective_effects/instabilities.py

@@ -8,7 +8,7 @@ General calculations about instabilities
 
 from scipy.constants import c, m_e, e, pi, epsilon_0
 
-def MBI_threshold(ring, sigma, R, b):
+def mbi_threshold(ring, sigma, R, b):
     """
     Compute the microbunching instability (MBI) threshold for a bunched beam
     considering the steady-state parallel plate model [1][2].
@@ -41,4 +41,47 @@ def MBI_threshold(ring, sigma, R, b):
         sigma**(1/3) / ( c * e * R**(1/3) ))
     I = N*e/ring.T0
     
-    return I
\ No newline at end of file
+    return I
+
+def cbi_threshold(ring, I, Vrf, f, beta, Ncav=1):
+    """
+    Compute the longitudinal and transverse coupled bunch instability 
+    thresolds driven by HOMs [1].
+
+    Parameters
+    ----------
+    ring : Synchrotron object
+    I : float
+        Total beam current in [A].
+    Vrf : float
+        Total RF voltage in [V].
+    f : float
+        Frequency of the HOM in [Hz].
+    beta : array-like of shape (2,)
+        Horizontal and vertical beta function at the HOM position in [m].
+    Ncav : int, optional
+        Number of RF cavity.
+
+    Returns
+    -------
+    Zlong : float
+        Maximum longitudinal impedance of the HOM in [ohm].
+    Zxdip : float
+        Maximum horizontal dipolar impedance of the HOM in [ohm/m].
+    Zydip : float
+        Maximum vertical dipolar impedance of the HOM in [ohm/m].
+        
+    References
+    ----------
+    [1] : Ruprecht, Martin, et al. "Calculation of Transverse Coupled Bunch 
+    Instabilities in Electron Storage Rings Driven By Quadrupole Higher Order 
+    Modes." 7th Int. Particle Accelerator Conf.(IPAC'16), Busan, Korea. 
+    """
+    
+    fs = ring.synchrotron_tune(Vrf)*ring.f0
+    Zlong = fs/(f*ring.ac*ring.tau[2]) * (2*ring.E0) / (ring.f0 * I * Ncav)
+    Zxdip = 1/(ring.tau[0]*beta[0]) * (2*ring.E0) / (ring.f0 * I * Ncav)
+    Zydip = 1/(ring.tau[1]*beta[1]) * (2*ring.E0) / (ring.f0 * I * Ncav)
+    
+    return (Zlong, Zxdip, Zydip)
+    
\ No newline at end of file

collective_effects/wakefield.py

@@ -17,6 +17,8 @@ from mbtrack2.tracking.element import Element
 from copy import deepcopy
 from scipy.integrate import trapz
 from scipy.interpolate import interp1d
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+
 
 class ComplexData:
     """
@@ -691,7 +693,7 @@ class ImpedanceModel(Element):
         self.update_name_list()
             
     def plot_area(self, Z_type="Zlong", component="real", sigma=None, 
-                  attr_list=None):
+                  attr_list=None, zoom=False):
         """
         Plot the contributions of different kind of wakefields.
 
@@ -705,7 +707,9 @@ class ImpedanceModel(Element):
             RMS bunch length in [s] to use for the spectral density. If equal
             to None, the spectral density is not plotted.
         attr_list : list or array of str, optional
-            Attributes to plot. 
+            Attributes to plot.
+        zoom : bool
+            If True, add a zoomed plot on top right corner.
 
         """
         if attr_list is None:
@@ -766,7 +770,28 @@ class ImpedanceModel(Element):
         ax.set_xlabel("Frequency [GHz]")
         ax.set_ylabel(label_list[leg] + " - " + component + " part")
         ax.set_title(label_list[leg] + " - " + component + " part")
+        
+        if zoom is True:
+            in_ax = inset_axes(ax,
+                            width="30%", # width = 30% of parent_bbox
+                            height=1.5, # height : 1 inch
+                            loc=1)
             
+            total_imp = 0
+            for index in  sorted_index:
+                attr = attr_list[index]
+                # Set all impedances with common indexes using + zero_impedance
+                try:
+                    sum_imp = getattr(getattr(self, attr), Z_type) + zero_impedance
+                    in_ax.fill_between(sum_imp.data.index*1e-3, total_imp, 
+                                    total_imp + sum_imp.data[component]*1e-9)
+                    total_imp += sum_imp.data[component]*1e-9
+                except AttributeError:
+                    pass
+            in_ax.set_xlim([0, 200])
+            in_ax.set_xlabel("Frequency [kHz]")
+            in_ax.set_ylabel(r"$[G\Omega]$")
+                
         return fig
             
     def effective_impedance(self, m, mu, sigma, M, tuneS, xi=None, 

tracking/synchrotron.py

@@ -237,7 +237,7 @@ class Synchrotron:
                     self.optics.local_dispersion[3]**2*self.sigma_delta)**0.5
         return sigma
     
-    def synchrotron_tune(Vrf):
+    def synchrotron_tune(self, Vrf):
         """
         Compute synchrotron tune from RF voltage
         
@@ -251,7 +251,7 @@ class Synchrotron:
         tuneS : float
             Synchrotron tune.
         """
-        phase = np.pi - np.arcsin(ring.U0 / Vrf)
-        tuneS = np.sqrt( - (Vrf / ring.E0) * (ring.h * ring.ac) / (2*np.pi) 
+        phase = np.pi - np.arcsin(self.U0 / Vrf)
+        tuneS = np.sqrt( - (Vrf / self.E0) * (self.h * self.ac) / (2*np.pi) 
                         * np.cos(phase) )
         return tuneS
\ No newline at end of file