From 652b68a103ae2881526d8901d91f3f0dc822127e Mon Sep 17 00:00:00 2001
From: Gamelin Alexis <alexis.gamelin@synchrotron-soleil.fr>
Date: Thu, 27 Oct 2022 16:27:51 +0200
Subject: [PATCH] Add function to plot lcbi stability diagram

Plotting function for longitudinal coupled bunch stability diagram for any set of CavityResonator objects.
Add loaded shunt impedance property to CavityResonator.
Change CavityResonator.Z method to use loaded shunt impedance.
---
 mbtrack2/instability/__init__.py      |   3 +-
 mbtrack2/instability/instabilities.py | 113 +++++++++++++++++++++-----
 mbtrack2/tracking/rf.py               |   7 +-
 3 files changed, 102 insertions(+), 21 deletions(-)

diff --git a/mbtrack2/instability/__init__.py b/mbtrack2/instability/__init__.py
index 054a9bd..5a5ee82 100644
--- a/mbtrack2/instability/__init__.py
+++ b/mbtrack2/instability/__init__.py
@@ -8,4 +8,5 @@ from mbtrack2.instability.instabilities import (mbi_threshold,
                                                 lcbi_growth_rate_mode,
                                                 lcbi_growth_rate,
                                                 rwmbi_growth_rate,
-                                                rwmbi_threshold)
\ No newline at end of file
+                                                rwmbi_threshold,
+                                                lcbi_stability_diagram)
\ No newline at end of file
diff --git a/mbtrack2/instability/instabilities.py b/mbtrack2/instability/instabilities.py
index bfecc53..195f035 100644
--- a/mbtrack2/instability/instabilities.py
+++ b/mbtrack2/instability/instabilities.py
@@ -4,6 +4,7 @@ General calculations about instability thresholds.
 """
 
 import numpy as np
+import matplotlib.pyplot as plt
 from scipy.constants import c, m_e, e, pi, epsilon_0
 import math
 
@@ -88,12 +89,12 @@ def cbi_threshold(ring, I, Vrf, f, beta, Ncav=1):
     
     return (Zlong, Zxdip, Zydip)
 
-def lcbi_growth_rate_mode(ring, I, Vrf, M, mu, fr=None, Rs=None, QL=None, Z=None):
+def lcbi_growth_rate_mode(ring, I, Vrf, M, mu, fr=None, RL=None, QL=None, Z=None):
     """
     Compute the longitudinal coupled bunch instability growth rate driven by
-    HOMs for a given coupled bunch mode mu [1].
+    an impedance for a given coupled bunch mode mu [1].
     
-    Use either the resonator model (fr,Rs,QL) or an Impedance object (Z).
+    Use either a list of resonators (fr, RL, QL) or an Impedance object (Z).
 
     Parameters
     ----------
