From b83abb5f2fe4a874729453a0d80fb44de15da105 Mon Sep 17 00:00:00 2001
From: KeonHKim <alphaover2pi@korea.ac.kr>
Date: Fri, 19 Jul 2024 00:08:11 +0900
Subject: [PATCH] Reconsidering the fundamental theorem of beam loading

---
 mbtrack2/impedance/resistive_wall.py |  18 +-
 mbtrack2/tracking/wakepotential.py   |   8 +-
 mbtrack2/utilities/misc.py           | 262 ++++++++++++++-------------
 3 files changed, 149 insertions(+), 139 deletions(-)

diff --git a/mbtrack2/impedance/resistive_wall.py b/mbtrack2/impedance/resistive_wall.py
index 5c7d0db..65bdae6 100644
--- a/mbtrack2/impedance/resistive_wall.py
+++ b/mbtrack2/impedance/resistive_wall.py
@@ -171,11 +171,17 @@ class CircularResistiveWall(WakeField):
         if exact == True:
             idx2 = time == 0
             idx3 = np.logical_not(np.logical_or(idx1, idx2))
-            idx4 = np.isclose(0, time, atol=atol)
+            # idx4 = np.isclose(time, 0, atol=atol) & (time >= 0)
             factor = self.Z0 * c / (3 * np.pi * self.radius**2) * self.length
-            wl[idx2] = 3 * factor
+            if np.any(idx2):
+                # fundamental theorem of beam loading
+                wl[idx2] = 3 * factor / 2
             wl[idx3] = self.__LongWakeExact(time[idx3], factor)
-            wl[idx4] *= 0.5
+            # if np.any(idx4):
+            #     closest_to_zero_idx = np.argmin(time[idx4])
+            #     # Get the actual index in the original array
+            #     closest_to_zero_idx = np.where(idx4)[0][closest_to_zero_idx]
+            #     wl[closest_to_zero_idx] *= 0.5
         else:
             idx2 = np.logical_not(idx1)
             wl[idx2] = self.__LongWakeApprox(time[idx2])
@@ -249,13 +255,13 @@ class CircularResistiveWall(WakeField):
 
     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
+            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) *
+            np.sqrt(self.Z0 * c * self.rho / (np.pi * t)) *
             self.length)
         return wt
 
diff --git a/mbtrack2/tracking/wakepotential.py b/mbtrack2/tracking/wakepotential.py
index 3003b1f..691d961 100644
--- a/mbtrack2/tracking/wakepotential.py
+++ b/mbtrack2/tracking/wakepotential.py
@@ -710,8 +710,8 @@ class LongRangeResistiveWall(Element):
 
         """
         wl = (1 / (4 * pi * self.radius) * np.sqrt(self.Z0 * self.rho /
-                                                   (c*pi)) /
-              t**(3 / 2)) * self.length * -1
+                                                   (c*pi*t**3))
+              * self.length) * -1
         return wl
 
     def Wdip(self, t, plane):
@@ -739,8 +739,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 8e688d0..8f35315 100644
--- a/mbtrack2/utilities/misc.py
+++ b/mbtrack2/utilities/misc.py
@@ -8,7 +8,7 @@ from pathlib import Path
 import numpy as np
 import pandas as pd
 from scipy.interpolate import interp1d
-from scipy.special import gamma, factorial, hyp2f1
+from scipy.special import factorial, gamma, hyp2f1
 
 from mbtrack2.impedance.wakefield import Impedance
 from mbtrack2.utilities.spectrum import spectral_density
@@ -192,27 +192,30 @@ def yokoya_elliptic(x_radius, y_radius):
     
     References
     ----------
-    [1] : Phys. Rev. Accel. Beams 22, 121001 (2019)
-    [2] : Phys. Rev. E 47, 656 (1993)
+    [1] : M. Migliorati, L. Palumbo, C. Zannini, N. Biancacci, and
+    V. G. Vaccaro, "Resistive wall impedance in elliptical multilayer vacuum
+    chambers." Phys. Rev. Accel. Beams 22, 121001 (2019).
+    [2] : R.L. Gluckstern, J. van Zeijts, and B. Zotter, "Coupling impedance
+    of beam pipes of general cross section." Phys. Rev. E 47, 656 (1993).
 
     """
