Example-64: AT ID

# Import

import torch
from torch import Tensor

from pathlib import Path
from tqdm import tqdm

import matplotlib
from matplotlib import pyplot as plt
from matplotlib.patches import Rectangle
matplotlib.rcParams['text.usetex'] = True

from model.library.element import Element
from model.library.line import Line
from model.library.corrector import Corrector
from model.library.quadrupole import Quadrupole
from model.library.matrix import Matrix
from model.library.kickmap import KM

from model.command.external import load_lattice
from model.command.build import build
from model.command.tune import tune
from model.command.tune import chromaticity
from model.command.orbit import dispersion
from model.command.orbit import ORM
from model.command.twiss import twiss
from model.command.advance import advance
from model.command.coupling import coupling

import at
from model.interface.at import convert

# Set data type and device

Element.dtype = dtype = torch.float64
Element.device = device = torch.device('cpu')

# Load lattice (ELEGANT table)
# Note, lattice is allowed to have repeated elements

path = Path('elettra.lte')
data = load_lattice(path)

# Build and setup lattice

ring:Line = build('RING', 'ELEGANT', data)

# Flatten sublines

ring.flatten()

# Remove all marker elements but the ones starting with MLL (long straight section centers)

ring.remove_group(pattern=r'^(?!MSS_)(?!MLL_).*', kinds=['Marker'])


# Set sextupole integration order and step size

ring.order = (('Sextupole', 1), )
ring.ns = (('Sextupole', 0.01), )

# Set linear dipoles

def apply(element:Element) -> None:
    element.linear = True

ring.apply(apply, kinds=['Dipole'])

# Insert correctors

for name, *_ in ring.layout():
    if name.startswith('CH'):
        corrector = Corrector(f'{name}_CXY', factor=1)
        ring.split((1 + 1, None, [name], None), paste=[corrector])

# Merge drifts

ring.merge()

# Change lattice start start

ring.start = "BPM_S01_01"

# Split BPMs

ring.split((None, ['BPM'], None, None))

# Roll lattice

ring.roll(1)

# Splice

ring.splice()

# Describe

ring.describe

{'BPM': 168,
 'Drift': 744,
 'Dipole': 156,
 'Sextupole': 240,
 'Quadrupole': 120,
 'Corrector': 24,
 'Marker': 24}

# Compute tunes (fractional part)

nux, nuy = tune(ring, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype, device=device)
etaqx, etapx, etaqy, etapy = dispersion(ring, orbit, [], limit=1)

# Compute twiss parameters

ax, bx, ay, by = twiss(ring, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux, muy = advance(ring, [], alignment=False, matched=True).T

# Compute coupling

c = coupling(ring, [])

# Compute chromaticity

psi = chromaticity(ring, [], matched=True)

# Define ID (linear model)

A = torch.tensor([[-0.0344386, 0., 0., 0.], [0., -0.0445673, 0., 0.], [0., 0., 0.056303, 0.], [0., 0., 0., 0.0804237]], dtype=dtype)

CENTER = 0.0
ID = Matrix('ID', length=0.0, A=A[torch.triu(torch.ones_like(A, dtype=torch.bool))].tolist())

# Insert ID (linear model)

error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=CENTER*1.0E-3)
error.splice()
error.describe

{'BPM': 168,
 'Drift': 745,
 'Dipole': 156,
 'Sextupole': 240,
 'Quadrupole': 120,
 'Corrector': 24,
 'Marker': 24,
 'Matrix': 1}

# Generate AT lattice

lattice = convert(error, energy=2.4, alignment=False)

# Compute optics

refpts = lattice.uint32_refpts(at.Monitor)
print(len(refpts))
_, ringdata, elemdata = lattice.get_optics(refpts=refpts, get_chrom=True)

168

# Set optics at BPMs

nux_id_lm_at, nuy_id_lm_at = ringdata['tune']
psi_id_lm_at = ringdata['chromaticity']

ax_id_lm_at, ay_id_lm_at = elemdata['alpha'].T
bx_id_lm_at, by_id_lm_at = elemdata['beta'].T
mux_id_lm_at, muy_id_lm_at = elemdata['mu'].T
etaqx_id_lm_at, etapx_id_lm_at, etaqy_id_lm_at, etapy_id_lm_at = elemdata['dispersion'].T

# Compute tunes (fractional part)

