Example-10: Dipole element factory (cylindrical multipole)

In this example dipole factory is illustrated for dipole with multipole cylindrical multipoles.

The dipole hamiltonian is:

$ \begin{align} & H(q_x, q_y, q_s, p_x, p_y, p_s; s) = \frac{p_s}{\beta} - t(s)(q_x p_y - q_y p_x) - (1 + h(s) q_x) \left(\sqrt{P_s^2 - P_x^2 - P_y^2 - \frac{1}{\beta^2 \gamma^2}} + a_s(q_x, q_y, q_s; s)\right) \\ & \\ & P_s = p_s + 1/\beta - \varphi(q_x, q_y, q_s; s) \\ & P_x = p_x - a_x(q_x, q_y, q_s; s) \\ & P_y = p_y - a_y(q_x, q_y, q_s; s) \\ \\ & (a_x, a_y, a_s) = (0, 0, -\frac{1}{1 + q_x/\rho}\left(\frac{q_x}{ \rho} + \frac{q_x^2}{2 \rho^2} \right) + U(q_x, q_y; \rho) )\\ & \varphi = 0 \\ & t = 0 \\ & h = \frac{1}{\rho} = \frac{\alpha}{l} \end{align} $

The constructed element signature is:

def dipole(qsps:Array, length:Array, angle:Array, kq_n:Array, kq_s:Array, ks_n:Array, ks_s:Array, ko_n:Array, ko_s:Array) -> Array:
    ...

Note, no fringe effects are icluded.

Cylindrical potential is precomputed for quadrupole, sextuipole and octupole components up to degree 10 in transverse coordinates.

import jax
from jax import jit
from jax import jacrev

from elementary.util import ptc
from elementary.util import beta
from elementary.dipole import dipole_factory

jax.numpy.set_printoptions(linewidth=256, precision=12)

# Set data type

jax.config.update("jax_enable_x64", True)

# Set device

device, *_ = jax.devices('cpu')
jax.config.update('jax_default_device', device)

# Set initial condition

(q_x, q_y, q_s) = qs = jax.numpy.array([-0.01, 0.005, 0.001])
(p_x, p_y, p_s) = ps = jax.numpy.array([0.001, 0.001, -0.0001])
qsps = jax.numpy.hstack([qs, ps])

# Define generic dipole element
# Note, exact should be set to false

gamma = 10**3
element = jit(dipole_factory(exact=False, multipole=True, beta=beta(gamma), gamma=gamma, order=2**1, iterations=1E3))

# Compare with PTC

length = jax.numpy.float64(1.0)
angle = jax.numpy.float64(0.05)

kq_n = jax.numpy.float64(-2.0)
kq_s = jax.numpy.float64(+1.5)
ks_n = jax.numpy.float64(-50.0)
ks_s = jax.numpy.float64(+75.0)
ko_n = jax.numpy.float64(-100.0)
ko_s = jax.numpy.float64(+500.0)

print(res := element(qsps, length, angle, kq_n, kq_s, ks_n, ks_s, ko_n, ko_s))
print(ref := ptc(qsps, 'sbend', {'l': float(length), 'angle': float(angle), 'knl': f'{{0.0,{float(kq_n*length)}, {float(ks_n*length)}, {float(ko_n*length)}}}', 'ksl': f'{{0.0,{float(kq_s*length)}, {float(ks_s*length)}, {float(ko_s*length)}}}', 'kill_ent_fringe': 'true', 'kill_exi_fringe': 'true'}, gamma=gamma))
print(jax.numpy.allclose(res, ref))

[-0.017200420757 -0.003042779236  0.001519392779 -0.017822768071 -0.015620313086 -0.0001        ]
[-0.017200420757 -0.003042779236  0.00151939277  -0.017822768072 -0.015620313086 -0.0001        ]
True