-    
-    if x_radius==0 or y_radius==0:
-        raise ValueError("Both radii must be non-zero values.")
-    elif x_radius==np.inf and y_radius==np.inf:
+
+    if (x_radius <= 0) or (y_radius <= 0):
+        raise ValueError("Both radii must be non-zero positive values.")
+    elif (x_radius == np.inf) and (y_radius == np.inf):
         raise ValueError("Both radii have infinite values.")
-    elif x_radius==np.inf:
+    elif x_radius == np.inf:
         yoklong = 1.0
-        yokxdip = np.pi**2/24
-        yokydip = np.pi**2/12
-        yokxquad = -np.pi**2/24
-        yokyquad = np.pi**2/24
-    elif y_radius==np.inf:
+        yokxdip = np.pi**2 / 24
+        yokydip = np.pi**2 / 12
+        yokxquad = -np.pi**2 / 24
+        yokyquad = np.pi**2 / 24
+    elif y_radius == np.inf:
         yoklong = 1.0
-        yokxdip = np.pi**2/12
-        yokydip = np.pi**2/24
-        yokxquad = np.pi**2/24
-        yokyquad = -np.pi**2/24
+        yokxdip = np.pi**2 / 12
+        yokydip = np.pi**2 / 24
+        yokxquad = np.pi**2 / 24
+        yokyquad = -np.pi**2 / 24
     else:
         if y_radius < x_radius:
             small_semiaxis = y_radius
@@ -221,53 +224,7 @@ def yokoya_elliptic(x_radius, y_radius):
             small_semiaxis = x_radius
             large_semiaxis = y_radius
 
-        # Our equations are approximately valid only for qr (ratio) values
-        # less than or equal to 0.8.
-        qr_th = 0.8
-        F = np.sqrt(large_semiaxis**2 - small_semiaxis**2)
-        mu_b = np.arccosh(large_semiaxis/F)
-        qr = (large_semiaxis - small_semiaxis) / (large_semiaxis + small_semiaxis)
-
-        def function_ff(small_semiaxis, F, mu_b, ip, il):
-            fflong = ( 2*np.sqrt(2)*small_semiaxis/(np.pi*F)
-                * (-1)**ip/np.cosh(2*ip*mu_b)
-                * (-1)**il/np.cosh(2*il*mu_b) )
-            
-            ffdx = ( np.sqrt(2)*small_semiaxis**3/(np.pi*F**3)
-                * (-1)**ip*(2*ip+1)/np.cosh((2*ip+1)*mu_b)
-                * (-1)**il*(2*il+1)/np.cosh((2*il+1)*mu_b) )
-            
-            ffdy = ( np.sqrt(2)*small_semiaxis**3/(np.pi*F**3)
-                * (-1)**ip*(2*ip+1)/np.sinh((2*ip+1)*mu_b)
-                * (-1)**il*(2*il+1)/np.sinh((2*il+1)*mu_b) )
-            
-            ffquad = ( np.sqrt(2)*small_semiaxis**3/(np.pi*F**3)
-                * (-1)**ip*(2*ip)**2/np.cosh(2*ip*mu_b)
-                * (-1)**il/np.cosh(2*il*mu_b) )
-            return (fflong, ffdx, ffdy, ffquad)
-
-        def function_L(mu_b, ip, il):
-            Lm = ( np.sqrt(2)*np.pi*np.exp(-(2*abs(ip-il)+1)*mu_b)
-                * gamma(0.5+abs(ip-il))/(gamma(0.5)*factorial(abs(ip-il)))
-                * hyp2f1(0.5, abs(ip-il)+0.5, abs(ip-il)+1, np.exp(-4*mu_b))
-                + np.sqrt(2)*np.pi*np.exp(-(2*ip+2*il+1)*mu_b)
-                * gamma(0.5+ip+il)/(gamma(0.5)*factorial(ip+il))
-                * hyp2f1(0.5, ip+il+0.5, ip+il+1, np.exp(-4*mu_b)) )
-        
-            Ldx = ( np.sqrt(2)*np.pi*np.exp(-(2*abs(ip-il)+1)*mu_b)
-                * gamma(0.5+abs(ip-il))/(gamma(0.5)*factorial(abs(ip-il)))
-                * hyp2f1(0.5, abs(ip-il)+0.5, abs(ip-il)+1, np.exp(-4*mu_b))
-                + np.sqrt(2)*np.pi*np.exp(-(2*ip+2*il+3)*mu_b)
-                * gamma(1.5+ip+il)/(gamma(0.5)*factorial(ip+il+1))
-                * hyp2f1(0.5, ip+il+1.5, ip+il+2, np.exp(-4*mu_b)) )
-        
-            Ldy = ( np.sqrt(2)*np.pi*np.exp(-(2*abs(ip-il)+1)*mu_b)
-                * gamma(0.5+abs(ip-il))/(gamma(0.5)*factorial(abs(ip-il)))
-                * hyp2f1(0.5, abs(ip-il)+0.5, abs(ip-il)+1, np.exp(-4*mu_b))
-                - np.sqrt(2)*np.pi*np.exp(-(2*ip+2*il+3)*mu_b)
-                * gamma(1.5+ip+il)/(gamma(0.5)*factorial(ip+il+1))
-                * hyp2f1(0.5, ip+il+1.5, ip+il+2, np.exp(-4*mu_b)) )
-            return (Lm, Ldx, Ldy)
+        qr = (large_semiaxis-small_semiaxis) / (large_semiaxis+small_semiaxis)
 
         if qr == 0:
             yoklong = 1.0
