Intrabeam scattering tracking implementation

Hello Salah,

Nicely done! Let's use this interface to review your code. One major thing is that documentation is missing, but we can fix this later when you solve all the other requests. Here are my requests on your IBS implementation right now.

1. (lines 13-15) You could move these imports inside your IBS class because they are only used by the IBS part of the code.
2. (line 535) The number of partitions cannot be selected without modifying the source code.
3. (line 536) I believe your self.dt is the same as self.ring.T0.
4. (line 538) You can import the classical electron radius from scipy.constants.physical_constants. However, this should be more general. Consider a situation when the code is used to compute the IBS of a hadron/proton accelerator. In that case, would you still use classical electron radius in the IBS formulas?
5. (line 541) There's no way to select the IBS model to compute T_ibs without modifying the source code.
6. (line 551,552) Why not use the same assignment as in line 546 (beta functions).
7. (lines 563-584) All three functions are the same. The only difference is their arguments. It should be a single function f (a, b, q), and you call it in your code f (a, b, q), f (1/a, b/a, q/a), and f (1/b, b/a, q/b). This will reduce the number of lines in the code and make it more readable.
8. (line 604-612) This is tracking. It should be inside the track() function. Functions get_T_ibs() should compute the T_ibs.
9. (line 621, 622) Same as (line 551, 552)
10. (line 666-676) Same consideration as (line 563-584)
11. (line 545, 614, 716, 770) These two functions share some similarities. You probably should create a separate function that does this shared part (like averaging different positions 's'). Then you call get_T_CIMP() or get_T_Piwinski() from it.
12. (line 797-815) Same consideration as (line 563-584, line 666-676)
13. (line 854) All the tracking related functionality (manipulation of x, xp, y, yp, delta, tau) and related 'kicks' should be implemented here. Model selection should be done when the object is initialised.

mbtrack2/tracking/element.py

@@ -10,9 +10,9 @@ from copy import deepcopy
 from functools import wraps
 
 import numpy as np
-from scipy.constants import elementary_charge, m_e
+from scipy.constants import elementary_charge, m_e, m_p, c, epsilon_0
 import scipy.integrate as quad
-from scipy.special import lpmv, hyp2f1
+from scipy.special import hyp2f1
 from mbtrack2.tracking.particles import Beam
 
 
@@ -515,61 +515,95 @@ def transverse_map_sector_generator(ring, positions):
 
 class IntrabeamScattering(Element):
     """
-    Intrabeam scattering method will compute the growth rate each turn and
+    Intrabeam scattering class will compute the growth rate each turn and
     increase Bunch emittance 
 
-    parameter:
-    ring: Synchrotron object
-    bunch: Bunch of Beam object
-    -----------------------------------
-    current_model: str, class variable
-        to be defined before tracking to include preferred model for tracking
+    Parameters:
+    -----------
+    ring: Object
+        Synchrotron object
+    bunch: obj
+        Bunch or Beam object
+    Current_model: Object
+        Preffered computational model for computing the growth rate T_(x,y,p)
             PS: the standard model of Piwinski
             PM: the modified model of Piwinski
             Bane: the Bane model
             CIMP: the CIMP model
+    n_points: int
+        Number of partition of the optics functions (depends on the lattice
+        complexity)
+    Particle: str
+        elementary particle of the Synchrotron 
+        "electron"
+        "proton"
+    
+    Methods:
+    --------
+    get_T_piwinski(bunch)
+    get_T_piwinski_M(bunch)
+    get_T_CIMP(bunch)
+    get_T_Bane(bunch)
+    kick(bunch, T_p, T_x, T_y)
+    track(bunch)
+
     """
-    def __init__(self, ring, bunch, current_model):
+    def __init__(self, ring, bunch, current_model, n_points, particle):
         self.ring = ring
         self.bunch = bunch
-        self.n = 1000
-        self.dt = 1/ self.ring.f0
+        self.n_points = n_points
         self.dt = self.ring.T0
         self.Spi = np.sqrt(np.pi)
-        self.r_0 = 2.8179403227 * 1e-15
-        self.s = np.linspace(0, self.ring.L, self.n)
+        self.s = np.linspace(0, self.ring.L, self.n_points)
         self.N = self.bunch.current * self.dt / elementary_charge
-        self.current_model = "PS"
+        self.current_model = ""
+        self.particle = ""
         self.sigma_s = np.std(self.bunch['tau'])
         self.sigma_p = np.std(self.bunch['delta'])
         self.sigma_px = np.std(self.bunch['xp'])
         self.sigma_py = np.std(self.bunch['yp'])
         self.beta_x, self.beta_y = self.ring.optics.beta(self.s)
