from itertools import count
from typing import Optional
import numpy as np
from matplotlib import patches
from ..lattice import Lattice
from .base import BaseElement
from .utils import straight_element
[docs]class Sextupole(BaseElement):
"""Sextupole element.
Args:
k: Strength in meters^-3.
l: Element length in meters.
name (optional): Element name.
Attributes:
k: Sextupole strength in meters^-3.
l: Element length in meters.
[docs][docs] length: Element length in meters.
m: Element phase space transfer matrix.
name: Element name.
"""
def __init__(self, k: float, l: float, name: Optional[str] = None):
self.k = k
self.l = l
if name is None:
name = f"sextupole_{next(self._instance_count)}"
super().__init__("k", "l", "name")
self.name = name
def _non_linear_term(self, phase_coords: np.ndarray) -> np.ndarray:
# see http://cern.ch/mad8/doc/phys_guide.pdf chapter 5.5.3
out = np.zeros(phase_coords.shape)
# strength depends on dp/p
k = self.k / (1 + phase_coords[4])
# k = self.k
out[1] = -0.5 * k * self.l * (phase_coords[0] ** 2 - phase_coords[2] ** 2)
out[3] = k * self.l * (phase_coords[0] * phase_coords[2])
return out
def _get_length(self) -> float:
return self.l
def _get_transfer_matrix(self) -> np.ndarray:
out = np.identity(5)
out[0, 1] = self.l
out[2, 3] = self.l
return out
[docs] def slice(self, n_sextupoles: int) -> Lattice:
"""Slice the element into a many smaller elements.
Args:
n_sextupoles: Number of :py:class:`Sextupole` elements.
Returns:
:py:class:`~accelerator.lattice.Lattice` of sliced :py:class:`Sextupole` elements.
"""
out = [
Sextupole(
self.k,
self.l / n_sextupoles,
name=f"{self.name}_slice_{i}",
)
for i in range(n_sextupoles)
]
return Lattice(out)
def _get_patch(self, s: float) -> patches.Patch:
label = "Sextupole"
colour = "tab:green"
return patches.Rectangle((s, -1), self.length, 2, facecolor=colour, label=label)
@staticmethod
def _dxztheta_ds(theta: float, d_s: float) -> np.ndarray:
return straight_element(theta, d_s)
[docs]class SextupoleThin(BaseElement):
"""Thin Sextupole element.
Args:
k: Strength in meters^-2.
name (optional): Element name.
Attributes:
k: Sextupole strength in meters^-2.
length: Element length in meters.
m: Element phase space transfer matrix.
name: Element name.
"""
_instance_count = count(0)
def __init__(self, k: float, name: Optional[str] = None):
self.k = k
if name is None:
name = f"sextupole_thin_{next(self._instance_count)}"
super().__init__("k", "name")
self.name = name
def _non_linear_term(self, phase_coords: np.ndarray) -> np.ndarray:
out = np.zeros(phase_coords.shape)
out[1] = -0.5 * self.k * (phase_coords[0] ** 2 - phase_coords[2] ** 2)
out[3] = self.k * (phase_coords[0] * phase_coords[2])
return out
def _get_length(self) -> float:
return 0
def _get_transfer_matrix(self) -> np.ndarray:
return np.identity(5)
def _get_patch(self, s: float) -> patches.Patch:
label = "Thin Sextupole"
colour = "tab:green"
return patches.FancyArrowPatch(
(s, 1),
(s, -1),
arrowstyle=patches.ArrowStyle("|-|"),
label=label,
edgecolor=colour,
facecolor=colour,
)
@staticmethod
def _dxztheta_ds(
theta: float, d_s: float # pylint: disable=unused-argument
) -> np.ndarray:
return np.array([0, 0, 0])