@@ -276,79 +233,126 @@ def yokoya_elliptic(x_radius, y_radius):
             yokxquad = 0.0
             yokyquad = 0.0
 
-        # Beyond this threshold (qr_th), they are not valid (they can even diverge),
-        # but should converge to Yokoya form factors of parallel plates.
-        # Fortunately, beyond this threshold, they should show asymptotic behavior,
-        # so we perform linear interpolation.
-        elif qr > qr_th:
-            yoklong_pp = 1.0
-            yokxdip_pp = np.pi**2/24
-            yokydip_pp = np.pi**2/12
-            yokxquad_pp = -np.pi**2/24
-            yokyquad_pp = np.pi**2/24
-
-            small_semiaxis_th = large_semiaxis*(1-qr_th)/(1+qr_th)
-            F_th = np.sqrt(large_semiaxis**2 - small_semiaxis_th**2)
-            mu_b_th = np.arccosh(large_semiaxis/F_th)
-
-            ip_range = np.arange(51)
-            il_range = np.arange(51)
-            ip, il = np.meshgrid(ip_range, il_range, indexing="ij")
-
-            ff_values = np.array(function_ff(small_semiaxis_th, F_th, mu_b_th, ip, il))
-            L_values = np.array(function_L(mu_b_th, ip, il))
-
-            coeff_long = np.where( (ip == 0) & (il == 0), 0.25,
-                                    np.where((ip == 0) | (il == 0), 0.5, 1.0) )
-            coeff_quad = np.where(il == 0, 0.5, 1.0)
-
-            yoklong_th = np.sum(coeff_long * ff_values[0] * L_values[0])
-            yokxdip_th = np.sum(ff_values[1] * L_values[1])
-            yokydip_th = np.sum(ff_values[2] * L_values[2])
-            yokxquad_th = -np.sum(coeff_quad * ff_values[3] * L_values[0])
-            yokyquad_th = -yokxquad_th
-            
-            if y_radius > x_radius:
-                yokxdip_th, yokydip_th = yokydip_th, yokxdip_th
-                yokxquad_th, yokyquad_th = yokyquad_th, yokxquad_th
-                yokxdip_pp, yokydip_pp = yokydip_pp, yokxdip_pp
-                yokxquad_pp, yokyquad_pp = yokyquad_pp, yokxquad_pp
-
-            qr_array = np.array([qr_th, 1.0])
-            yoklong_array = np.array([yoklong_th, yoklong_pp])
-            yokxdip_array = np.array([yokxdip_th, yokxdip_pp])
-            yokydip_array = np.array([yokydip_th, yokydip_pp])
-            yokxquad_array = np.array([yokxquad_th, yokxquad_pp])
-            yokyquad_array = np.array([yokyquad_th, yokyquad_pp])
-
-            yoklong = np.interp(qr, qr_array, yoklong_array)
-            yokxdip = np.interp(qr, qr_array, yokxdip_array)
-            yokydip = np.interp(qr, qr_array, yokydip_array)
-            yokxquad = np.interp(qr, qr_array, yokxquad_array)
-            yokyquad = np.interp(qr, qr_array, yokyquad_array)
-
         else:
+            # Define form factor functions
+            def function_ff(small_semiaxis, F, mu_b, ip, il):
+                coeff_fflong = 2 * np.sqrt(2) * small_semiaxis / (np.pi * F)
+                coeff_fftrans = np.sqrt(2) * small_semiaxis**3 / (np.pi * F**3)
+
+                fflong = (coeff_fflong * (-1)**ip / np.cosh(2 * ip * mu_b) *
+                          (-1)**il / np.cosh(2 * il * mu_b))
+                ffdx = (coeff_fftrans * (-1)**ip * (2*ip + 1) / np.cosh(
+                    (2*ip + 1) * mu_b) * (-1)**il * (2*il + 1) / np.cosh(
+                        (2*il + 1) * mu_b))
+                ffdy = (coeff_fftrans * (-1)**ip * (2*ip + 1) / np.sinh(
+                    (2*ip + 1) * mu_b) * (-1)**il * (2*il + 1) / np.sinh(
+                        (2*il + 1) * mu_b))
+                ffquad = (coeff_fftrans * (-1)**ip * (2 * ip)**2 /
+                          np.cosh(2 * ip * mu_b) * (-1)**il /
+                          np.cosh(2 * il * mu_b))
+
+                return (fflong, ffdx, ffdy, ffquad)
+
+            def function_L(mu_b, ip, il):
+                common_L = (np.sqrt(2) * np.pi *
+                            np.exp(-(2 * abs(ip - il) + 1) * mu_b) *
+                            gamma(0.5 + abs(ip - il)) /
+                            (gamma(0.5) * factorial(abs(ip - il))) *
+                            hyp2f1(0.5,
+                                   abs(ip - il) + 0.5,
+                                   abs(ip - il) + 1, np.exp(-4 * mu_b)))
+                val_m = (
+                    np.sqrt(2) * np.pi * np.exp(-(2*ip + 2*il + 1) * mu_b) *
+                    gamma(0.5 + ip + il) / (gamma(0.5) * factorial(ip + il)) *
+                    hyp2f1(0.5, ip + il + 0.5, ip + il + 1, np.exp(-4 * mu_b)))
+                val_d = (
+                    np.sqrt(2) * np.pi * np.exp(-(2*ip + 2*il + 3) * mu_b) *
+                    gamma(1.5 + ip + il) /
+                    (gamma(0.5) * factorial(ip + il + 1)) *
+                    hyp2f1(0.5, ip + il + 1.5, ip + il + 2, np.exp(-4 * mu_b)))
+
+                Lm = common_L + val_m
+                Ldx = common_L + val_d
+                Ldy = common_L - val_d
+
+                return (Lm, Ldx, Ldy)
+
             ip_range = np.arange(51)
             il_range = np.arange(51)
             ip, il = np.meshgrid(ip_range, il_range, indexing="ij")
 
-            ff_values = np.array(function_ff(small_semiaxis, F, mu_b, ip, il))
-            L_values = np.array(function_L(mu_b, ip, il))
-
-            coeff_long = np.where( (ip == 0) & (il == 0), 0.25,
-                                    np.where((ip == 0) | (il == 0), 0.5, 1.0) )
+            coeff_long = np.where((ip == 0) & (il == 0), 0.25,
+                                  np.where((ip == 0) | (il == 0), 0.5, 1.0))
             coeff_quad = np.where(il == 0, 0.5, 1.0)
 