-        self.dispX = self.ring.optics.dispX(self.s)
-        self.dispY = self.ring.optics.dispY(self.s)
-        self.disppX = self.ring.optics.disppX(self.s)
-        self.disppY = self.ring.optics.disppY(self.s)
-        self.alphaX = self.ring.optics.alpha(self.s)[0]
-        self.alphaY = self.ring.optics.alpha(self.s)[1]
-        self.emit_x = self.bunch.emit[0]
-        self.emit_y = self.bunch.emit[1]
-
-    # @Element.parallel
+        self.dispX, self.disppX, self.dispY, self.disppY = self.ring.optics.dispersion(self.s)
+        self.alphaX, self.alphaY = self.ring.optics.alpha(self.s)
+        self.emit_x, self.emit_y, self.emit_z = self.bunch.emit
+        self.H_x = (1 / self.beta_x) * (self.dispX**2 + 
+                                       ((self.beta_x * self.disppX) + (self.alphaX * self.dispX))**2)
+        self.H_y = (1 / self.beta_y) * (self.dispY**2 + 
+                                       ((self.beta_y * self.disppY) + (self.alphaY * self.dispY))**2)
+        self.H = (1 / self.sigma_p**2) + (self.H_x / self.emit_x) + (self.H_y / self.emit_y)
+        self.sigma_H = np.sqrt(1 / self.H)        
+        self.d = np.sqrt(np.std(bunch['x'])**2 + np.std(bunch['y'])**2)
+        self.K = (elementary_charge**2) / (4 * np.pi * epsilon_0 * c**2)
+        self.r_0 = 2.8179403227 * 1e-15
+        if self.particle == "electron":
+            self.r_0 = self.K/m_e
+        elif self.particle == "proton":
+            self.r_0 = self.K/m_p
+        self.A = (self.r_0**2 * self.N) / (64 * np.pi**2 * self.ring.beta**3 * self.ring.gamma**4 * self.emit_x * self.emit_y * self.sigma_s * self.sigma_p)
+
     def get_T_piwinski(self, bunch):
