From 741567e34a0ce53a6065b84824095f7a76f8802e Mon Sep 17 00:00:00 2001
From: Keon Hee KIM <alphaover2pi@korea.ac.kr>
Date: Thu, 25 Jul 2024 18:47:50 +0200
Subject: [PATCH] Faster CircularResistiveWall

---
 mbtrack2/impedance/resistive_wall.py          | 244 +++++++++---------
 .../particles_electromagnetic_fields.py       |   2 +-
 mbtrack2/tracking/wakepotential.py            |   9 +-
 mbtrack2/utilities/misc.py                    |   4 +-
 4 files changed, 136 insertions(+), 123 deletions(-)

diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py
index dd99bf1..91fa7e1 100644
--- a/mbtrack2/impedance/resistive_wall.py
+++ b/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.
 
         """
diff --git a/mbtrack2/tracking/particles_electromagnetic_fields.py b/mbtrack2/tracking/particles_electromagnetic_fields.py
index ac9577b..3780ebb 100644
--- a/mbtrack2/tracking/particles_electromagnetic_fields.py
+++ b/mbtrack2/tracking/particles_electromagnetic_fields.py
@@ -4,7 +4,7 @@ For example, it can be applied to space charge, beam-beam force, electron lenses
 This is largely adapted from a fork of PyHEADTAIL https://github.com/gubaidulinvadim/PyHEADTAIL.
 Only the fastest Fadeeva implementation of the error function is used here.
 See  Oeftiger, A., de Maria, R., Deniau, L., Li, K., McIntosh, E., Moneta, L., Hegglin, S., Aviral, A. (2016).
-Review of CPU and GPU Fadeeva Implementations. https://cds.cern.ch/record/2207430/files/wepoy044.pdf 
+Review of CPU and GPU Faddeeva Implementations. https://cds.cern.ch/record/2207430/files/wepoy044.pdf 
 """
 
 from functools import wraps
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 3003b1f..3fdb19b 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/mbtrack2/tracking/wakepotential.py
@@ -709,9 +709,8 @@ class LongRangeResistiveWall(Element):
             Wake function in [V/C].
 
         """
-        wl = (1 / (4 * pi * self.radius) * np.sqrt(self.Z0 * self.rho /
-                                                   (c*pi)) /
-              t**(3 / 2)) * self.length * -1
+        wl = (1 / (4 * pi * self.radius) *
+              np.sqrt(self.Z0 * self.rho / (c * pi * t**3)) * self.length) * -1
         return wl
 
     def Wdip(self, t, plane):
@@ -739,8 +738,8 @@ class LongRangeResistiveWall(Element):
         else:
             raise ValueError()
 
-        wdip = (1 / (pi * r3**3) * np.sqrt(self.Z0 * c * self.rho / pi) /
-                t**(1 / 2) * self.length)
+        wdip = (1 / (pi * r3**3) * np.sqrt(self.Z0 * c * self.rho / (pi*t)) *
+                self.length)
         return wdip
 
     def update_tables(self, beam):
diff --git a/mbtrack2/utilities/misc.py b/mbtrack2/utilities/misc.py
index f884b17..8f35315 100644
--- a/mbtrack2/utilities/misc.py
+++ b/mbtrack2/utilities/misc.py
@@ -289,7 +289,7 @@ def yokoya_elliptic(x_radius, y_radius):
             # less than or equal to 0.8.
             qr_th = 0.8
 
-            if qr < qr_th:
+            if qr <= qr_th:
                 F = np.sqrt(large_semiaxis**2 - small_semiaxis**2)
                 mu_b = np.arccosh(large_semiaxis / F)
 
@@ -307,7 +307,7 @@ def yokoya_elliptic(x_radius, y_radius):
                     yokxdip, yokydip = yokydip, yokxdip
                     yokxquad, yokyquad = yokyquad, yokxquad
 
-            # Beyond the threshold (qr >= qr_th), they may not be valid,
+            # Beyond the threshold (qr > qr_th), they may not be valid,
             # but should converge to Yokoya form factors of parallel plates.
             # Fortunately, beyond this threshold, they should show asymptotic behavior,
             # so we perform linear interpolation.
-- 
GitLab