Faster CircularResistiveWall

Update CircularResistiveWall class in resistive_wall.py to replace numerical integrations with an analytical forms for improved performance.

• Numerical integrations are removed.
• LongWakeExact and TransWakeExact are used even in the case where time > 20*t0.
• "atol" for LongWakeExact is removed.

I believe "atol" is not needed when considering the fundamental theorem of beam loading. This is because Wl(0) = Wl(0+)/2 is comes from the equation Wl(0) = [Wl(0-)*l/2 + Wl(0+)*l/2] / l, which represents the average Wl value of a very narrow bin (as the length of bin l -> 0) at the origin. Defining "atol" for the fundamental theorem of beam loading introduces the possibility that more than one point could be divided by 2, which I believe is mathematically and physically unacceptable. Thus, it is correct to consider the fundamental theorem of beam loading only at the origin. PyHEADTAIL also considers the fundamental theorem of beam loading only at the origin. Please see def function_longitudinal at the link (line 341).

mbtrack2/impedance/resistive_wall.py

@@ -5,15 +5,15 @@ Define resistive wall elements based on the WakeField class.
 
 import numpy as np
 from scipy.constants import c, epsilon_0, mu_0
-from scipy.integrate import quad
 
 from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction
+from mbtrack2.tracking.particles_electromagnetic_fields import _wofz
 
 
 def skin_depth(frequency, rho, mu_r=1, epsilon_r=1):
     """
     General formula for the skin depth.
-    
+
     Parameters
     ----------
     frequency : array of float
@@ -24,12 +24,12 @@ def skin_depth(frequency, rho, mu_r=1, epsilon_r=1):
         Relative magnetic permeability.
     epsilon_r : float, optional
         Relative electric permittivity.
-    
+
     Returns
     -------
     delta : array of float
         Skin depth in [m].
-    
+
     """
 
     delta = (np.sqrt(2 * rho / (np.abs(2 * np.pi * frequency) * mu_r * mu_0)) *
@@ -43,49 +43,46 @@ def skin_depth(frequency, rho, mu_r=1, epsilon_r=1):
 class CircularResistiveWall(WakeField):
     """
     Resistive wall WakeField element for a circular beam pipe.
-    
+
     Impedance from approximated formulas from Eq. (2.77) of Chao book [1].
-    Wake function formulas from [2].
-        
+    Wake function formulas from [2, 3] and the fundamental theorem of
+    beam loading from [4].
+
     Parameters
     ----------
     time : array of float
         Time points where the wake function will be evaluated in [s].
     frequency : array of float
         Frequency points where the impedance will be evaluated in [Hz].
-    length : float 
+    length : float
         Beam pipe length in [m].
     rho : float
         Resistivity in [ohm.m].
-    radius : float 
+    radius : float
         Beam pipe radius in [m].
     exact : bool, optional
-        If False, approxmiated formulas are used for the wake function 
+        If False, approxmiated formulas are used for the wake function
         computations.
-    atol : float, optional
-        Absolute tolerance used to enforce fundamental theorem of beam loading
-        for the exact expression of the longitudinal wake function.
-        Default is 1e-20.
-    
+        The default is True.
+
     References
     ----------
-    [1] : Chao, A. W. (1993). Physics of collective beam instabilities in high 
+    [1] : Chao, A. W. (1993). Physics of collective beam instabilities in high
     energy accelerators. Wiley.
-    [2] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
-    interbunch collective beam motions in storage rings." Nuclear Instruments 
-    and Methods in Physics Research Section A: Accelerators, Spectrometers, 
+    [2] : Ivanyan, Mikayel I., and Vasili M. Tsakanov. "Analytical treatment of
+    resistive wake potentials in round pipes." Nuclear Instruments and Methods
+    in Physics Research Section A: Accelerators, Spectrometers, Detectors and
+    Associated Equipment 522, no. 3 (2004): 223-229.
+    [3] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and
+    interbunch collective beam motions in storage rings." Nuclear Instruments
+    and Methods in Physics Research Section A: Accelerators, Spectrometers,
     Detectors and Associated Equipment 806 (2016): 221-230.
+    [4] : Zotter, Bruno W., and Semyon A. Kheifets (1998). Impedances and wakes
+    in high-energy particle accelerators. World Scientific.
 
     """
 
-    def __init__(self,
-                 time,
-                 frequency,
-                 length,
-                 rho,
-                 radius,
-                 exact=False,
-                 atol=1e-20):
+    def __init__(self, time, frequency, length, rho, radius, exact=True):
         super().__init__()
 
         self.length = length
@@ -100,7 +97,7 @@ class CircularResistiveWall(WakeField):
         Z2 = c / omega * length * (1 + np.sign(frequency) * 1j) * rho / (
             np.pi * radius**3 * skin_depth(frequency, rho))
 
-        Wl = self.LongitudinalWakeFunction(time, exact, atol)
+        Wl = self.LongitudinalWakeFunction(time, exact)
         Wt = self.TransverseWakeFunction(time, exact)
 
         Zlong = Impedance(variable=frequency,
@@ -123,134 +120,151 @@ class CircularResistiveWall(WakeField):
         super().append_to_model(Wxdip)
         super().append_to_model(Wydip)
 
-    def LongitudinalWakeFunction(self, time, exact=False, atol=1e-20):
+    def LongitudinalWakeFunction(self, time, exact=True):
         """
-        Compute the longitudinal wake function of a circular resistive wall 
-        using Eq. (22), or approxmiated expression Eq. (24), of [1]. The 
-        approxmiated expression is valid if the time is large compared to the 
-        characteristic time t0.
-        
-        If some time value is smaller than atol, then the fundamental theorem 
-        of beam loading is applied: Wl(0) = Wl(0+)/2.
+        Compute the longitudinal wake function of a circular resistive wall
+        using Eq. (11), of [1], or approxmiated expression Eq. (24), of [2].
+        The approxmiated expression is valid if the time is large compared
+        to the characteristic time t0.
+
+        Eq. (11) in [1] is completely identical to Eq. (22) in [2].
+
+        The real parts of the last two terms of Eq. (11) in [1] are the same,
+        and the imaginary parts have the same magnitude but opposite signs.
+        Therefore, the former term was doubled, the latter term was eliminated,
+        and only the real part was taken to speed up the calculation.
+
+        The fundamental theorem of beam loading [3] is applied for the exact
+        expression of the longitudinal wake function: Wl(0) = Wl(0+)/2.
 
         Parameters
         ----------
         time : array of float
             Time points where the wake function is evaluated in [s].
         exact : bool, optional
-            If True, the exact expression is used. The default is False.
-        atol : float, optional
-            Absolute tolerance used to enforce fundamental theorem of beam loading
-            for the exact expression of the longitudinal wake function.
-            Default is 1e-20.
+            If True, the exact expression is used. The default is True.
 
         Returns
         -------
         wl : array of float
             Longitudinal wake function in [V/C].
-            
+
         References
         ----------
-        [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
-        interbunch collective beam motions in storage rings." Nuclear Instruments 
-        and Methods in Physics Research Section A: Accelerators, Spectrometers, 
+        [1] : Ivanyan, Mikayel I., and Vasili M. Tsakanov. "Analytical treatment of
+        resistive wake potentials in round pipes." Nuclear Instruments and Methods
+        in Physics Research Section A: Accelerators, Spectrometers, Detectors and
+        Associated Equipment 522, no. 3 (2004): 223-229.
+        [2] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and
+        interbunch collective beam motions in storage rings." Nuclear Instruments
+        and Methods in Physics Research Section A: Accelerators, Spectrometers,
         Detectors and Associated Equipment 806 (2016): 221-230.
+        [3] : Zotter, Bruno W., and Semyon A. Kheifets (1998). Impedances and wakes
+        in high-energy particle accelerators. World Scientific.
         """
         wl = np.zeros_like(time)
         idx1 = time < 0
-        wl[idx1] = 0
         if exact == True:
-            idx2 = time > 20 * self.t0
+            idx2 = time == 0
             idx3 = np.logical_not(np.logical_or(idx1, idx2))
-            wl[idx3] = self.__LongWakeExact(time[idx3], atol)
+            factor = (self.Z0 * c / (3 * np.pi * self.radius**2) * self.length)
+            if np.any(idx2):
+                # fundamental theorem of beam loading
+                wl[idx2] = 3 * factor / 2
+            wl[idx3] = self.__LongWakeExact(time[idx3], factor)
         else:
             idx2 = np.logical_not(idx1)
-        wl[idx2] = self.__LongWakeApprox(time[idx2])
+            wl[idx2] = self.__LongWakeApprox(time[idx2])
         return wl
 
-    def TransverseWakeFunction(self, time, exact=False):
+    def TransverseWakeFunction(self, time, exact=True):
         """
-        Compute the transverse wake function of a circular resistive wall 
-        using Eq. (25), or approxmiated expression Eq. (26), of [1]. The 
-        approxmiated expression is valid if the time is large compared to the 
-        characteristic time t0.
-        
-        Exact expression (Eq. (25) from [1]) is corrected by factor (c * t0).
+        Compute the transverse wake function of a circular resistive wall
+        using Eq. (11), of [1], or approxmiated expression Eq. (26), of [2].
+        The approxmiated expression is valid if the time is large compared
+        to the characteristic time t0.
+
+        Eq. (11) in [1] is completely identical to Eq. (25) in [2].
+
+        There are typos in both Eq. (11) in [1] and Eq. (25) in [2].
+        Corrected the typos in the last two terms of exact expression for
+        transverse wake function in Eq. (11), of [1].
+        It is necessary to multiply Eq. (25) in [2] by c*t0.
+
+        The real parts of the last two terms of Eq. (11) in [1] are the same,
+        and the imaginary parts have the same magnitude but opposite signs.
+        Therefore, the former term was doubled, the latter term was eliminated,
+        and only the real part was taken to speed up the calculation.
 
         Parameters
         ----------
         time : array of float
             Time points where the wake function is evaluated in [s].
         exact : bool, optional
-            If True, the exact expression is used. The default is False.
+            If True, the exact expression is used. The default is True.
 
         Returns
         -------
         wt : array of float
-            Transverse wake function in [V/C].
-            
+            Transverse wake function in [V/C/m].
+
         References
         ----------
-        [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
-        interbunch collective beam motions in storage rings." Nuclear Instruments 
-        and Methods in Physics Research Section A: Accelerators, Spectrometers, 
+        [1] : Ivanyan, Mikayel I., and Vasili M. Tsakanov. "Analytical treatment of
+        resistive wake potentials in round pipes." Nuclear Instruments and Methods
+        in Physics Research Section A: Accelerators, Spectrometers, Detectors and
+        Associated Equipment 522, no. 3 (2004): 223-229.
+        [2] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and
+        interbunch collective beam motions in storage rings." Nuclear Instruments
+        and Methods in Physics Research Section A: Accelerators, Spectrometers,
         Detectors and Associated Equipment 806 (2016): 221-230.
         """
         wt = np.zeros_like(time)
         idx1 = time < 0
-        wt[idx1] = 0
         if exact == True:
-            idx2 = time > 20 * self.t0
+            idx2 = time == 0
             idx3 = np.logical_not(np.logical_or(idx1, idx2))
-            wt[idx3] = self.__TransWakeExact(time[idx3])
+            factor = ((self.Z0 * c**2 * self.t0) /
+                      (3 * np.pi * self.radius**4) * self.length)
+            wt[idx3] = self.__TransWakeExact(time[idx3], factor)
         else:
             idx2 = np.logical_not(idx1)
-        wt[idx2] = self.__TransWakeApprox(time[idx2])
+            wt[idx2] = self.__TransWakeApprox(time[idx2])
         return wt
 
-    def __LongWakeExact(self, time, atol):
-        wl = np.zeros_like(time)
-        factor = 4 * self.Z0 * c / (np.pi * self.radius**2) * self.length
-        for i, t in enumerate(time):
-            val, err = quad(lambda z: self.__function(t, z), 0, np.inf)
-            wl[i] = factor * (
-                np.exp(-t / self.t0) / 3 * np.cos(np.sqrt(3) * t / self.t0) -
-                np.sqrt(2) / np.pi * val)
-            if np.isclose(0, t, atol=atol):
-                wl[i] = wl[i] / 2
+    def __LongWakeExact(self, t, factor):
+        w1re, _ = _wofz(0, np.sqrt(2 * t / self.t0))
+        w2re, _ = _wofz(
+            np.cos(np.pi / 6) * np.sqrt(2 * t / self.t0),
+            np.sin(np.pi / 6) * np.sqrt(2 * t / self.t0))
+
+        wl = factor * (4 * np.exp(-1 * t / self.t0) *
+                       np.cos(np.sqrt(3) * t / self.t0) + w1re - 2*w2re)
         return wl
 
-    def __TransWakeExact(self, time):
-        wt = np.zeros_like(time)
-        factor = ((8 * self.Z0 * c**2 * self.t0) / (np.pi * self.radius**4) *
-                  self.length)
-        for i, t in enumerate(time):
-            val, err = quad(lambda z: self.__function2(t, z), 0, np.inf)
-            wt[i] = factor * (
-                1 / 12 *
-                (-1 * np.exp(-t / self.t0) * np.cos(np.sqrt(3) * t / self.t0) +
-                 np.sqrt(3) * np.exp(-t / self.t0) *
-                 np.sin(np.sqrt(3) * t / self.t0)) - np.sqrt(2) / np.pi * val)
+    def __TransWakeExact(self, t, factor):
+        w1re, _ = _wofz(0, np.sqrt(2 * t / self.t0))
+        w2re, w2im = _wofz(
+            np.cos(np.pi / 6) * np.sqrt(2 * t / self.t0),
+            np.sin(np.pi / 6) * np.sqrt(2 * t / self.t0))
+
+        wt = factor * (2 * np.exp(-1 * t / self.t0) *
+                       (np.sqrt(3) * np.sin(np.sqrt(3) * t / self.t0) -
+                        np.cos(np.sqrt(3) * t / self.t0)) + w1re + 2 *
+                       (np.cos(-np.pi / 3) * w2re - np.sin(-np.pi / 3) * w2im))
         return wt
 
     def __LongWakeApprox(self, t):
-        wl = -1 * (1 / (4 * np.pi * self.radius) * np.sqrt(self.Z0 * self.rho /
-                                                           (c * np.pi)) /
-                   t**(3 / 2)) * self.length
+        wl = -1 * (1 / (4 * np.pi * self.radius) *
+                   np.sqrt(self.Z0 * self.rho /
+                           (c * np.pi * t**3)) * self.length)
         return wl
 
     def __TransWakeApprox(self, t):
         wt = (1 / (np.pi * self.radius**3) *
-              np.sqrt(self.Z0 * c * self.rho / np.pi) / t**(1 / 2) *
-              self.length)
+              np.sqrt(self.Z0 * c * self.rho / (np.pi * t)) * self.length)
         return wt
 
-    def __function(self, t, x):
-        return ((x**2 * np.exp(-1 * (x**2) * t / self.t0)) / (x**6 + 8))
-
-    def __function2(self, t, x):
-        return ((-1 * np.exp(-1 * (x**2) * t / self.t0)) / (x**6 + 8))
-
 
 class Coating(WakeField):
 
@@ -264,7 +278,7 @@ class Coating(WakeField):
                  approx=False):
         """
         WakeField element for a coated circular beam pipe.
-        
+
         The longitudinal and tranverse impedances are computed using formulas
         from [1].
 
@@ -287,8 +301,8 @@ class Coating(WakeField):
 
         References
         ----------
-        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive 
-        wall impedance on beam dynamics in the Future Circular e+ e− Collider." 
+        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive
+        wall impedance on beam dynamics in the Future Circular e+ e− Collider."
         Physical Review Accelerators and Beams 21.4 (2018): 041001.
 
         """