-            yoklong = np.sum(coeff_long * ff_values[0] * L_values[0])
-            yokxdip = np.sum(ff_values[1] * L_values[1])
-            yokydip = np.sum(ff_values[2] * L_values[2])
-            yokxquad = -np.sum(coeff_quad * ff_values[3] * L_values[0])
-            yokyquad = -yokxquad
+            # Our equations are approximately valid only for qr (ratio) values
+            # less than or equal to 0.8.
+            qr_th = 0.8
+
+            if qr <= qr_th:
+                F = np.sqrt(large_semiaxis**2 - small_semiaxis**2)
+                mu_b = np.arccosh(large_semiaxis / F)
+
+                ff_values = np.array(
+                    function_ff(small_semiaxis, F, mu_b, ip, il))
+                L_values = np.array(function_L(mu_b, ip, il))
+
+                yoklong = np.sum(coeff_long * ff_values[0] * L_values[0])
+                yokxdip = np.sum(ff_values[1] * L_values[1])
+                yokydip = np.sum(ff_values[2] * L_values[2])
+                yokxquad = -np.sum(coeff_quad * ff_values[3] * L_values[0])
+                yokyquad = -yokxquad
+
+                if y_radius > x_radius:
+                    yokxdip, yokydip = yokydip, yokxdip
+                    yokxquad, yokyquad = yokyquad, yokxquad
+
+            # 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.
+            else:
+                yoklong_pp = 1.0
+                yokxdip_pp = np.pi**2 / 24
+                yokydip_pp = np.pi**2 / 12
+                yokxquad_pp = -np.pi**2 / 24
+                yokyquad_pp = np.pi**2 / 24
+
+                small_semiaxis_th = large_semiaxis * (1-qr_th) / (1+qr_th)
+                F_th = np.sqrt(large_semiaxis**2 - small_semiaxis_th**2)
+                mu_b_th = np.arccosh(large_semiaxis / F_th)
+
+                ff_values_th = np.array(
+                    function_ff(small_semiaxis_th, F_th, mu_b_th, ip, il))
+                L_values_th = np.array(function_L(mu_b_th, ip, il))
+
+                yoklong_th = np.sum(coeff_long * ff_values_th[0] *
+                                    L_values_th[0])
+                yokxdip_th = np.sum(ff_values_th[1] * L_values_th[1])
+                yokydip_th = np.sum(ff_values_th[2] * L_values_th[2])
+                yokxquad_th = -np.sum(
+                    coeff_quad * ff_values_th[3] * L_values_th[0])
+                yokyquad_th = -yokxquad_th
+
+                if y_radius > x_radius:
+                    yokxdip_th, yokydip_th = yokydip_th, yokxdip_th
+                    yokxquad_th, yokyquad_th = yokyquad_th, yokxquad_th
+                    yokxdip_pp, yokydip_pp = yokydip_pp, yokxdip_pp
+                    yokxquad_pp, yokyquad_pp = yokyquad_pp, yokxquad_pp
+
+                qr_array = np.array([qr_th, 1.0])
+                yoklong_array = np.array([yoklong_th, yoklong_pp])
+                yokxdip_array = np.array([yokxdip_th, yokxdip_pp])
+                yokydip_array = np.array([yokydip_th, yokydip_pp])
+                yokxquad_array = np.array([yokxquad_th, yokxquad_pp])
+                yokyquad_array = np.array([yokyquad_th, yokyquad_pp])
+
+                yoklong = np.interp(qr, qr_array, yoklong_array)
+                yokxdip = np.interp(qr, qr_array, yokxdip_array)
+                yokydip = np.interp(qr, qr_array, yokydip_array)
+                yokxquad = np.interp(qr, qr_array, yokxquad_array)
+                yokyquad = np.interp(qr, qr_array, yokyquad_array)
 
-            if y_radius > x_radius:
-                yokxdip, yokydip = yokydip, yokxdip
-                yokxquad, yokyquad = yokyquad, yokxquad
-    
     return (yoklong, yokxdip, yokydip, yokxquad, yokyquad)
 
 
-- 
GitLab