Source code for mbtrack2.impedance.csr

# -*- coding: utf-8 -*-
"""
Define coherent synchrotron radiation (CSR) wakefields in various models.
"""

import mpmath as mp
import numpy as np
from numpy.typing import NDArray
from scipy.constants import c, epsilon_0, mu_0, pi
from scipy.special import gamma

from mbtrack2.impedance.wakefield import Impedance, WakeField, WakeFunction
from mbtrack2.utilities.misc import power_minus1
from mbtrack2.utilities.synchrotron import Synchrotron


@np.vectorize
def _find_Yk_roots(x: float, k: int) -> NDArray:
    """
    Find the roots of the characteristic Eq.(5.18) in [1].
    Then the roots are selected based on the sign of x.
    Roots with nonzero imaginary part are discarded.

    Parameters
    ----------
    x : float
    k : int
    
    Returns
    -------
    Yk : float
    
    References
    ---------- 
    [1] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern,
    Longitudinal wakefield for an electron moving on a circular orbit,
    Particle Accelerators 57, 9 (1997).
    """
    coefficients = [1, 0, 0, -6 * x / k**(3 / 2), -3]
    roots = np.roots(coefficients)
    real_roots = roots[roots.imag == 0]
    real_roots = np.abs(real_roots)
    selected_root = np.where(x > 0, np.max(real_roots), np.min(real_roots))
    return selected_root


