Define an independent wofz function to avoid unnecessary complex number calculations

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.special import wofz
+from scipy.special import wofz as _scipy_wofz
 
 from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction
 
@@ -238,25 +238,39 @@ class CircularResistiveWall(WakeField):
             idx2 = np.logical_not(idx1)
             wt[idx2] = self.__TransWakeApprox(time[idx2])
         return wt
+    
+    def __wofz(self, z):
+        """
+        Compute the Faddeeva function w(z) = exp(-z**2) * erfc(-i*z).
+
+        Returns
+        -------
+        tuple
+            Real and imaginary parts of the Faddeeva function.
+        """
+        res = _scipy_wofz(z)
+        return res.real, res.imag
 
     def __LongWakeExact(self, t, factor):
+        w1re, _ = self.__wofz( 1j * np.sqrt(2 * t / self.t0) )
+        w2re, _ = self.__wofz( np.exp(1j * 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)
-             + np.real( wofz( 1j *
-                        np.sqrt(2 * t / self.t0) ) )
-             - 2 * np.real( wofz( np.exp(1j * np.pi / 6) *
-                            np.sqrt(2 * t / self.t0) ) ) )
+             + w1re - 2 * w2re )
         return wl
 
     def __TransWakeExact(self, t, factor):
+        w1re, _ = self.__wofz( 1j * np.sqrt(2 * t / self.t0) )
+        w2re, w2im = self.__wofz( np.exp(1j * 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) )
-             + np.real( wofz( 1j *
-                        np.sqrt(2 * t / self.t0) ) )
-             + 2 * np.real( np.exp(-1j * np.pi / 3) *
-                            wofz( np.exp(1j * np.pi / 6) *
-                            np.sqrt(2 * 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):

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