@@ -106,12 +107,12 @@ def lcbi_growth_rate_mode(ring, I, Vrf, M, mu, fr=None, Rs=None, QL=None, Z=None
         Nomber of bunches in the beam.
     mu : int
         Coupled bunch mode number (= 0, ..., M-1).
-    fr : float, optional
-        Frequency of the HOM in [Hz].
-    Rs : float, optional
-        Shunt impedance of the HOM in [Ohm].
-    QL : float, optional
-        Loaded quality factor of the HOM.
+    fr : float or list, optional
+        Frequency of the resonator in [Hz].
+    RL : float or list, optional
+        Loaded shunt impedance of the resonator in [Ohm].
+    QL : float or list, optional
+        Loaded quality factor of the resonator.
     Z : Impedance, optional
         Longitunial impedance to consider.
 
@@ -131,11 +132,27 @@ def lcbi_growth_rate_mode(ring, I, Vrf, M, mu, fr=None, Rs=None, QL=None, Z=None
     nu_s = ring.synchrotron_tune(Vrf)
     factor = ring.eta * I / (4 * np.pi * ring.E0 * nu_s)
     
+    if isinstance(fr, (float, int)):
+        fr = np.array([fr])
+    elif isinstance(fr, list):
+        fr = np.array(fr)
+    if isinstance(RL, (float, int)):
+        RL = np.array([RL])
+    elif isinstance(RL, list):
+        RL = np.array(RL)
+    if isinstance(QL, (float, int)):
+        QL = np.array([QL])
+    elif isinstance(QL, list):
+        QL = np.array(QL)
+        
     if Z is None:
         omega_r = 2 * np.pi * fr
-        n_max = int(10 * omega_r / (ring.omega0 * M))
+        n_max = int(10 * omega_r.max() / (ring.omega0 * M))
         def Zr(omega):
-            return np.real(Rs / (1 + 1j * QL * (omega_r/omega - omega/omega_r)))
+            Z = 0
+            for i in range(len(fr)):
+                Z += np.real(RL[i] / (1 + 1j * QL[i] * (omega_r[i]/omega - omega/omega_r[i])))
+            return Z
     else:
         fmax = Z.data.index.max()
         n_max = int(2 * np.pi * fmax / (ring.omega0 * M))
@@ -151,12 +168,12 @@ def lcbi_growth_rate_mode(ring, I, Vrf, M, mu, fr=None, Rs=None, QL=None, Z=None
 
     return factor * sum_val
     
-def lcbi_growth_rate(ring, I, Vrf, M, fr=None, Rs=None, QL=None, Z=None):
+def lcbi_growth_rate(ring, I, Vrf, M, fr=None, RL=None, QL=None, Z=None):
     """
     Compute the maximum growth rate for longitudinal coupled bunch instability 
-    driven by HOMs [1].
+    driven an impedance [1].
     
-    Use either the resonator model (fr,Rs,QL) or an Impedance object (Z).
+    Use either a list of resonators (fr, RL, QL) or an Impedance object (Z).
 
     Parameters
     ----------
@@ -167,11 +184,11 @@ def lcbi_growth_rate(ring, I, Vrf, M, fr=None, Rs=None, QL=None, Z=None):
         Total RF voltage in [V].
     M : int
         Nomber of bunches in the beam.
-    fr : float, optional
+    fr : float or list, optional
         Frequency of the HOM in [Hz].
-    Rs : float, optional
-        Shunt impedance of the HOM in [Ohm].
-    QL : float, optional
+    RL : float or list, optional
+        Loaded shunt impedance of the HOM in [Ohm].
+    QL : float or list, optional
         Loaded quality factor of the HOM.
     Z : Impedance, optional
         Longitunial impedance to consider.
@@ -195,12 +212,70 @@ def lcbi_growth_rate(ring, I, Vrf, M, fr=None, Rs=None, QL=None, Z=None):
     """
     growth_rates = np.zeros(M)
     for i in range(M):
-        growth_rates[i] = lcbi_growth_rate_mode(ring, I, Vrf, M, i, fr=fr, Rs=Rs, QL=QL, Z=Z)
+        growth_rates[i] = lcbi_growth_rate_mode(ring, I, Vrf, M, i, fr=fr, RL=RL, QL=QL, Z=Z)
     
     growth_rate = np.max(growth_rates)
     mu = np.argmax(growth_rates)
     
     return growth_rate, mu, growth_rates
+
+def lcbi_stability_diagram(ring, I, Vrf, M, modes, cavity_list, detune_range):
+    """
+    Plot longitudinal coupled bunch instability stability diagram for a 
+    arbitrary list of CavityResonator objects around a detuning range.
+    
+    Last object in the cavity_list is assumed to be the one with the variable 
+    detuning.
+
+    Parameters
+    ----------
+    ring : Synchrotron object
+        Ring parameters.
+    I : float
+        Total beam current in [A].
+    Vrf : float
+        Total RF voltage in [V].
+    M : int
+        Nomber of bunches in the beam.
+    modes : list
+        Coupled bunch mode numbers to consider.
+    cavity_list : list
+        List of CavityResonator objects to consider, which can be:
+            - active cavities
+            - passive cavities
+            - HOMs
+            - mode dampers
+    detune_range : array
+        Detuning range (list of points) of the last CavityResonator object.
+
+    Returns
+    -------
+    fig : Figure
+        Show the shunt impedance threshold for different coupled bunch modes.
+
+    """
+    
+    Rth = np.zeros_like(detune_range)
+    fig, ax = plt.subplots()
+
+    for mu in modes:
+        fixed_gr = 0
+        for cav in cavity_list[:-1]:
+            fixed_gr += lcbi_growth_rate_mode(ring, I=I, Vrf=Vrf, mu=mu, fr=cav.fr, RL=cav.RL, QL=cav.QL, M=M)
+        
+        cav = cavity_list[-1]
+        for i, det in enumerate(detune_range):
+            gr = lcbi_growth_rate_mode(ring, I=I, Vrf=Vrf, mu=mu, fr=cav.m*ring.f1 + det, RL=cav.RL, QL=cav.QL, M=M)
+            Rth[i] = (1/ring.tau[2] - fixed_gr) * cav.Rs / gr
+
+        ax.plot(detune_range*1e-3, Rth*1e-6, label="$\mu$ = " + str(int(mu)))
+    
+    plt.xlabel(r"$\Delta f$ [kHz]")
+    plt.ylabel(r"$R_{s,max}$ $[M\Omega]$")
+    plt.yscale("log")
+    plt.legend()
+        
+    return fig
     
 def rwmbi_growth_rate(ring, current, beff, rho_material, plane='x'):
     """
diff --git a/mbtrack2/tracking/rf.py b/mbtrack2/tracking/rf.py
index b677034..aa4eb68 100644
--- a/mbtrack2/tracking/rf.py
+++ b/mbtrack2/tracking/rf.py
@@ -578,6 +578,11 @@ class CavityResonator():
     @Rs.setter
     def Rs(self, value):
         self.Rs_per_cavity = value / self.Ncav
+        
+    @property
+    def RL(self):
+        """Loaded shunt impedance [ohm]"""
+        return self.Rs / (1 + self.beta)
 
     @property
     def Q(self):
@@ -729,7 +734,7 @@ class CavityResonator():
     
     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))
+        return self.RL/(1 + 1j*self.QL*(self.fr/f - f/self.fr))
     
     def set_optimal_detune(self, I0):
         """
-- 
GitLab