Updated CircularResistiveWall class to replace numerical integrations with an...

Updated CircularResistiveWall class to replace numerical integrations with an analytical forms for improved performance.

mbtrack2/impedance/resistive_wall.py

@@ -5,7 +5,7 @@ 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 scipy.special import wofz
 
 from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction
 
@@ -45,7 +45,7 @@ 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].
         
     Parameters
     ----------
@@ -71,7 +71,11 @@ class CircularResistiveWall(WakeField):
     ----------
     [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 
+    [2] : Ivanyan, Mikayel I. and Tsakanov, Vasili M. "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.
@@ -126,9 +130,11 @@ class CircularResistiveWall(WakeField):
     def LongitudinalWakeFunction(self, time, exact=False, atol=1e-20):
         """
         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.
+        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 equivalent to Eq. (22) in [2].
         
         If some time value is smaller than atol, then the fundamental theorem 
         of beam loading is applied: Wl(0) = Wl(0+)/2.
@@ -151,18 +157,25 @@ class CircularResistiveWall(WakeField):
             
         References
         ----------
-        [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
+        [1] : Ivanyan, Mikayel I. and Tsakanov, Vasili M. "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.
         """
         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)
+            idx4 = np.isclose(0, time, atol=atol)
+            factor = 4 * self.Z0 * c / (np.pi * self.radius**2) * self.length
+            wl[idx2] = 0.25 * factor
+            wl[idx3] = self.__LongWakeExact(time[idx3], factor)
+            wl[idx4] *= 0.5
         else:
             idx2 = np.logical_not(idx1)
             wl[idx2] = self.__LongWakeApprox(time[idx2])
@@ -171,11 +184,13 @@ class CircularResistiveWall(WakeField):
     def TransverseWakeFunction(self, time, exact=False):
         """
         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.
+        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.
         
-        Exact expression (Eq. (25) from [1]) is corrected by factor (c * t0).
+        Eq. (11) in [1] is completely equivalent to Eq. (25) in [2].
+        Corrected the typos in the last two terms in exact expression
+        (Eq. (11), of [1]).
 
         Parameters
         ----------
@@ -191,51 +206,50 @@ class CircularResistiveWall(WakeField):
             
         References
         ----------
-        [1] : Skripka, Galina, et al. "Simultaneous computation of intrabunch and 
+        [1] : Ivanyan, Mikayel I. and Tsakanov, Vasili M. "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 = ((8 * self.Z0 * c**2 * self.t0) / (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])
         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):
+        wl = np.real(factor / 12 * ( 4 * np.exp(-1 * t / self.t0)
+            * np.cos(np.sqrt(3) * t / self.t0)
+            + wofz(1j * np.sqrt(2 * t / self.t0))
+            - wofz(np.exp(1j * np.pi / 6) * np.sqrt(2 * t / self.t0))
+            - wofz(-1 * np.exp(-1j * np.pi / 6) * np.sqrt(2 * t / self.t0)) ) )
         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):
+        wt = np.real(factor / 24 * ( 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) )
+            + wofz(1j * np.sqrt(2 * t / self.t0))
+            + np.exp(-1j * np.pi / 3) * wofz( np.exp(1j * np.pi / 6)
+            * np.sqrt(2 * t / self.t0) )
+            + np.exp(1j * np.pi / 3) * wofz( -1 * np.exp(-1j * np.pi / 6)
+            * np.sqrt(2 * t / self.t0) ) ) )
         return wt
 
     def __LongWakeApprox(self, t):
-        wl = -1 * (1 / (4 * np.pi * self.radius) * np.sqrt(self.Z0 * self.rho /
-                                                           (c * np.pi)) /
+        wl = -1 * (1 / (4 * np.pi * self.radius) *
+            np.sqrt(self.Z0 * self.rho / (c * np.pi)) /
             t**(3 / 2)) * self.length
         return wl
 
@@ -245,12 +259,6 @@ class CircularResistiveWall(WakeField):
             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):