Add RFBucket class

tracking/rfbucket.py

+# -*- coding: utf-8 -*-
+"""
+Created on Wed Mar 16 17:54:36 2022
+
+@author: foosang
+"""
+
+import numpy as np
+import matplotlib.pyplot as plt
+from scipy.optimize import fsolve
+
+class RFBucket:
+    """
+    Calculate the RF bucket for a given synchronous phase. Valid for both
+    sine and cosine RF convention, but limited to the following intervals:
+        sine convention : [0,180°],
+        cosine convention : [-90°,90°].
+        
+    Parameters
+    ----------
+    phi_s : float
+        Synchronous phase in [rad].
+    convention : {"cos", "sin"}, optional
+        RF convention corresponding to input phi_s. Use "cos" or "sin" for a 
+        cosine or sine convention, respectively. The default is "cos".
+        
+    Attributes
+    ----------
+    phi_s : float
+        Synchronous phase in [rad].
+    convention : str
+        RF convention.
+    above_transition : bool
+        Depending on phi_s and convention, above_transition is True when the
+        system is above the transition energy and is False if the system is 
+        below the transition energy.
+    """
+  
+    def __init__(self, phi_s, convention="cos"):
+        if convention == "sin":
+            if phi_s<0 or phi_s>np.pi:
+                raise ValueError("The synchronous phase is outside the valid range.")
+            else:
+                if phi_s>=0 and phi_s<=np.pi/2:
+                    self.above_transition = False
+                else:
+                    self.above_transition = True
+        elif convention == "cos":
+            if phi_s<-np.pi/2 or phi_s>np.pi/2:
+                raise ValueError("The synchronous phase is outside the valid range.")
+            else:
+                if phi_s>=-np.pi/2 and phi_s<=0:
+                    self.above_transition = False
+                else:
+                    self.above_transition = True
+        else:
+            raise ValueError("Invalid RF convention.")
+        
+        self.phi_s = phi_s
+        self.convention = convention
+        if self.convention == "cos":
+            if self.above_transition is True:
+                self._phi_s = np.pi - (np.pi/2 - phi_s)
+            else:
+                self._phi_s = np.pi/2 - phi_s
+                
+        elif self.convention == "sin":
+            if self.above_transition is True:
+                self._phi_s = np.pi - phi_s
+            else:
+                self._phi_s = phi_s
+        
+    def separatrix(self, phi):
+        """
+        Calculate the separatrix in longitudinal phase space.
+
+        Parameters
+        ----------
+        phi : array of float
+            Phase interval on which the separatrix will be calculated in [rad]
+
+        Returns
+        -------
+        y1 : array
+            The positive half of the separatrix in the unit of (dphi/dt)/Omega_s,
+            where dphi/dt is phase velocity and Omega_s is synchrotron 
+            oscillation frequency.
+        y2 : array
+            The negative half of the separatrix in the same unit as y1. 
+
+        """
+        
+        y1 = np.sqrt(2/np.cos(self._phi_s)*(np.cos(phi)+phi*np.sin(self._phi_s)-
+                np.cos(np.pi-self._phi_s)-(np.pi-self._phi_s)*np.sin(self._phi_s)))
+        
+        y2 = -1*y1
+        
+        return y1, y2
+    
+    def separatrix_root(self, x):
+        """
+        Function used for determining the root of the separatrix equation.
+
+        """
+        
+        value =  np.cos(x) + x*np.sin(self._phi_s) - \
+            np.cos(np.pi-self._phi_s) - (np.pi-self._phi_s)*np.sin(self._phi_s)
+        return value
+    
+    def plot_separatrix(self, phi, unit="deg"):
+        """
+        Plot separatrix
+
+        Parameters
+        ----------
+        phi : array of float
+            Phase interval on which the separatrix will be calculated in [rad]
+        unit : {"deg", "rad"}, optional
+            Unit of horizontal axis. Use "deg" for degree or "rad" for radian.
+            The default is "deg".
+
+        Returns
+        -------
+        fig : Figure
+            Figure object with the separatrix plotted on it.
+
+        """
+        
+        y1, y2 = self.separatrix(phi)
+        if self.convention == "cos":
+            if self.above_transition is True:
+                phi = phi - np.pi/2             
+            else:
+                phi = -1*(phi - np.pi/2)
+        
+        elif self.convention == "sin":
+            if self.above_transition is True:
+                phi = np.pi - phi
+        
+        fig, ax = plt.subplots()
+        
+        if unit == "deg":
+            phi = phi/np.pi * 180
+            xlabel = "$\\phi$ (deg)"
+        elif unit == "rad":
+            xlabel = "$\\phi$ (rad)"
+            
+        ax.plot(phi, y1, color="tab:blue")
+        ax.plot(phi, y2, color="tab:blue")
+        ax.set_xlabel(xlabel)
+        ax.set_ylabel('$\\dot{\\phi} / \\Omega_s$')
+        
+        return fig
+            
+    def bucket_height(self, ring, Vrf, E0, plane):
+        """
+        Calculate the height of the RF bucket which corresponds to the maximum
+        energy deviation.
+
+        Parameters
+        ----------
+        ring : Synchrotron object
+        Vrf : float
+            RF voltage in [V].
+        E0 : float
+            Beam energy in [eV].
+        plane : {"x", "y"}
+            The considered plane of motion of the beam. Use "x" for the
+            horizontal plane or "y" for the vertical plane.
+
+        Returns
+        -------
+        delta_max : float
+            Maximum energy deviation.
+
+        """
+        G = np.sqrt(0.5* abs(2*np.cos(self._phi_s)-(np.pi-2*self._phi_s)*
+                             np.sin(self._phi_s)))
+          
+        plane_dict = {'x':0, 'y':1}
+        index = plane_dict[plane]
+        beta = ring.optics.local_beta[index]
+        h = ring.h
+        eta = ring.eta
+        
+        delta_max = G * np.sqrt(2*beta**2*E0*Vrf/(np.pi*h*abs(eta))) / E0
+        
+        return delta_max
+    
+    def extreme_phases(self):
+        """
+        Calculate the two extreme phases on the separatrix.
+
+        Returns
+        -------
+        phi_min : float
+            The minimum phase in [rad].
+        phi_max : float
+            The maximum phase in [rad].
+
+        """
+        
+        if self._phi_s > np.pi/2 and self._phi_s <= np.pi:
+            x0 = 270 / 180 * np.pi # Iinitial estimate of the root
+            phi_min0 = np.pi - self._phi_s
+            phi_max0 = fsolve(self.separatrix_root, x0=x0)[0]
+            
+        elif self._phi_s >= 0 and self._phi_s <= np.pi/2:
+            x0 = -90 / 180 * np.pi # Initial estimate of the root
+            phi_max0 = np.pi - self._phi_s
+            phi_min0 = fsolve(self.separatrix_root, x0=x0)[0]
+            
+        if self.convention == "cos":
+            if self.above_transition is True:
+                phi_min = phi_min0 - np.pi/2
+                phi_max = phi_max0 - np.pi/2                
+            else:
+                phi_min = -1*(phi_max0 - np.pi/2)
+                phi_max = -1*(phi_min0 - np.pi/2)
+        
+        elif self.convention == "sin":
+            if self.above_transition is True:
+                phi_min = np.pi - phi_max0
+                phi_max = np.pi - phi_min0
+            else:
+                phi_min = phi_min0
+                phi_max = phi_max0
+            
+        return phi_min, phi_max
+    
+    
\ No newline at end of file