"""
A package dealing with particles electromagnetic fields.
For example, it can be applied to space charge, beam-beam force,
electron lenses or beam-ion instabilities.
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 Faddeeva Implementations.
https://cds.cern.ch/record/2207430/files/wepoy044.pdf
"""
from functools import wraps
from typing import Callable
import numpy as np
from numpy.typing import NDArray
from scipy.constants import epsilon_0, pi
from scipy.special import wofz as _scipy_wofz
[docs]
def _wofz(x: NDArray, y: NDArray) -> tuple[NDArray, NDArray]:
"""
Compute the Faddeeva function w(z) = exp(-z^2) * erfc(-i*z).
Parameters
----------
x : float
Real part of the argument.
y : float
Imaginary part of the argument.
Returns
-------
tuple
Real and imaginary parts of the Faddeeva function.
"""
res = _scipy_wofz(x + 1j*y)
return res.real, res.imag
[docs]
def _sqrt_sig(sig_x: float, sig_y: float) -> float:
"""
Compute the square root of the difference between
the squared transverse rms and vertical rms.
Parameters
----------
sig_x : float
Transverse rms of the distribution.
sig_y : float
Vertical rms of the distribution.
Returns
-------
float
Square root of the difference between
the squared transverse rms and vertical rms.
"""
return np.sqrt(2 * (sig_x*sig_x - sig_y*sig_y))
[docs]
def _efieldn_mit(x: NDArray, y: NDArray, sig_x: float,
sig_y: float) -> tuple[NDArray, NDArray]:
"""
Returns electromagnetic fields as E_x/Q, E_y/Q in (V/m/Coulomb).
Parameters
----------
x : np.ndarray
x coordinates in meters.
y : np.ndarray
y coordinates in meters.
sig_x : float
Transverse rms of the distribution in meters.
sig_y : float
Vertical rms of the distribution in meters.
Returns
-------
tuple
Normalized electromagnetic fields Ex/Q, Ey/Q
in the units of (V/m/Coulomb).
"""
sig_sqrt = _sqrt_sig(sig_x, sig_y)
w1re, w1im = _wofz(x / sig_sqrt, y / sig_sqrt)
ex = np.exp(-x * x / (2*sig_x*sig_x) - y * y / (2*sig_y*sig_y))
w2re, w2im = _wofz(x * sig_y / (sig_x*sig_sqrt),
y * sig_x / (sig_y*sig_sqrt))
denom = 2 * epsilon_0 * np.sqrt(pi) * sig_sqrt
return (w1im - ex*w2im) / denom, (w1re - ex*w2re) / denom
[docs]
def efieldn_gauss_round(x: NDArray,
y: NDArray,
sig_r: float,
sig_r2: None = None) -> tuple[NDArray, NDArray]:
"""
Computes the electromagnetic field of a round Gaussian distribution.
Parameters
----------
x : np.ndarray
x coordinates in meters.
y : np.ndarray
y coordinates in meters.
sig_r : float
Transverse rms of the distribution in meters.
Returns
-------
tuple
Normalized electromagnetic fields Ex/Q, Ey/Q
in the units of (V/m/Coulomb).
"""
r_squared = x*x + y*y
amplitude = -np.expm1(-r_squared /
(2*sig_r*sig_r)) / (2*pi*epsilon_0*r_squared)
return x * amplitude, y * amplitude
[docs]
def _efieldn_linearized(x: NDArray, y: NDArray, sig_x: float,
sig_y: float) -> tuple[NDArray, NDArray]:
"""
Computes linearized electromagnetic field.
Parameters
----------
x : np.ndarray
x coordinate in meters.
y : np.ndarray
y coordinate in meters.
sig_x : float
Vertical rms of the distribution in meters.
sig_y : float
Vertical rms of the distribution in meters.
Returns
-------
tuple
Normalized electromagnetic fields Ex/Q, Ey/Q
in the units of (V/m/Coulomb).
"""
a = np.sqrt(2) * sig_x
b = np.sqrt(2) * sig_y
amplitude = 1 / (pi * epsilon_0 * (a+b))
return x / a * amplitude, y / b * amplitude
[docs]
def add_sigma_check(efieldn: Callable) -> Callable:
"""
Wrapper for a normalized electromagnetic field function.
Adds the following actions before calculating the field:
1) Exchange x and y quantities if sig_x < sig_y.
2) Apply round beam field formula when sig_x is close to sig_y.
Parameters
----------
efieldn : function
Function to calculate normalized electromagnetic field.
Returns
-------
function
Wrapped function, including round beam and inverted sig_x/sig_y.
"""
sigmas_ratio_threshold = 1e-3
absolute_threshold = 1e-10
@wraps(efieldn)
def efieldn_checked(x, y, sig_x, sig_y, *args, **kwargs):
tol_kwargs = dict(rtol=sigmas_ratio_threshold, atol=absolute_threshold)
if np.allclose(sig_x, sig_y, **tol_kwargs):
if np.allclose(sig_y, 0, **tol_kwargs):
en_x = en_y = np.zeros(x.shape, dtype=np.float64)
else:
en_x, en_y = efieldn_gauss_round(x, y, sig_x, sig_y, *args,
**kwargs)
elif np.all(sig_x < sig_y):
en_y, en_x = efieldn(y, x, sig_y, sig_x, *args, **kwargs)
else:
en_x, en_y = efieldn(x, y, sig_x, sig_y, *args, **kwargs)
return en_x, en_y
return efieldn_checked
[docs]
def get_displaced_efield(efieldn: Callable, xr: NDArray, yr: NDArray,
sig_x: float, sig_y: float, mean_x: float,
mean_y: float) -> tuple[NDArray, NDArray]:
"""
Compute the charge-normalized electric field components of
a two-dimensional Gaussian charge distribution.
Parameters
----------
efieldn : function
Calculates electromagnetic field of a given distribution of charges.
xr : np.array
x coordinates in meters.
yr : np.array
y coordinates in meters.
sig_x : float
Horizontal rms size in meters.
sig_y : float
Vertical rms size in meters.
mean_x : float
Horizontal mean of the distribution in meters.
mean_y : float
Vertical mean of the distribution in meters.
Returns
-------
tuple
Charge-normalized electromagnetic fields with
a displaced center of the distribution.
"""
x = xr - mean_x
y = yr - mean_y
efieldn = add_sigma_check(efieldn)
en_x, en_y = efieldn(np.abs(x), np.abs(y), sig_x, sig_y)
en_x = np.abs(en_x) * np.sign(x)
en_y = np.abs(en_y) * np.sign(y)
return en_x, en_y