nux_id_lm, nuy_id_lm = tune(error, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_lm, etapx_id_lm, etaqy_id_lm, etapy_id_lm = dispersion(error, orbit, [], limit=1)

# Compute twiss parameters

ax_id_lm, bx_id_lm, ay_id_lm, by_id_lm = twiss(error, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux_id_lm, muy_id_lm = advance(error, [], alignment=False, matched=True).T

# Compute coupling

c_id_lm = coupling(error, [])

# Compute chromaticity

psi_id_lm = chromaticity(error, [], matched=True)

# Tune shifts

print((nux - nux_id_lm).numpy())
print((nuy - nuy_id_lm).numpy())
print()

print((nux - nux_id_lm_at).numpy())
print((nuy - nuy_id_lm_at).numpy())
print()

0.02566042669989943
-0.011228413019111483

0.025711104422837916
-0.011228053550816186

# Chromaticity

print(psi.numpy())
print(psi_id_lm.numpy())
print(psi_id_lm_at)

[2.02960315 2.01312355]
[1.71483742 2.01653988]
[1.71445621 2.01656829]

# AT coupling

index, *_ = lattice.get_uint32_index("MSS_S10")
rf = at.RFCavity("CAV", 0.0, 2_000_000.0, 499_654_096.666667, 432, 2_400_000_000.0, PassMethod="CavityPass")
lattice.insert(int(index) + 1, rf)
lattice.enable_6d()
_, data, _ = lattice.ohmi_envelope(refpts=refpts)
data.mode_emittances

array([ 2.40439332e-10, -1.21810204e-37,  1.71831614e-06])

# Define ID (load kick map)

CENTER = 0.0
ID = KM(name='ID', path=Path('id.mat'), energy=2.4, count=40, insertion=True)

# Insert ID (kick map)

error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=CENTER*1.0E-3)
error.splice()
error.describe

{'BPM': 168,
 'Drift': 745,
 'Dipole': 156,
 'Sextupole': 240,
 'Quadrupole': 120,
 'Corrector': 24,
 'Marker': 24,
 'KM': 1}

# Generate AT lattice

lattice = convert(error, energy=2.4, alignment=False)

# Compute optics

refpts = lattice.uint32_refpts(at.Monitor)
print(len(refpts))
_, ringdata, elemdata = lattice.get_optics(refpts=refpts, get_chrom=True)

168

# Set optics at BPMs

nux_id_km_at, nuy_id_km_at = ringdata['tune']
psi_id_km_at = ringdata['chromaticity']

ax_id_km_at, ay_id_km_at = elemdata['alpha'].T
bx_id_km_at, by_id_km_at = elemdata['beta'].T
mux_id_km_at, muy_id_km_at = elemdata['mu'].T
etaqx_id_km_at, etapx_id_km_at, etaqy_id_km_at, etapy_id_km_at = elemdata['dispersion'].T

# Compute tunes (fractional part)

nux_id_km, nuy_id_km = tune(error, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_km, etapx_id_km, etaqy_id_km, etapy_id_km = dispersion(error, orbit, [], limit=1)

# Compute twiss parameters

ax_id_km, bx_id_km, ay_id_km, by_id_km = twiss(error, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux_id_km, muy_id_km = advance(error, [], alignment=False, matched=True).T

# Compute coupling

c_id_km = coupling(error, [])

# Compute chromaticity

psi_id_km = chromaticity(error, [], matched=True)

# Tune shifts

print((nux - nux_id_km).numpy())
print((nuy - nuy_id_km).numpy())
print()

print((nux - nux_id_km_at).numpy())
print((nuy - nuy_id_km_at).numpy())
print()

0.025636790558097755
-0.011262894479845187

0.02568748173391394
-0.011262534209166342

# Chromaticity

print(psi.numpy())
print(psi_id_km.numpy())
print(psi_id_km_at)

[2.02960315 2.01312355]
[1.74225348 1.99775283]
[1.74187601 1.99778124]

# AT coupling

index, *_ = lattice.get_uint32_index("MSS_S10")
rf = at.RFCavity("CAV", 0.0, 2_000_000.0, 499_654_096.666667, 432, 2_400_000_000.0, PassMethod="CavityPass")
lattice.insert(int(index) + 1, rf)
lattice.enable_6d()
_, data, _ = lattice.ohmi_envelope(refpts=refpts)
data.mode_emittances

array([2.40431968e-10, 2.77736999e-25, 1.71831568e-06])