From 2bf4ef518314b7956c55edd1a2a9bff422006ccf Mon Sep 17 00:00:00 2001
From: KeonHKim <alphaover2pi@korea.ac.kr>
Date: Mon, 22 Jul 2024 22:49:57 +0900
Subject: [PATCH] Define an independent wofz function to avoid unnecessary
complex number calculations
---
mbtrack2/impedance/resistive_wall.py | 36 +++++++++++++------
.../particles_electromagnetic_fields.py | 2 +-
2 files changed, 26 insertions(+), 12 deletions(-)
diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py
index 21a513d..7928c48 100644
--- a/mbtrack2/impedance/resistive_wall.py
+++ b/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):
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
--
GitLab