@@ -319,10 +333,10 @@ class Coating(WakeField):
 
     def LongitudinalImpedance(self, f, approx):
         """
-        Compute the longitudinal impedance of a coating using Eq. (5), or 
-        approxmiated expression Eq. (8), of [1]. The approxmiated expression 
-        is valid if the skin depth of the coating is large compared to the 
-        coating thickness. 
+        Compute the longitudinal impedance of a coating using Eq. (5), or
+        approxmiated expression Eq. (8), of [1]. The approxmiated expression
+        is valid if the skin depth of the coating is large compared to the
+        coating thickness.
 
         Parameters
         ----------
@@ -335,11 +349,11 @@ class Coating(WakeField):
         -------
         Zl : array
             Longitudinal impedance values in [ohm].
-            
+
         References
         ----------
-        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive 
-        wall impedance on beam dynamics in the Future Circular e+ e− Collider." 
+        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive
+        wall impedance on beam dynamics in the Future Circular e+ e− Collider."
         Physical Review Accelerators and Beams 21.4 (2018): 041001.
 
         """
@@ -368,10 +382,10 @@ class Coating(WakeField):
 
     def TransverseImpedance(self, f, approx):
         """
-        Compute the transverse impedance of a coating using Eq. (6), or 
-        approxmiated expression Eq. (9), of [1]. The approxmiated expression 
-        is valid if the skin depth of the coating is large compared to the 
-        coating thickness. 
+        Compute the transverse impedance of a coating using Eq. (6), or
+        approxmiated expression Eq. (9), of [1]. The approxmiated expression
+        is valid if the skin depth of the coating is large compared to the
+        coating thickness.
 
         Parameters
         ----------
@@ -384,11 +398,11 @@ class Coating(WakeField):
         -------
         Zt : array
             Transverse impedance values in [ohm].
-            
+
         References
         ----------
-        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive 
-        wall impedance on beam dynamics in the Future Circular e+ e− Collider." 
+        [1] : Migliorati, M., E. Belli, and M. Zobov. "Impact of the resistive
+        wall impedance on beam dynamics in the Future Circular e+ e− Collider."
         Physical Review Accelerators and Beams 21.4 (2018): 041001.
 
         """