# -*- coding: utf-8 -*-
"""
Module where the class to store the optic functions is defined.
"""
from typing import Any
import at
import matplotlib.pyplot as plt
import numpy as np
from matplotlib.axes import Axes
from numpy.typing import NDArray
from scipy.interpolate import interp1d
[docs]
class Optics:
"""
Class used to handle optic functions.
Parameters
----------
lattice_file : str, optional if local_beta, local_alpha and
local_dispersion are specified.
AT lattice file path.
tracking_loc : float, optional
Longitudinal position where the tracking is in the AT lattice in [m].
Only used if an AT lattice is loaded to compute the optic functions at
this location.
Default is 0.
local_beta : array of shape (2,), optional if lattice_file is specified.
Beta function at the location of the tracking. Default is mean beta if
lattice has been loaded.
local_alpha : array of shape (2,), optional if lattice_file is specified.
Alpha function at the location of the tracking. Default is mean alpha
if lattice has been loaded.
local_dispersion : array of shape (4,), optional if lattice_file is
specified.
Dispersion function and derivative at the location of the tracking.
Default is zero if lattice has been loaded.
Attributes
----------
use_local_values : bool
True if no lattice has been loaded.
local_gamma : array of shape (2,)
Gamma function at the location of the tracking.
lattice : AT lattice
average_beta : array of shape (2,)
H and V average beta functions.
Methods
-------
load_from_AT(lattice_file, **kwargs)
Load a lattice from accelerator toolbox (AT).
setup_interpolation()
Setup interpolation of the optic functions.
beta(position)
Return beta functions at specific locations given by position.
alpha(position)
Return alpha functions at specific locations given by position.
gamma(position)
Return gamma functions at specific locations given by position.
dispersion(position)
Return dispersion functions at specific locations given by position.
plot(self, var, option, n_points=1000)
Plot optical variables.
"""
[docs]
def __init__(self,
lattice_file: str | None = None,
tracking_loc: float = 0,
local_beta: NDArray | None = None,
local_alpha: NDArray | None = None,
local_dispersion: NDArray | None = None,
**kwargs):
if lattice_file is not None:
self.lattice_file = lattice_file
self.use_local_values = False
self.tracking_loc = tracking_loc
self.load_from_AT(lattice_file, **kwargs)
if local_beta is None:
self._local_beta = self.beta(self.tracking_loc)
else:
self._local_beta = local_beta
if local_alpha is None:
self._local_alpha = self.alpha(self.tracking_loc)
else:
self._local_alpha = local_alpha
if local_dispersion is None:
self.local_dispersion = self.dispersion(self.tracking_loc)
else:
self.local_dispersion = local_dispersion
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
else:
self.lattice_file = None
self.use_local_values = True
self.tracking_loc = None
self._local_beta = local_beta
self._local_alpha = local_alpha
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
self.local_dispersion = local_dispersion
def __repr__(self) -> str:
return (
f"Optics(lattice_file={self.lattice_file}, tracking_loc={self.tracking_loc}, "
f"local_beta={self._local_beta}, local_alpha={self._local_alpha}, "
f"local_dispersion={self.local_dispersion})")
def __str__(self) -> str:
return (f"Optics Configuration:\n"
f" Lattice File: {self.lattice_file}\n"
f" Tracking Location: {self.tracking_loc}\n"
f" Local Beta: {self._local_beta}\n"
f" Local Alpha: {self._local_alpha}\n"
f" Local Dispersion: {self.local_dispersion}")
[docs]
def load_from_AT(self, lattice_file: str, **kwargs):
"""
Load a lattice from accelerator toolbox (AT).
Parameters
----------
lattice_file : str
AT lattice file path.
n_points : int or float, optional
Minimum number of points to use for the optic function arrays.
periodicity : int, optional
Lattice periodicity, if not specified the AT lattice periodicity is
used.
"""
self.n_points = int(kwargs.get("n_points", 1e3))
periodicity = kwargs.get("periodicity")
self.lattice = at.load_lattice(lattice_file)
if self.lattice.radiation:
self.lattice.radiation_off()
lattice = self.lattice.slice(slices=self.n_points)
refpts: Any = np.arange(0, len(lattice))
_, tune, chrom, twiss = at.linopt(lattice,
refpts=refpts,
get_chrom=True)
if periodicity is None:
self.periodicity = lattice.periodicity
else:
self.periodicity = periodicity
if self.periodicity > 1:
periods = np.arange(0, self.periodicity)
shift = periods * twiss.s_pos[-1]
shift = np.repeat(shift, len(twiss.s_pos))
pos = np.tile(twiss.s_pos.T, self.periodicity) + shift
else:
pos = twiss.s_pos
self.position = pos
self.beta_array = np.tile(twiss.beta.T, self.periodicity)
self.alpha_array = np.tile(twiss.alpha.T, self.periodicity)
self.dispersion_array = np.tile(twiss.dispersion.T, self.periodicity)
self.mu_array = np.tile(twiss.mu.T, self.periodicity)
self.position = np.append(self.position, self.lattice.circumference)
self.beta_array = np.append(self.beta_array,
self.beta_array[:, 0:1],
axis=1)
self.alpha_array = np.append(self.alpha_array,
self.alpha_array[:, 0:1],
axis=1)
self.dispersion_array = np.append(self.dispersion_array,
self.dispersion_array[:, 0:1],
axis=1)
self.mu_array = np.append(self.mu_array, self.mu_array[:, 0:1], axis=1)
self.gamma_array = (1 + self.alpha_array**2) / self.beta_array
self.tune = tune * self.periodicity
self.chro = chrom * self.periodicity
self.ac = at.get_mcf(self.lattice)
self.mu_array[:, -1] = (np.floor(self.mu_array[:, -2] /
(2 * np.pi)) + self.tune) * 2 * np.pi
self.setup_interpolation()
[docs]
def setup_interpolation(self):
"""Setup interpolation of the optic functions."""
self.betaX = interp1d(self.position,
self.beta_array[0, :],
kind='linear')
self.betaY = interp1d(self.position,
self.beta_array[1, :],
kind='linear')
self.alphaX = interp1d(self.position,
self.alpha_array[0, :],
kind='linear')
self.alphaY = interp1d(self.position,
self.alpha_array[1, :],
kind='linear')
self.gammaX = interp1d(self.position,
self.gamma_array[0, :],
kind='linear')
self.gammaY = interp1d(self.position,
self.gamma_array[1, :],
kind='linear')
self.dispX = interp1d(self.position,
self.dispersion_array[0, :],
kind='linear')
self.disppX = interp1d(self.position,
self.dispersion_array[1, :],
kind='linear')
self.dispY = interp1d(self.position,
self.dispersion_array[2, :],
kind='linear')
self.disppY = interp1d(self.position,
self.dispersion_array[3, :],
kind='linear')
self.muX = interp1d(self.position, self.mu_array[0, :], kind='linear')
self.muY = interp1d(self.position, self.mu_array[1, :], kind='linear')
@property
def local_beta(self) -> NDArray:
"""
Return beta function at the location defined by the lattice file.
"""
return self._local_beta
@local_beta.setter
def local_beta(self, beta_array):
"""
Set the values of beta function. Gamma function is automatically
recalculated after the new value of beta function is set.
Parameters
----------
beta_array : array of shape (2,)
Beta function in the horizontal and vertical plane.
"""
self._local_beta = beta_array
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
@property
def local_alpha(self) -> NDArray:
"""
Return alpha function at the location defined by the lattice file.
"""
return self._local_alpha
@local_alpha.setter
def local_alpha(self, alpha_array):
"""
Set the value of beta functions. Gamma function is automatically
recalculated after the new value of alpha function is set.
Parameters
----------
alpha_array : array of shape (2,)
Alpha function in the horizontal and vertical plane.
"""
self._local_alpha = alpha_array
self._local_gamma = (1 + self._local_alpha**2) / self._local_beta
@property
def local_gamma(self) -> NDArray:
"""
Return beta function at the location defined by the lattice file.
"""
return self._local_gamma
[docs]
def beta(self, position: NDArray | float) -> NDArray:
"""
Return beta functions at specific locations given by position. If no
lattice has been loaded, local values are returned.
Parameters
----------
position : array or float
Longitudinal position at which the beta functions are returned.
Returns
-------
beta : array
Beta functions.
"""
if self.use_local_values:
return np.outer(self.local_beta, np.ones_like(position))
beta = [self.betaX(position), self.betaY(position)]
return np.array(beta)
[docs]
def alpha(self, position: NDArray | float) -> NDArray:
"""
Return alpha functions at specific locations given by position. If no
lattice has been loaded, local values are returned.
Parameters
----------
position : array or float
Longitudinal position at which the alpha functions are returned.
Returns
-------
alpha : array
Alpha functions.
"""
if self.use_local_values:
return np.outer(self.local_alpha, np.ones_like(position))
alpha = [self.alphaX(position), self.alphaY(position)]
return np.array(alpha)
[docs]
def gamma(self, position: NDArray | float) -> NDArray:
"""
Return gamma functions at specific locations given by position. If no
lattice has been loaded, local values are returned.
Parameters
----------
position : array or float
Longitudinal position at which the gamma functions are returned.
Returns
-------
gamma : array
Gamma functions.
"""
if self.use_local_values:
return np.outer(self.local_gamma, np.ones_like(position))
gamma = [self.gammaX(position), self.gammaY(position)]
return np.array(gamma)
[docs]
def dispersion(self, position: NDArray | float) -> NDArray:
"""
Return dispersion functions at specific locations given by position.
If no lattice has been loaded, local values are returned.
Parameters
----------
position : array or float
Longitudinal position at which the dispersion functions are
returned.
Returns
-------
dispersion : array
Dispersion functions.
"""
if self.use_local_values:
return np.outer(self.local_dispersion, np.ones_like(position))
dispersion = [
self.dispX(position),
self.disppX(position),
self.dispY(position),
self.disppY(position)
]
return np.array(dispersion)
[docs]
def mu(self, position: NDArray | float) -> NDArray:
"""
Return phase advances at specific locations given by position.
If no lattice has been loaded, None is returned.
Parameters
----------
position : array or float
Longitudinal position at which the phase advances are returned.
Returns
-------
mu : array
Phase advances.
"""
if self.use_local_values:
return np.outer(np.array([0, 0]), np.ones_like(position))
mu = [self.muX(position), self.muY(position)]
return np.array(mu)
[docs]
def plot(self,
var: str,
option: str,
n_points: int = 1000,
ax: Axes | None = None) -> Axes:
"""
Plot optical variables.
Parameters
----------
var : {"beta", "alpha", "gamma", "dispersion", "mu"}
Optical variable to be plotted.
option : str
If var = "beta", "alpha" and "gamma", option = {"x","y"} specifying
the axis of interest.
If var = "dispersion", option = {"x","px","y","py"} specifying the
axis of interest for the dispersion function or its derivative.
n_points : int
Number of points on the plot. The default value is 1000.
ax : Axes, optional
Axes where the plot is displayed. If None, a new figure is created.
Return
------
ax : Axes
Axes with the plot on it.
"""
var_dict = {
"beta": self.beta,
"alpha": self.alpha,
"gamma": self.gamma,
"dispersion": self.dispersion,
"mu": self.mu
}
if var == "dispersion":
option_dict = {"x": 0, "px": 1, "y": 2, "py": 3}
label = ["D$_{x}$ (m)", "D'$_{x}$", "D$_{y}$ (m)", "D'$_{y}$"]
ylabel = label[option_dict[option]]
elif var == "beta" or var == "alpha" or var == "gamma" or var == "mu":
option_dict = {"x": 0, "y": 1}
label_dict = {
"beta": "$\\beta$",
"alpha": "$\\alpha$",
"gamma": "$\\gamma$",
"mu": "$\\mu$"
}
if option == "x": label_sup = "$_{x}$"
elif option == "y": label_sup = "$_{y}$"
unit = {
"beta": " (m)",
"alpha": "",
"gamma": " (m$^{-1}$)",
"mu": ""
}
ylabel = label_dict[var] + label_sup + unit[var]
else:
raise ValueError("Variable name is not found.")
if self.use_local_values is not True:
position = np.linspace(0, self.lattice.circumference,
int(n_points))
else:
position = np.linspace(0, 1)
var_list = var_dict[var](position)[option_dict[option]]
if ax is None:
fig, ax = plt.subplots()
ax.plot(position, var_list)
ax.set_xlabel("position (m)")
ax.set_ylabel(ylabel)
return ax
@property
def average_beta(self) -> NDArray:
"""
Return average beta functions.
If self.use_local_values, self.local_beta is returned.
Returns
-------
average_beta : array of shape (2,)
H and V average beta functions.
"""
if self.use_local_values:
return self.local_beta
L = self.position[-1]
position = np.linspace(0, L, int(self.n_points))
length = position[1:] - position[:-1]
center = (position[1:] + position[:-1]) / 2
beta = self.beta(center)
beta_H_star = 1 / L * (length * beta[0, :]).sum()
beta_V_star = 1 / L * (length * beta[1, :]).sum()
return np.array([beta_H_star, beta_V_star])