# -*- coding: utf-8 -*-
"""
Module where the BeamLoadingEquilibrium class is defined.
"""
import matplotlib.pyplot as plt
import numpy as np
from numpy.typing import NDArray
from scipy.constants import c
from scipy.fft import rfft, rfftfreq
from scipy.integrate import trapezoid
from scipy.optimize import root
from mbtrack2.tracking.rf import CavityResonator
from mbtrack2.utilities.spectrum import gaussian_bunch
from mbtrack2.utilities.synchrotron import Synchrotron
[docs]
class BeamLoadingEquilibrium():
"""Class used to compute beam equilibrium profile for a given storage ring
and a list of RF cavities of any harmonic. The class assumes an uniform
filling of the storage ring. Based on [1] which is an extension of [2].
Parameters
----------
ring : Synchrotron object
cavity_list : list of CavityResonator objects
I0 : beam current in [A].
auto_set_MC_theta : if True, allow class to change cavity phase for
CavityResonator objetcs with m = 1 (i.e. main cavities)
F : list of form factor amplitude
PHI : list of form factor phase
B1 : lower intergration boundary
B2 : upper intergration boundary
N : int or float, optinal
Number of points used for longitudinal grid.
Default is 1000.
References
----------
[1] Gamelin, A., & Yamamoto, N. (2021). Equilibrium bunch density
distribution with multiple active and passive RF cavities.
IPAC'21 (MOPAB069).
[2] Venturini, M. (2018). Passive higher-harmonic rf cavities with general
settings and multibunch instabilities in electron storage rings.
Physical Review Accelerators and Beams, 21(11), 114404.
"""
[docs]
def __init__(self,
ring: Synchrotron,
cavity_list: list[CavityResonator],
I0: float,
auto_set_MC_theta: bool = False,
F: list[float] | None = None,
PHI: list[float] | None = None,
B1: float = -0.2,
B2: float = 0.2,
N: int = 1000):
self.ring = ring
self.cavity_list = cavity_list
self.I0 = I0
self.n_cavity = len(cavity_list)
self.auto_set_MC_theta = auto_set_MC_theta
if F is None:
self.F = np.ones((self.n_cavity, ))
else:
self.F = F
if PHI is None:
self.PHI = np.zeros((self.n_cavity, ))
else:
self.PHI = PHI
self.B1 = B1
self.B2 = B2
self.mpi = False
self.N = N
self.z0 = np.linspace(self.B1, self.B2, self.N)
self.tau0 = self.z0 / c
self.rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0) / c
self.sampling = self.tau0[1] - self.tau0[0]
self.freq0 = rfftfreq(self.N, self.sampling)
# Define constants for scaled potential u(z)
self.u0 = self.ring.U0 / (self.ring.ac * self.ring.sigma_delta**2 *
self.ring.E0 * self.ring.L)
self.ug = np.zeros((self.n_cavity, ))
self.ub = np.zeros((self.n_cavity, ))
self.update_potentials()
[docs]
def update_rho(self):
uexp_vals = self.uexp(self.z0)
A = trapezoid(uexp_vals, x=self.z0)
self.rho0 = uexp_vals / A
[docs]
def update_potentials(self):
"""Update potentials with cavity and ring data."""
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
self.ug[i] = cavity.Vg / (self.ring.ac * self.ring.sigma_delta**2 *
self.ring.E0 * self.ring.L *
self.ring.k1 * cavity.m)
self.ub[i] = 2 * self.I0 * cavity.Rs / (
self.ring.ac * self.ring.sigma_delta**2 * self.ring.E0 *
self.ring.L * self.ring.k1 * cavity.m * (1 + cavity.beta))
[docs]
def energy_balance(self) -> float:
"""Return energy balance for the synchronous particle
(z = 0 ,delta = 0)."""
delta = self.ring.U0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
delta += cavity.Vb(
self.I0) * self.F[i] * np.cos(cavity.psi - self.PHI[i])
delta -= cavity.Vg * np.cos(cavity.theta_g)
return delta
[docs]
def center_of_mass(self) -> float:
"""Return center of mass position in [s]"""
CM = np.average(self.z0, weights=self.rho0)
return CM / c
[docs]
def u(self, z: NDArray) -> NDArray:
"""Scaled potential u(z)"""
pot = self.u0 * z
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
pot += -self.ug[i] * (
np.sin(self.ring.k1 * cavity.m * z + cavity.theta_g) -
np.sin(cavity.theta_g))
pot += self.ub[i] * self.F[i] * np.cos(cavity.psi) * (
np.sin(self.ring.k1 * cavity.m * z + cavity.psi - self.PHI[i])
- np.sin(cavity.psi - self.PHI[i]))
return pot
[docs]
def du_dz(self, z: NDArray) -> NDArray:
"""Partial derivative of the scaled potential u(z) by z"""
pot = self.u0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
pot += -self.ug[i] * self.ring.k1 * cavity.m * np.cos(
self.ring.k1 * cavity.m * z + cavity.theta_g)
pot += self.ub[i] * self.F[i] * self.ring.k1 * cavity.m * np.cos(
cavity.psi) * np.cos(self.ring.k1 * cavity.m * z + cavity.psi -
self.PHI[i])
return pot
[docs]
def uexp(self, z: NDArray) -> NDArray:
return np.exp(-1 * self.u(z))
[docs]
def integrate_func(self, f: NDArray, g: NDArray) -> NDArray:
"""Return Integral[f*g]/Integral[f] between B1 and B2"""
A = trapezoid(f * g, x=self.z0)
B = trapezoid(f, x=self.z0)
return A / B
[docs]
def to_solve(self, x: NDArray | list, CM: bool = True) -> NDArray:
"""System of non-linear equation to solve to find the form factors F
and PHI at equilibrum.
The system is composed of Eq. (B6) and (B7) of [1] for each cavity.
If auto_set_MC_theta == True, the system also find the main cavity
phase to impose energy balance or cancel center of mass offset.
If CM is True, the system imposes zero center of mass offset,
if False, the system imposes energy balance.
"""
# Update values of F, PHI and theta
if self.auto_set_MC_theta:
self.F = x[:-1:2]
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
if cavity.m == 1:
cavity.theta = x[-1]
cavity.set_generator(self.I0)
self.update_potentials()
else:
self.F = x[::2]
self.PHI = x[1::2]
self.update_rho()
# Compute system
if self.auto_set_MC_theta:
res = np.zeros((self.n_cavity * 2 + 1, ))
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
res[2 *
i] = self.F[i] * np.cos(self.PHI[i]) - self.integrate_func(
self.uexp(self.z0),
np.cos(self.ring.k1 * cavity.m * self.z0))
res[2*i +
1] = self.F[i] * np.sin(self.PHI[i]) - self.integrate_func(
self.uexp(self.z0),
np.sin(self.ring.k1 * cavity.m * self.z0))
# Factor 1e-8 or 1e12 for better convergence
if CM is True:
res[self.n_cavity * 2] = self.center_of_mass() * 1e12
else:
res[self.n_cavity * 2] = self.energy_balance() * 1e-8
else:
res = np.zeros((self.n_cavity * 2, ))
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
res[2 *
i] = self.F[i] * np.cos(self.PHI[i]) - self.integrate_func(
self.uexp(self.z0),
np.cos(self.ring.k1 * cavity.m * self.z0))
res[2*i +
1] = self.F[i] * np.sin(self.PHI[i]) - self.integrate_func(
self.uexp(self.z0),
np.sin(self.ring.k1 * cavity.m * self.z0))
return res
[docs]
def plot_rho(self):
"""Plot the bunch equilibrium profile"""
plt.plot(self.tau0 * 1e12, self.rho0)
plt.xlabel(r"$\tau$ [ps]")
plt.title("Equilibrium bunch profile")
[docs]
def voltage(self, z: NDArray) -> NDArray:
"""Return the RF system total voltage at position z"""
Vtot = 0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
Vtot += cavity.VRF(z, self.I0, self.F[i], self.PHI[i])
return Vtot
[docs]
def dV(self, z: NDArray) -> NDArray:
"""Return derivative of the RF system total voltage at position z"""
Vtot = 0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
Vtot += cavity.dVRF(z, self.I0, self.F[i], self.PHI[i])
return Vtot
[docs]
def ddV(self, z: NDArray) -> NDArray:
"""Return the second derivative of the RF system\
total voltage at position z"""
Vtot = 0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
Vtot += cavity.ddVRF(z, self.I0, self.F[i], self.PHI[i])
return Vtot
[docs]
def deltaVRF(self, z: NDArray) -> NDArray:
"""Return the generator voltage minus beam loading\
voltage of the total RF system at position z"""
Vtot = 0
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
Vtot += cavity.deltaVRF(z, self.I0, self.F[i], self.PHI[i])
return Vtot
[docs]
def plot_dV(self):
"""Plot the derivative of RF system total voltage"""
plt.plot(self.z0, self.dV(self.z0))
plt.xlabel("z [m]")
plt.ylabel("Total RF voltage (V)")
[docs]
def plot_voltage(self):
"""Plot the RF system total voltage between z1 and z2"""
plt.plot(self.z0, self.voltage(self.z0))
plt.xlabel("z [m]")
plt.ylabel("Total RF voltage (V)")
[docs]
def std_rho(self) -> float:
"""Return the rms bunch equilibrium size in [m]"""
average = np.average(self.z0, weights=self.rho0)
variance = np.average((self.z0 - average)**2, weights=self.rho0)
return np.sqrt(variance)
[docs]
def beam_equilibrium(self,
x0: NDArray | list | None = None,
tol: float = 1e-4,
method: str = 'hybr',
options: dict = None,
plot: bool = False,
CM: bool = True,
verbiose: bool = False):
"""Solve system of non-linear equation to find the form factors F
and PHI at equilibrum. Can be used to compute the equilibrium bunch
profile.
Parameters
----------
x0 : initial guess
tol : tolerance for termination of the algorithm
method : method used by scipy.optimize.root to solve the system
options : options given to scipy.optimize.root
plot : if True, plot the equilibrium bunch profile
CM : if True, the system imposes zero center of mass offset,
if False, the system imposes energy balance.
verbiose : if True, print out informations about convergence.
Returns
-------
sol : OptimizeResult object representing the solution
"""
if x0 is None:
x0 = [1, 0] * self.n_cavity
if self.auto_set_MC_theta:
x0 = x0 + [self.cavity_list[0].theta]
if verbiose:
if CM:
print("The initial center of mass offset is " +
str(self.center_of_mass() * 1e12) + " ps")
else:
print("The initial energy balance is " +
str(self.energy_balance()) + " eV")
sol = root(lambda x: self.to_solve(x, CM),
x0,
tol=tol,
method=method,
options=options)
# Update values of F, PHI and theta_g
if self.auto_set_MC_theta:
self.F = sol.x[:-1:2]
for i in range(self.n_cavity):
cavity = self.cavity_list[i]
if cavity.m == 1:
cavity.theta = sol.x[-1]
else:
self.F = sol.x[::2]
self.PHI = sol.x[1::2]
if verbiose:
if CM:
print("The final center of mass offset is " +
str(self.center_of_mass() * 1e12) + " ps")
else:
print("The final energy balance is " +
str(self.energy_balance()) + " eV")
if not sol.success:
print("The algorithm has converged: " + str(sol.success))
if plot:
self.plot_rho()
return sol
[docs]
def PTBL_threshold(self,
I0: float,
m: int | None = None,
MC_index: int = 0,
HC_index: int = 1) -> tuple[float, float, float]:
"""
Compute the periodic transient beam loading (PTBL) instability
threshold based on [1].
The beam_equilibrium method should be called before the PTBL_threshold
one.
Eq. (22) and (23) have been modified for cosine voltage convention
for the main cavity.
Parameters
----------
I0 : float
Beam total current in [A].
m : int, optional
Number of bunches in the beam.
The default is None. If None, uniform filling is assumed.
MC_index : int, optional
Index of the main cavity in cavity_list. The default is 0.
HC_index : int, optional
Index of the harmonic cavity in cavity_list. The default is 1.
Returns
-------
eta : float
Amplification factor, if bigger than 1 then the beam is possibly
in the PTBL regime, Eq. (22) of [1].
RQth : float
R/Q threshold, Eq. (24) of [1].
f : float
f function, Eq. (23) of [1].
Reference
---------
[1] : He, T. (2022). Novel perturbation method for judging the
stability of the equilibrium solution in the presence of passive
harmonic cavities. Physical Review Accelerators and Beams, 25(9),
094402.
"""
if m is None:
m = self.ring.h
MC = self.cavity_list[MC_index]
HC = self.cavity_list[HC_index]
a = np.exp(-HC.wr * self.ring.T0 / (2 * HC.Q))
theta = HC.detune * 2 * np.pi * self.ring.T0
dtheta = np.arcsin((1-a) * np.cos(theta / 2) /
(np.sqrt(1 + a**2 - 2 * a * np.cos(theta))))
k = np.arange(1, m)
d_k = np.exp(-1 * HC.wr * self.ring.T0 * (k-1) / (2 * HC.Q * m))
theta_k = theta/2 + dtheta - (k-1) / m * theta
eps_k = 1 - np.sin(np.pi / 2 - k * 2 * np.pi / m)
num = np.sum(eps_k * d_k * np.cos(theta_k))
f = num / (m * np.sqrt(1 + a**2 - 2 * a * np.cos(theta)))
eta = (2 * np.pi * HC.m**2 * self.F[HC_index] * self.ring.h * I0 *
HC.Rs / HC.Q * f /
(MC.Vc * np.sin(MC.theta - self.PHI[HC_index] / HC.m)))
RQth = (MC.Vc * np.sin(MC.theta - self.PHI[HC_index] / HC.m) /
(2 * np.pi * HC.m**2 * self.F[1] * self.ring.h * I0 * f))
return (eta, RQth, f)
@property
def R_factor(self) -> float:
"""
Touschek lifetime ratio as defined in [1].
Reference
---------
[1] : Byrd, J. M., and M. Georgsson. "Lifetime increase using passive
harmonic cavities in synchrotron light sources." Physical Review
Special Topics-Accelerators and Beams 4.3 (2001): 030701.
"""
rho0 = gaussian_bunch(self.tau0, self.ring.sigma_0) / c
R = trapezoid(rho0**2, self.tau0) / trapezoid(self.rho0**2, self.tau0)
return R
@property
def bunch_spectrum(self) -> NDArray:
"""Return the bunch equilibrium spectrum."""
dft_rho = rfft(self.rho0) * self.sampling / self.N
return np.abs(dft_rho) / np.max(np.abs(dft_rho))
[docs]
def plot_bunch_spectrum(self) -> None:
"""Plot the bunch equilibrium spectrum"""
plt.plot(self.freq0 * 1e-9, self.bunch_spectrum)
plt.xlabel(r"Frequency [GHz]")
plt.xlim([0, self.freq0.max() * 1e-9 / 3])
plt.title("Equilibrium bunch spectrum")