-        from scipy.constants import elementary_charge, m_e
-        import scipy.integrate as quad
-        from scipy.special import lpmv, hyp2f1
-        A = (self.r_0**2 *  self.N) / (64 * np.pi**2 * self.ring.beta**3 * self.ring.gamma**4 * self.emit_x * self.emit_y * self.sigma_s * self.sigma_p)
-        
-        h = (1 / self.sigma_p**2) + (self.dispX**2 / (self.beta_x * self.emit_x)) + (self.dispY**2 / (self.beta_y * self.emit_y))
+        """
+        Compute the growth rate T_(x,y,p) using Piwinski standard formalism
 
-        sigma_h = np.sqrt(1 / h) 
-        d = np.sqrt(np.std(self.bunch['x'])**2 + np.std(self.bunch['y'])**2) 
+        Parameters
+        ----------
+        bunch : Object
+            bunch or beam object.
 
+        Returns
+        -------
+        T_p, T_x, T_y: float
+            growth rate on z,x,y respectively
+            
+        Reference:
+        ----------
+        [1]A. Piwinski, Intra-Beam-Scattering, (1974).
+        """
+        h = (1 / self.sigma_p**2) + (self.dispX**2 / (self.beta_x * self.emit_x)) + (self.dispY**2 / (self.beta_y * self.emit_y))
+        sigma_h = np.sqrt(1 / h) 
         a_bar = (sigma_h / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
         b_bar = (sigma_h / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
-
-        q_bar = sigma_h * self.ring.beta * np.sqrt(2 * d / self.r_0) 
+        q_bar = sigma_h * self.ring.beta * np.sqrt(2 * self.d / self.r_0) 
         # -------------------------------------------------------------------------
         def scattering_1bq(u, i, a_bar, b_bar, q_bar):
             P2 = (1/a_bar[i]**2) + ((1 - (1/a_bar[i]**2)) * u**2)
@@ -595,7 +629,7 @@ class IntrabeamScattering(Element):
             return f_abq
         #-------------------------------------------------------------------------
         vabq, v1aq, v1bq = [], [], []
-        for i in range(self.n):
+        for i in range(self.n_points):
             el_abq, err = quad.quad(scattering_abq, 0, 1, args=(i, a_bar, b_bar, q_bar))
             el_1aq, err = quad.quad(scattering_1aq, 0, 1, args=(i, a_bar, b_bar, q_bar))
             el_1bq, err = quad.quad(scattering_1bq, 0, 1, args=(i, a_bar, b_bar, q_bar))
@@ -603,35 +637,37 @@ class IntrabeamScattering(Element):
             v1aq.append(el_1aq)
             v1bq.append(el_1bq)
 
-
-        T_p = A * (vabq * (sigma_h**2 / self.sigma_p**2)) 
-        T_x = A * (v1bq + (vabq * ((self.dispX**2 * sigma_h**2) / (self.beta_x * self.emit_x))))
-        T_y = A * (v1aq + (vabq * ((self.dispY**2 * sigma_h**2) / (self.beta_y * self.emit_y))))
+        T_p = self.A * (vabq * (sigma_h**2 / self.sigma_p**2)) 
+        T_x = self.A * (v1bq + (vabq * ((self.dispX**2 * sigma_h**2) / (self.beta_x * self.emit_x))))
+        T_y = self.A * (v1aq + (vabq * ((self.dispY**2 * sigma_h**2) / (self.beta_y * self.emit_y))))
         T_p = np.average(T_p)
         T_x = np.average(T_x)
         T_y = np.average(T_y)
         return T_p, T_x, T_y 
         
     def get_T_CIMP(self, bunch):
-        from scipy.constants import elementary_charge, m_e
-        import scipy.integrate as quad
-        from scipy.special import lpmv, hyp2f1
-        H_x = (1 / self.beta_x) * (self.dispX**2 + 
-                                       ((self.beta_x * self.disppX) + (self.alphaX * self.dispX))**2)
-        H_y = (1 / self.beta_y) * (self.dispY**2 + 
-                                       ((self.beta_y * self.disppY) + (self.alphaY * self.dispY))**2)
-        H = (1 / self.sigma_p**2) + (H_x / self.emit_x) + (H_y / self.emit_y)
-        sigma_H = np.sqrt(1 / H)
-        a = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
-        b = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
-        
-        A = (self.r_0**2 * self.N) / (64 * np.pi**2 * self.ring.beta**3 * self.ring.gamma**4 * self.emit_x * self.emit_y * self.sigma_s * self.sigma_p)
-        
-        d = np.sqrt(np.std(bunch['x'])**2 + np.std(bunch['y'])**2)
-
-        q = sigma_H * self.ring.beta * np.sqrt(2 * d / self.r_0)
+        """
+        Compute the growth rate T_(x,y,p) using Completely integrated modified 
+        Piwinski formalism
 
+        Parameters
+        ----------
+        bunch : Object
+            bunch or beam object.
 
+        Returns
+        -------
+        T_p, T_x, T_y: float
+            growth rate on z,x,y respectively
+            
+        Reference:
+        ----------
+        [1]K. Kubo, S. K. Mtingwa, and A. Wolski, Intrabeam Scattering Formulas
+        for High Energy Beams, Phys. Rev. ST Accel. Beams 8, 081001 (2005).
+        """
+        a = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
+        b = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
+        q = self.sigma_H * self.ring.beta * np.sqrt(2 * self.d / self.r_0)
         # ---------------------------------------
         def Puv(u, v, x):
 
@@ -648,15 +684,6 @@ class IntrabeamScattering(Element):
             else:
                 g_val = np.sqrt(np.pi / u) * ((Puv(0, -.5, x_arg)) - ((3/2) * (Puv(-1, -.5, x_arg))))
             return g_val
-            
-        def g(u):
-            if u <= 10:
-                vg = (2.691 * (1 - (0.2288964 / u))) / ((1 + 0.16 * u) * (1 + (1.35 * np.exp(-u/0.2))))
-            else: 
-                vg = 0
-            return vg
-        
-
         
         def scattering_1bq(i, a, b, q_bar):
             f_1bq = ((-4 * np.pi**(3/4) * np.log(q_bar[i] / a[i])) / (1 / a[i])) * g_func(b[i] / a[i])
@@ -670,9 +697,8 @@ class IntrabeamScattering(Element):
             f_abq = -(scattering_1aq(i, a, b, q_bar) * (1/b[i]**2)) - (scattering_1bq(i, a, b, q_bar) * (1/a[i]**2))
             return f_abq
         
-        
         vabq, v1aq, v1bq = [], [], []
-        for i in range(self.n):
+        for i in range(self.n_points):
             el_abq = scattering_abq(i, a, b, q)
             el_1aq = scattering_1aq(i, a, b, q)
             el_1bq = scattering_1bq(i, a, b, q)
@@ -680,40 +706,36 @@ class IntrabeamScattering(Element):
             v1aq.append(el_1aq)
             v1bq.append(el_1bq)
         
-        
-        T_p = A * (vabq * (sigma_H**2 / self.sigma_p**2))
-        T_x = A * (v1bq + (vabq * ((self.dispX**2 * sigma_H**2) / (self.beta_x * self.emit_x))))
-        T_y = A * (v1aq + (vabq * ((self.dispY**2 * sigma_H**2) / (self.beta_y * self.emit_y))))
-
-
-        g_ab, g_ba = [], []
-        for i in range(self.n):
-            val_ab = g_func(a[i] / b[i])
-            val_ba = g_func(b[i] / a[i]) 
-            g_ab.append(val_ab)
-            g_ba.append(val_ba)
-
+        T_p = self.A * (vabq * (self.sigma_H**2 / self.sigma_p**2))
+        T_x = self.A * (v1bq + (vabq * ((self.dispX**2 * self.sigma_H**2) / (self.beta_x * self.emit_x))))
+        T_y = self.A * (v1aq + (vabq * ((self.dispY**2 * self.sigma_H**2) / (self.beta_y * self.emit_y))))
         T_p = np.average(T_p)
         T_x = np.average(T_x)
         T_y = np.average(T_y)
         return T_p, T_x, T_y
 
-
     def get_T_Bane(self, bunch):
-        from scipy.constants import elementary_charge, m_e
-        import scipy.integrate as quad
-        from scipy.special import lpmv, hyp2f1
+        """
+        Compute the growth rate T_(x,y,p) using Bane formalism 
 
-        H_x = (1 / self.beta_x) * (self.dispX**2 + 
-                                       ((self.beta_x * self.disppX) + (self.alphaX * self.dispX))**2)
-        H_y = (1 / self.beta_y) * (self.dispY**2 + 
-                                       ((self.beta_y * self.disppY) + (self.alphaY * self.dispY))**2)
-        H = (1 / self.sigma_p**2) + (H_x / self.emit_x) + (H_y / self.emit_y)
-        sigma_H = np.sqrt(1 / H)
-        a = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
-        b = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
-        
+        Parameters
+        ----------
+        bunch : Object
+            bunch or beam object.
 
+        Returns
+        -------
+        T_p, T_x, T_y: float
+            growth rate on z,x,y respectively
+            
+        Reference:
+        ----------
+        [1]K. L. F. Bane, H. Hayano, K. Kubo, T. Naito, T. Okugi, and J. Urakawa, 
+        Intrabeam Scattering Analysis of Measurements at KEK’s Accelerator Test 
+        Facility Damping Ring, Phys. Rev. ST Accel. Beams 5, 084403 (2002).
+        """
+        a = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
+        b = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
         C_log = np.log((self.ring.gamma**2 * self.emit_x * np.sqrt(self.beta_y * self.emit_y)) / 
                        (self.r_0 * self.beta_x))
         C_a = a / b
@@ -723,39 +745,43 @@ class IntrabeamScattering(Element):
             return g_val
 
         integrate_gval = []
-        for j in range(self.n):
+        for j in range(self.n_points):
             reslt, err = quad.quad(g_func, 0, np.inf, args=(j, C_a))
             integrate_gval.append(reslt)
 
-        T_pp = (self.r_0**2 * self.N * C_log * sigma_H * integrate_gval * (self.beta_x * 
+        T_pp = (self.r_0**2 * self.N * C_log * self.sigma_H * integrate_gval * (self.beta_x * 
                                                                            self.beta_y)**(-1/4)) / (16 * self.ring.gamma**3 * 
                                                                                                                          self.emit_x**(3/4) * self.emit_y**(3/4) * self.sigma_s * self.sigma_p**3)
         T_p = np.average(T_pp)
-        T_x = (self.sigma_p**2 * H_x * T_pp) / self.emit_x
-        T_y = (self.sigma_p**2 * H_y * T_pp) / self.emit_y
-
+        T_x = (self.sigma_p**2 * self.H_x * T_pp) / self.emit_x
+        T_y = (self.sigma_p**2 * self.H_y * T_pp) / self.emit_y
         T_x = np.average(T_x)
         T_y = np.average(T_y)
         return T_p, T_x, T_y
         
     def get_T_piwinski_M(self, bunch):
-        from scipy.constants import elementary_charge, m_e
-        import scipy.integrate as quad
-        from scipy.special import lpmv, hyp2f1
-
-        A = (self.r_0**2 * self.N) / (64 * np.pi**2 * self.ring.beta**3 * self.ring.gamma**4 * self.emit_x * self.emit_y * self.sigma_s * self.sigma_p)
+        """
+        Compute the growth rate T_(x,y,p) Modified Piwinski formalism 
 
-        d = np.sqrt(np.std(bunch['x'])**2 + np.std(bunch['y'])**2)
+        Parameters
+        ----------
+        bunch : Object
+            bunch or beam object.
 
-        H_x = (1 / self.beta_x) * (self.dispX**2 + 
-                                       ((self.beta_x * self.disppX) + (self.alphaX * self.dispX))**2)
-        H_y = (1 / self.beta_y) * (self.dispY**2 + 
-                                       ((self.beta_y * self.disppY) + (self.alphaY * self.dispY))**2)
-        H = (1 / self.sigma_p**2) + (H_x / self.emit_x) + (H_y / self.emit_y)
-        sigma_H = np.sqrt(1 / H)
-        a = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
-        b = (sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
-        q = sigma_H * self.ring.beta * np.sqrt( 2 * d / self.r_0)
+        Returns
+        -------
+        T_p, T_x, T_y: float
+            growth rate on z,x,y respectively
+            
+        Reference:
+        ----------
+        [1]K. L. F. Bane, A Simplified Model of Intrabeam Scattering, in 8th
+        European Particle Accelerator Conference (Paris, France, 2002), 
+        p. 1443.
+        """
+        a = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_x / self.emit_x)
+        b = (self.sigma_H / self.ring.gamma) * np.sqrt(self.beta_y / self.emit_y)
+        q = self.sigma_H * self.ring.beta * np.sqrt( 2 * self.d / self.r_0)
         ######--------------------defining scattering functions
         def scattering_1bq(u, i, a, b, q):
             P2 = (1/a[i]**2) + ((1 - (1/a[i]**2)) * u**2)
@@ -779,8 +805,7 @@ class IntrabeamScattering(Element):
         #####----------------------------
         #computing intergrals
         vabq, v1aq, v1bq = [], [], []
-        for i in range(self.n):
-
+        for i in range(self.n_points):
             el_1aq, err = quad.quad(scattering_1aq, 0, 1, args=(i, a, b, q))
             el_1bq, err = quad.quad(scattering_1bq, 0, 1, args=(i, a, b, q))
             el_abq = -(el_1aq * (1/b[i]**2)) - (el_1bq * (1/a[i]**2))
@@ -790,12 +815,10 @@ class IntrabeamScattering(Element):
     
         vabq = np.array(vabq)
         v1aq = np.array(v1aq)
-        v1bq = np.array(v1bq)
-    
-    
-        T_p = A * (vabq * (sigma_H**2 / self.sigma_p**2))
-        T_x = A * (v1bq + (vabq * ((self.dispX**2 * sigma_H**2) / (self.beta_x * self.emit_x))))
-        T_y = A * (v1aq + (vabq * ((self.dispY**2 * sigma_H**2) / (self.beta_y * self.emit_y))))
+        v1bq = np.array(v1bq)    
+        T_p = self.A * (vabq * (self.sigma_H**2 / self.sigma_p**2))
+        T_x = self.A * (v1bq + (vabq * ((self.dispX**2 * self.sigma_H**2) / (self.beta_x * self.emit_x))))
+        T_y = self.A * (v1aq + (vabq * ((self.dispY**2 * self.sigma_H**2) / (self.beta_y * self.emit_y))))
         T_p = np.average(T_p)
         T_x = np.average(T_x)
         T_y = np.average(T_y)
@@ -815,10 +838,14 @@ class IntrabeamScattering(Element):
     @Element.parallel
     def track(self, bunch):
         """
-        Tracking method of Intrabeamscattering
+        Tracking method of Intrabeamscattering takes T_(x,y,p) and apply 
+        momentum change to the coordinates of macroparticles inside the bunch
+        
         parameters:
-        bunch: Bunch or Beam object
-        return 1/Tx, 1/Ty, 1/Tp
+        -----------
+        bunch: obj 
+            Bunch or Beam object
+        
         """
         if self.current_model == "PS":
             T_p, T_x, T_y = self.get_T_piwinski(bunch)
@@ -832,4 +859,4 @@ class IntrabeamScattering(Element):
         if self.current_model == "Bane":
             T_p, T_x, T_y = self.get_T_Bane(bunch)
             self.kick(bunch, T_p, T_x, T_y)
-
+