[docs] class FreeSpaceCSR(WakeField): """ Free space steady-state coherent synchrotron radiation Wakefield element, based on [1-3]. Notes: - Impedance is computed using Eq. (A10) of [2]. - The WakeField model stores Zlong and Wcsr. - The Wcsr component (type 'csr') holds the antiderivative of the wake function rather than the wake function itself, because the direct wake diverges at tau=0 making numerical convolution inaccurate. - WakePotential class uses the antiderivative convolved with the derivative of the bunch profile instead. - The direct wake function is available via LongitudinalWakeFunction for analysis but is not added to the model. Parameters ---------- time : array of float Time points where the wake function antiderivative is evaluated in [s]. Must span the bunch extent for tracking. frequency : array of float Frequency points where the impedance will be evaluated in [Hz]. length : float Length of the impedance element to consider in [m]. radius : float Dipole radius of curvature in [m]. ring : Synchrotron Ring object, used to access the Lorentz factor gamma. References ---------- [1] : Faltens, A., & Laslett, L. J. (1973). Longitudinal coupling impedance of a stationary electron ring in a cylindrical geometry. part. Accel., 4, 151-157. [2] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron radiation using mesh." Physical Review Special Topics-Accelerators and Beams 7.5 (2004): 054403. [3] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern, Longitudinal wakefield for an electron moving on a circular orbit, Particle Accelerators 57, 9 (1997). """
[docs] def __init__(self, time: NDArray, frequency: NDArray, length: float, radius: float, ring: Synchrotron): super().__init__() self.length = length self.radius = radius self.Z0 = mu_0 * c self.gamma = ring.gamma Zl = self.LongitudinalImpedance(frequency) Wl_antiderivative = self.LongitudinalWakeAntiderivative(time) Zlong = Impedance(variable=frequency, function=Zl, component_type='long') Wlong_antiderivative = WakeFunction(variable=time, function=Wl_antiderivative, component_type='csr') super().append_to_model(Zlong) super().append_to_model(Wlong_antiderivative)
[docs] def LongitudinalImpedance(self, frequency: NDArray) -> NDArray: """ Compute the free space steady-state CSR impedance. Based on Eq. (A10) of [1]. This formula is valid only if omega << (3 * gamma^3 * c) / (2 * R). Parameters ---------- frequency : float array Frequency in [Hz]. Returns ------- Zl : complex array Longitudinal impedance in [ohm]. References ---------- [1] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron radiation using mesh." Physical Review Special Topics-Accelerators and Beams 7.5 (2004): 054403. """ Zl = (self.Z0 * self.length / (2 * np.pi) * gamma(2 / 3) * (-1j * 2 * np.pi * frequency / (3 * c * self.radius**2))**(1 / 3)) return Zl
[docs] def LongitudinalWakeFunction(self, time: NDArray) -> NDArray: """ Compute the free-space CSR wake function following Eq. (13) of [1]. For analysis only. Not added to the WakeField model because the wake diverges at tau=0. Use the Wcsr antiderivative component for tracking (see LongitudinalWakeAntiderivative). Parameters ---------- time : float array Time in [s]. Returns ------- Wl : float array Longitudinal wake function in [V/C]. References ---------- [1] : E. L. Saldin, E. A. Schneidmiller, and M. V. Yurkov, On the coherent radiation of an electron bunch moving in an arc of a circle, Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 398, 373 (1997). [2] : J. Qiang, C. Mitchell, and R. D. Ryne, A fast high-order method to calculate wakefields in an electron beam, Nuclear Instruments and Methods in Physics Research Section A: Accelerators, Spectrometers, Detectors and Associated Equipment 682, 49 (2012). """ s_bar = time * c * self.gamma**3 / self.radius sq = np.sqrt(64 + 144 * s_bar**2) u = np.cbrt(sq + 12*s_bar) - np.cbrt(sq - 12*s_bar) A = (u**2 / 4 - 1) / 2 / (1 + u**2 / 4)**3 B = (1/6 - u**2 / 18 - u**4 / 96) C = (1 + u**2 / 4)**3 * (1 + u**2 / 12)**2 # Amplitude coefficient. Attention if antiderivative method is used the # coefficient is different. It should be 1/3 of this value. ampl_coef = -self.length / (pi* epsilon_0) * self.gamma**4 / self.radius**2 res = ampl_coef * (A + B/C) return res
[docs] def LongitudinalWakeAntiderivative(self, time: NDArray, points_around_zero: int = 100_000 ) -> NDArray: """ Compute the antiderivative of the free space CSR wakefunction. We follow Eq. (3.13) of [1]. The antiderivative should be used to calculate the part of the convolution integral around 0. Parameters ---------- time : float array Time in [s]. Returns ------- Wl_int : float array Antiderivative of the longitudinal wake function per unit length in [V/C/m]. References ---------- [1] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern, Longitudinal wakefield for an electron moving on a circular orbit, Particle Accelerators 57, 9 (1997). """ def _antiderivative(time): mu = -3 * self.gamma**3 * c * time / (2 * self.radius) lambda_var = np.sqrt(1 + mu**2) omega = mu + lambda_var antiderivative = ( 9 / 16 * (-2 / mu + 1 / (mu*lambda_var) * (omega**(1 / 3) + omega**(-1 / 3)) + 2 / lambda_var * (omega**(2 / 3) - omega**(-2 / 3)))) return antiderivative antiderivative = _antiderivative(time) antiderivative[time >= 0] = 0 # The first negative sample sits next to the tau = 0 singular edge. negative_indices = np.flatnonzero(time < 0) if negative_indices.size > 0 and time.size > 1: dtau = np.abs(time[1] - time[0]) smallest_negative = negative_indices[-1] t_around_zero = np.linspace(-dtau / 2, 0, points_around_zero) antiderivative_tmp = _antiderivative(t_around_zero) antiderivative_tmp[-1] = 0 antiderivative[smallest_negative] = ( np.trapezoid(antiderivative_tmp, t_around_zero) / (dtau/2)) ampl_coef = -2 * self.length / ( 9*pi*epsilon_0) * self.gamma / self.radius / c return ampl_coef * antiderivative
[docs] class ParallelPlatesCSR(WakeField): """ Perfectly conducting parallel plates steady-state coherent synchrotron radiation Wakefield element, based on [1,2]. The WakeField model stores three components: - Zlong (full impedance, free space + shielding, Eq. A1 of [1]) - Wlong (shielding correction only, Murphy et al. [2] Eqs. 5.21-5.22) - Wcsr (free-space CSR antiderivative from an internally constructed FreeSpaceCSR instance). Together Wlong and Wcsr give the full time-domain wake, consistent with Zlong. WakePotential class handles both components automatically. Parameters ---------- time : array of float Time points spanning the bunch extent in [s]. The analytical wake components are stored internally on the corresponding convolution lag grid. frequency : array of float Frequency points where the impedance will be evaluated in [Hz]. length : float Length of the impedance element to consider in [m]. radius : float Dipole radius of curvature in [m]. distance : float Full-gap vertical distance between the parallel plates in [m]. ring : Synchrotron Ring object, used to access the Lorentz factor gamma. k_max : int, optional Number of terms in the image-charge sum for the wake function. The default is 100. Attributes ---------- threshold : float Shielding threshold in the parallel plates model in [Hz]. References ---------- [1] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron radiation using mesh." Physical Review Special Topics-Accelerators and Beams 7.5 (2004): 054403. [2] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern, Longitudinal wakefield for an electron moving on a circular orbit, Particle Accelerators 57, 9 (1997). """
[docs] def __init__(self, time: NDArray, frequency: NDArray, length: float, radius: float, distance: float, ring: Synchrotron, k_max: int = 100): super().__init__() self.length = length self.radius = radius self.distance = distance self.Z0 = mu_0 * c self.gamma = ring.gamma self.k_max = k_max Zl = self.LongitudinalImpedance(frequency) Wl = self.LongitudinalWakeFunction(time, k_max=self.k_max) Zlong = Impedance(variable=frequency, function=Zl, component_type='long') Wlong = WakeFunction(variable=time, function=Wl, component_type="long") super().append_to_model(Zlong) super().append_to_model(Wlong) # Add free-space CSR antiderivative (Wcsr) for a complete time-domain model. csr_fs = FreeSpaceCSR(time=time, frequency=frequency, length=length, radius=radius, ring=ring) super().append_to_model(csr_fs.Wcsr)
@property def threshold(self) -> float: """Shielding threshold in the parallel plates model in [Hz].""" return (3*c) / (2 * np.pi) * (self.radius / self.distance**3)**0.5
[docs] def LongitudinalImpedance(self, frequency: NDArray, tol: float = 1e-5) -> NDArray: """ Compute the CSR impedance using the perfectly conducting parallel plates steady-state model. The impedance includes both the free space contribution and the shielding effect from the plates. Impedance is computed using Eq. (A1) of [1]. Parameters ---------- frequency : float array Frequency in [Hz]. tol : float, optinal Desired maximum final error on sum_func. Returns ------- Zl : complex array Longitudinal impedance in [ohm]. References ---------- [1] : Agoh, T., and K. Yokoya. "Calculation of coherent synchrotron radiation using mesh." Physical Review Special Topics-Accelerators and Beams 7.5 (2004): 054403. """ Zl = np.zeros(frequency.shape, dtype=np.complex128) constant = (2 * np.pi * self.Z0 * self.length / self.distance * (2 / self.radius)**(1 / 3)) for i, f in enumerate(frequency): k = 2 * mp.pi * f / c sum_value = mp.nsum(lambda p: self.sum_func(p, k), [0, mp.inf], tol=tol, method='r+s') Zl[i] = constant * (1 / k)**(1 / 3) * complex(sum_value) return Zl
[docs] def sum_func(self, p: int, k: float) -> complex: """ Utility function for LongitudinalImpedance. Parameters ---------- p : int k : float Returns ------- sum_value : mpc """ xp = (2*p + 1) * mp.pi / self.distance * (self.radius / 2 / k**2)**(1 / 3) Ai = mp.airyai(xp**2) Bi = mp.airybi(xp**2) Aip = mp.airyai(xp**2, 1) Bip = mp.airybi(xp**2, 1) return Aip * (Aip + 1j*Bip) + xp**2 * Ai * (Ai + 1j*Bi)
[docs] def _G2(self, s: NDArray, k_max: int = 100) -> NDArray: """ Normalised Green's function for the parallel plates CSR wakefunction. Implements Eq. (5.22b) of [1]. Parameters ---------- s : float array Longitudinal distance in [m]. Returns ------- G2 : float array Green's function. References ---------- [1] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern, Longitudinal wakefield for an electron moving on a circular orbit, Particle Accelerators 57, 9 (1997). """ delta = self.distance / self.radius / 2 x = -s / (2 * self.radius * delta**(3 / 2)) k = np.arange(1, k_max, 1) X, K = np.meshgrid(x, k) Y = _find_Yk_roots(x=X, k=K) return (2 * power_minus1(K + 1) / K**2 * (4 * Y**4 * (3 - Y**4) / (1 + Y**4)**3)).sum(axis=0)
[docs] def LongitudinalWakeFunction(self, time: NDArray, k_max: int = 100) -> NDArray: """ Compute the parallel-plates shielding correction to the CSR wake function following Eqs. (5.21-5.22) of [1]. Computes only the shielding correction; the free-space term is not included. In the WakeField model the free-space contribution is provided automatically by the Wcsr component. Parameters ---------- time : float array Time in [s]. k_max : int, optional Number of terms in the G2 sum. The default is 100. Returns ------- Wl : float array Longitudinal wake function (shielding correction only) in [V/C]. References ---------- [1] : J. B. Murphy, S. Krinsky, and R. L. Gluckstern, Longitudinal wakefield for an electron moving on a circular orbit, Particle Accelerators 57, 9 (1997). """ ampl_coef = -4 * self.length / (8*pi*epsilon_0) / self.distance**2 G2 = self._G2(time * c, k_max=k_max) return ampl_coef * G2