ELETTRA-08: Local ID correction (local correction: tunes, twiss and dispersion)
[1]:
# In this example model free local ID correction is performed using tunes, twiss and dispersion
[2]:
# Import
import torch
from torch import Tensor
from pathlib import Path
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.quadrupole import Quadrupole
from model.library.matrix import Matrix
from model.command.external import load_lattice
from model.command.build import build
from model.command.tune import tune
from model.command.orbit import dispersion
from model.command.twiss import twiss
from model.command.advance import advance
from model.command.coupling import coupling
from model.command.wrapper import Wrapper
from model.command.wrapper import forward
from model.command.wrapper import inverse
from model.command.wrapper import normalize
[3]:
# Set data type and device
Element.dtype = dtype = torch.float64
Element.device = device = torch.device('cpu')
[4]:
# Load lattice (ELEGANT table)
# Note, lattice is allowed to have repeated elements
path = Path('elettra.lte')
data = load_lattice(path)
[5]:
# 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'^(?!MLL_).*', kinds=['Marker'])
# Replace all sextupoles with quadrupoles
def factory(element:Element) -> None:
table = element.serialize
table.pop('ms', None)
return Quadrupole(**table)
ring.replace_group(pattern=r'', factory=factory, kinds=['Sextupole'])
# Set linear dipoles
def apply(element:Element) -> None:
element.linear = True
ring.apply(apply, kinds=['Dipole'])
# Merge drifts
ring.merge()
# Change lattice start
ring.start = "BPM_S01_01"
# Split BPMs
ring.split((None, ['BPM'], None, None))
# Roll lattice
ring.roll(1)
# Splice lattice
ring.splice()
# Describe
ring.describe
[5]:
{'BPM': 168, 'Drift': 708, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 12}
[6]:
# Compute tunes (fractional part)
nux, nuy = tune(ring, [], matched=True, limit=1)
[7]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx, etapx, etaqy, etapy = dispersion(ring, orbit, [], limit=1)
[8]:
# Compute twiss parameters
ax, bx, ay, by = twiss(ring, [], matched=True, advance=True, full=False).T
[9]:
# Compute phase advances
mux, muy = advance(ring, [], alignment=False, matched=True).T
[10]:
# Compute coupling
c = coupling(ring, [])
[11]:
# Several local knobs can be used to correct ID effects
# Normal quadrupole correctors
nkn = ['OCT_S01_02', 'QF_S01_02', 'QD_S01_02', 'QD_S01_03', 'QF_S01_03', 'OCT_S01_03']
# Skew quadrupole correctors
nks = ['SD_S01_05', 'SH_S01_02', 'SH_S01_03', 'SD_S01_06']
[12]:
# Define tune observable
def observable_tune(kn, ks):
nux, nuy = tune(ring,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None),
matched=True)
return torch.stack([nux, nuy])
kn = torch.tensor(6*[0.0], dtype=dtype)
ks = torch.tensor(4*[0.0], dtype=dtype)
print(observable_tune(kn, ks))
print(torch.func.jacfwd(observable_tune, 0)(kn, ks))
print(torch.func.jacfwd(observable_tune, 1)(kn, ks))
tensor([0.2994, 0.1608], dtype=torch.float64)
tensor([[ 0.0382, 0.2439, 0.0873, 0.0873, 0.2439, 0.0382],
[-0.1558, -0.1247, -0.0525, -0.0525, -0.1247, -0.1558]],
dtype=torch.float64)
tensor([[0., 0., 0., 0.],
[0., 0., 0., 0.]], dtype=torch.float64)
[13]:
# Define twiss observable
def observable_twiss(kn, ks):
_, bx, _, by = twiss(ring,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None),
matched=True,
advance=True,
full=False,
convert=True).T
return torch.stack([bx, by]).T
kn = torch.tensor(6*[0.0], dtype=dtype)
ks = torch.tensor(4*[0.0], dtype=dtype)
print(observable_twiss(kn, ks))
print(torch.func.jacfwd(observable_twiss, 0)(kn, ks))
print(torch.func.jacfwd(observable_twiss, 1)(kn, ks))
tensor([[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317],
[ 4.1498, 2.3317],
[ 1.5214, 5.7966],
[ 2.9455, 4.5619],
[ 2.7752, 6.4643],
[ 5.4767, 6.2506],
[ 5.5205, 13.7764],
[10.0310, 5.1758],
[10.0413, 5.2348],
[ 2.3579, 19.8903],
[ 5.4767, 6.2506],
[ 1.9946, 12.4186],
[ 2.7752, 6.4643],
[ 2.9455, 4.5619],
[ 4.1498, 2.3317]], dtype=torch.float64)
tensor([[[-0.3099, -1.0966, -0.2110, 2.0799, 6.0433, 0.9979],
[ 2.1106, 1.5414, 0.5463, 0.3157, 1.0453, 1.5421]],
[[-0.2323, -1.2062, -0.3723, 0.5677, 1.7235, 0.3043],
[ 6.2351, 4.7589, 1.8256, 0.1322, 1.1163, 2.0612]],
[[ 0.7034, 4.6561, 1.6870, 0.8745, 2.1237, 0.2396],
[ 5.2657, 4.1662, 1.6936, -0.4625, -0.4717, -0.0486]],
...,
[[-0.3857, -1.9356, -0.5808, 1.1147, 3.3489, 0.5823],
[-7.4183, -5.8282, -2.3435, 0.4910, 0.2690, -0.4318]],
[[-0.6349, -3.6847, -1.2339, 0.5486, 1.8707, 0.3826],
[-5.0785, -4.1567, -1.7764, 0.9865, 1.7863, 1.7313]],
[[ 0.5198, 4.0743, 1.5998, 2.2819, 6.2002, 0.9092],
[-1.6946, -1.5203, -0.7302, 0.8487, 1.8759, 2.1969]]],
dtype=torch.float64)
tensor([[[0., 0., 0., 0.],
[0., 0., 0., 0.]],
[[0., 0., 0., 0.],
[0., 0., 0., 0.]],
[[0., 0., 0., 0.],
[0., 0., 0., 0.]],
...,
[[0., 0., 0., 0.],
[0., 0., 0., 0.]],
[[0., 0., 0., 0.],
[0., 0., 0., 0.]],
[[0., 0., 0., 0.],
[0., 0., 0., 0.]]], dtype=torch.float64)
[14]:
# Define dispersion observable
def observable_dispersion(kn, ks):
orbit = torch.tensor(4*[0.0], dtype=dtype)
etax, _, etay, _ = dispersion(ring,
orbit,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None))
return torch.stack([etax, etay]).T
kn = torch.tensor(6*[0.0], dtype=dtype)
ks = torch.tensor(4*[0.0], dtype=dtype)
print(observable_dispersion(kn, ks))
print(torch.func.jacfwd(observable_dispersion, 0)(kn, ks))
print(torch.func.jacfwd(observable_dispersion, 1)(kn, ks))
tensor([[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00],
[ 3.0380e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[-4.7351e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-6.6337e-06, 0.0000e+00],
[-3.0076e-06, 0.0000e+00],
[ 5.6430e-02, 0.0000e+00],
[ 2.8385e-02, 0.0000e+00],
[ 3.6084e-02, 0.0000e+00],
[ 3.6490e-02, 0.0000e+00],
[ 5.7064e-02, 0.0000e+00]], dtype=torch.float64)
tensor([[[ 9.7335e-07, 6.0423e-06, 2.1173e-06, 7.3312e-07, 1.7280e-06,
1.8313e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]],
[[ 6.5599e-07, 4.1426e-06, 1.4668e-06, 7.2980e-07, 1.8454e-06,
2.3522e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]],
[[-1.6495e-07, -6.3114e-07, -1.3640e-07, 1.1908e-06, 3.5055e-06,
5.9274e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]],
...,
[[-8.7080e-07, -5.4828e-06, -1.9379e-06, -9.1388e-07, -2.2912e-06,
-2.8621e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]],
[[-9.8106e-07, -6.2897e-06, -2.2471e-06, -1.4069e-06, -3.6711e-06,
-5.0143e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]],
[[ 6.0629e-07, 3.4972e-06, 1.1680e-06, -4.3526e-07, -1.4998e-06,
-3.0900e-07],
[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,
0.0000e+00]]], dtype=torch.float64)
tensor([[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 1.4255e-02, -1.1707e-06, -2.5867e-07, 2.7165e-04]],
[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 7.6600e-02, -5.3693e-06, 3.2503e-06, -2.4709e-02]],
[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[-6.4586e-02, 4.3706e-06, -3.5290e-06, 2.5280e-02]],
...,
[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 1.4670e-02, -2.0658e-06, -4.6029e-06, 2.4736e-02]],
[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 2.5641e-02, -1.1723e-06, 4.2357e-06, -2.6022e-02]],
[[ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00],
[ 1.4255e-02, -1.1707e-06, -2.5867e-07, 2.7165e-04]]],
dtype=torch.float64)
[15]:
# Note, for ideal lattice both tune and twiss observable derivatives with respect to skew quadrupole knobs are zero
# Only dispersion derivatives are not zero, which means that skew quadrupols are used to correct dispersion
[16]:
# Construct full target observable vector and corresponding responce matrix
def observable(knobs):
kn, ks = torch.split(knobs, [6, 4])
tunes = observable_tune(kn, ks)
betas = observable_twiss(kn, ks)
etas = observable_dispersion(kn, ks)
return torch.cat([tunes, betas.flatten(), etas.flatten()])
knobs = torch.tensor((6 + 4)*[0.0], dtype=dtype)
print((target := observable(knobs)).shape)
print((matrix := torch.func.jacfwd(observable)(knobs)).shape)
torch.Size([674])
torch.Size([674, 10])
[17]:
# To perform model free correction the response matrix should have full rank
print(torch.linalg.matrix_rank(matrix))
tensor(10)
[18]:
# Define ID model
# Note, only the flattened triangular part of the A and B matrices is passed
A = torch.tensor([[-0.03484222052711237, 1.0272120741819959E-7, -4.698931299341201E-9, 0.0015923185492594811],
[1.0272120579834892E-7, -0.046082787920135176, 0.0017792061173117564, 3.3551298301095784E-8],
[-4.6989312853101E-9, 0.0017792061173117072, 0.056853750760983084, -1.5929605363332683E-7],
[0.0015923185492594336, 3.3551298348653296E-8, -1.5929605261642905E-7, 0.08311631737263032]], dtype=dtype)
B = torch.tensor([[0.03649353186115209, 0.0015448347221877217, 0.00002719892025520868, -0.0033681183134964482],
[0.0015448347221877217, 0.13683886657005795, -0.0033198692682377406, 0.00006140578258682469],
[0.00002719892025520868, -0.0033198692682377406, -0.05260095308967722, 0.005019907688182885],
[-0.0033681183134964482, 0.00006140578258682469, 0.005019907688182885, -0.2531573249456863]], dtype=dtype)
ID = Matrix('ID',
length=0.0,
A=A[torch.triu(torch.ones_like(A, dtype=torch.bool))].tolist(),
B=B[torch.triu(torch.ones_like(B, dtype=torch.bool))].tolist())
[19]:
# Insert ID into the existing lattice
# This will replace the target marker
error = ring.clone()
error.flatten()
error.insert(ID, 'MLL_S01', position=0.0)
error.splice()
# Describe
error.describe
[19]:
{'BPM': 168,
'Drift': 708,
'Dipole': 156,
'Quadrupole': 360,
'Matrix': 1,
'Marker': 11}
[20]:
# Compute tunes (fractional part)
nux_id, nuy_id = tune(error, [], matched=True, limit=1)
[21]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id, etapx_id, etaqy_id, etapy_id = dispersion(error, orbit, [], limit=1)
[22]:
# Compute twiss parameters
ax_id, bx_id, ay_id, by_id = twiss(error, [], matched=True, advance=True, full=False).T
[23]:
# Compute phase advances
mux_id, muy_id = advance(error, [], alignment=False, matched=True).T
[24]:
# Compute coupling
c_id = coupling(error, [])
[25]:
# Tune shifts
print((nux - nux_id))
print((nuy - nuy_id))
tensor(0.0260, dtype=torch.float64)
tensor(-0.0114, dtype=torch.float64)
[26]:
# Coupling (minimal tune distance)
print(c)
print(c_id)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
[27]:
# Dispersion
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_id).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_id).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
[28]:
# Beta-beating
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_id)/bx).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_id)/by).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((bx - bx_id)/bx)**2).mean().sqrt())
print(100*(((by - by_id)/by)**2).mean().sqrt())
tensor(11.5994, dtype=torch.float64)
tensor(1.7916, dtype=torch.float64)
[29]:
# Phase advance
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_id)/mux).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_id)/muy).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((mux - mux_id)/mux)**2).mean().sqrt())
print(100*(((muy - muy_id)/muy)**2).mean().sqrt())
tensor(8.7941, dtype=torch.float64)
tensor(1.7778, dtype=torch.float64)
[30]:
# Define parametric observable vector (emulate tune measurement)
def observable_tune(kn, ks):
nux, nuy = tune(error,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None),
matched=True)
return torch.stack([nux, nuy])
def observable_twiss(kn, ks):
_, bx, _, by = twiss(error,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None),
matched=True,
advance=True,
full=False,
convert=True).T
return torch.stack([bx, by]).T
def observable_dispersion(kn, ks):
orbit = torch.tensor(4*[0.0], dtype=dtype)
etax, _, etay, _ = dispersion(error,
orbit,
[kn, ks],
('kn', None, nkn, None),
('ks', None, nks, None))
return torch.stack([etax, etay]).T
def observable(knobs):
kn, ks = torch.split(knobs, [6, 4])
tunes = observable_tune(kn, ks)
betas = observable_twiss(kn, ks)
etas = observable_dispersion(kn, ks)
return torch.cat([tunes, betas.flatten(), etas.flatten()])
[31]:
# Check the residual vector norm
knobs = torch.tensor((6 + 4)*[0.0], dtype=dtype)
print(((observable(knobs) - target)**2).sum())
tensor(64.2608, dtype=torch.float64)
[32]:
# Optimization loop (model free)
# Responce matrix (jacobian)
M = matrix.clone()
# Weighting covariance (sensitivity) matrix
epsilon = 1.0E-9
C = M @ M.T
C = C + epsilon*torch.eye(len(C), dtype=dtype)
# Cholesky decomposition
L = torch.linalg.cholesky(C)
# Whiten response
M = torch.linalg.solve_triangular(L, M, upper=False)
# Additional weights
# Can be used to extra weight selected observables, e.g. tunes
weights = torch.ones(len(M), dtype=dtype)
weights = weights.sqrt()
# Whiten response with additional weights
M = M*weights.unsqueeze(1)
# Iterative correction
lr = 0.75
# Initial value
knobs = torch.tensor((6 + 4)*[0.0], dtype=dtype)
# Correction loop
for _ in range(8):
value = observable(knobs)
residual = target - value
residual = torch.linalg.solve_triangular(L, residual.unsqueeze(-1), upper=False).squeeze(-1)
residual = residual*weights
delta = torch.linalg.lstsq(M, residual, driver="gels").solution
knobs += lr*delta
print(((value - target)**2).sum())
print()
tensor(64.2608, dtype=torch.float64)
tensor(3.6714, dtype=torch.float64)
tensor(0.3128, dtype=torch.float64)
tensor(0.0689, dtype=torch.float64)
tensor(0.0535, dtype=torch.float64)
tensor(0.0525, dtype=torch.float64)
tensor(0.0525, dtype=torch.float64)
tensor(0.0525, dtype=torch.float64)
[33]:
# Apply final corrections
kn, ks = torch.split(knobs, [6, 4])
result = error.clone()
result.flatten()
for name, knob in zip(nkn, kn):
result[name].kn = (result[name].kn + knob).item()
for name, knob in zip(nks, ks):
result[name].ks = (result[name].ks + knob).item()
result.splice()
[34]:
# Compute tunes (fractional part)
nux_result, nuy_result = tune(result, [], matched=True, limit=1)
[35]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_result, etapx_result, etaqy_result, etapy_result = dispersion(result, orbit, [], limit=1)
[36]:
# Compute twiss parameters
ax_result, bx_result, ay_result, by_result = twiss(result, [], matched=True, advance=True, full=False).T
[37]:
# Compute phase advances
mux_result, muy_result = advance(result, [], alignment=False, matched=True).T
[38]:
# Compute coupling
c_result = coupling(result, [])
[39]:
# Tune shifts
print((nux - nux_id).abs())
print((nuy - nuy_id).abs())
print()
print((nux - nux_result).abs())
print((nuy - nuy_result).abs())
print()
tensor(0.0260, dtype=torch.float64)
tensor(0.0114, dtype=torch.float64)
tensor(0.0031, dtype=torch.float64)
tensor(0.0106, dtype=torch.float64)
[40]:
# Coupling (minimal tune distance)
print(c)
print(c_id)
print(c_result)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
tensor(0.0009, dtype=torch.float64)
[41]:
# Dispersion
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_id).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_id).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_result).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_result).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
[42]:
# Beta-beating
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_id)/bx).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_id)/by).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_result)/bx).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_result)/by).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((bx - bx_id)/bx)**2).mean().sqrt())
print(100*(((by - by_id)/by)**2).mean().sqrt())
print()
print(100*(((bx - bx_result)/bx)**2).mean().sqrt())
print(100*(((by - by_result)/by)**2).mean().sqrt())
print()
tensor(11.5994, dtype=torch.float64)
tensor(1.7916, dtype=torch.float64)
tensor(0.0562, dtype=torch.float64)
tensor(0.3210, dtype=torch.float64)
[43]:
# Phase advance
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_id)/mux).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_id)/muy).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_result)/mux).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_result)/muy).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((mux - mux_id)/mux)**2).mean().sqrt())
print(100*(((muy - muy_id)/muy)**2).mean().sqrt())
print()
print(100*(((mux - mux_result)/mux)**2).mean().sqrt())
print(100*(((muy - muy_result)/muy)**2).mean().sqrt())
print()
tensor(8.7941, dtype=torch.float64)
tensor(1.7778, dtype=torch.float64)
tensor(0.3144, dtype=torch.float64)
tensor(0.3240, dtype=torch.float64)
[44]:
# From the above results, tunes are not corrected
# This can be solved by adding extra weights or by using global tune knobs (prefered)
[45]:
# Optimization loop (model free) with tune weights
# Responce matrix (jacobian)
M = matrix.clone()
# Weighting covariance (sensitivity) matrix
epsilon = 1.0E-9
C = M @ M.T
C = C + epsilon*torch.eye(len(C), dtype=dtype)
# Cholesky decomposition
L = torch.linalg.cholesky(C)
# Whiten response
M = torch.linalg.solve_triangular(L, M, upper=False)
# Additional weights
# Can be used to extra weight selected observables, e.g. tunes
weights = torch.ones(len(M), dtype=dtype)
weights = weights.sqrt()
# Tune weights (hand picked)
w_nux = 5.0
w_nuy = 8.0
weights[0] = w_nux
weights[1] = w_nuy
# Whiten response with additional weights
M = M*weights.unsqueeze(1)
# Iterative correction
lr = 0.75
# Initial value
knobs = torch.tensor((6 + 4)*[0.0], dtype=dtype)
# Correction loop
for _ in range(8):
value = observable(knobs)
residual = target - value
residual = torch.linalg.solve_triangular(L, residual.unsqueeze(-1), upper=False).squeeze(-1)
residual = residual*weights
delta = torch.linalg.lstsq(M, residual, driver="gels").solution
knobs += lr*delta
print(((value - target)**2).sum())
print()
tensor(64.2608, dtype=torch.float64)
tensor(28.4599, dtype=torch.float64)
tensor(23.3812, dtype=torch.float64)
tensor(22.3323, dtype=torch.float64)
tensor(22.1887, dtype=torch.float64)
tensor(22.1935, dtype=torch.float64)
tensor(22.2070, dtype=torch.float64)
tensor(22.2140, dtype=torch.float64)
[46]:
# Apply final corrections
kn, ks = torch.split(knobs, [6, 4])
result = error.clone()
result.flatten()
for name, knob in zip(nkn, kn):
result[name].kn = (result[name].kn + knob).item()
for name, knob in zip(nks, ks):
result[name].ks = (result[name].ks + knob).item()
result.splice()
[47]:
# Compute tunes (fractional part)
nux_result, nuy_result = tune(result, [], matched=True, limit=1)
[48]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_result, etapx_result, etaqy_result, etapy_result = dispersion(result, orbit, [], limit=1)
[49]:
# Compute twiss parameters
ax_result, bx_result, ay_result, by_result = twiss(result, [], matched=True, advance=True, full=False).T
[50]:
# Compute phase advances
mux_result, muy_result = advance(result, [], alignment=False, matched=True).T
[51]:
# Compute coupling
c_result = coupling(result, [])
[52]:
# Tune shifts
print((nux - nux_id).abs())
print((nuy - nuy_id).abs())
print()
print((nux - nux_result).abs())
print((nuy - nuy_result).abs())
print()
tensor(0.0260, dtype=torch.float64)
tensor(0.0114, dtype=torch.float64)
tensor(0.0001, dtype=torch.float64)
tensor(0.0001, dtype=torch.float64)
[53]:
# Coupling (minimal tune distance)
print(c)
print(c_id)
print(c_result)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
tensor(0.0009, dtype=torch.float64)
[54]:
# Dispersion
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_id).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_id).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_result).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_result).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
[55]:
# Beta-beating
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_id)/bx).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_id)/by).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_result)/bx).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_result)/by).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((bx - bx_id)/bx)**2).mean().sqrt())
print(100*(((by - by_id)/by)**2).mean().sqrt())
print()
print(100*(((bx - bx_result)/bx)**2).mean().sqrt())
print(100*(((by - by_result)/by)**2).mean().sqrt())
print()
tensor(11.5994, dtype=torch.float64)
tensor(1.7916, dtype=torch.float64)
tensor(1.2593, dtype=torch.float64)
tensor(4.2956, dtype=torch.float64)
[56]:
# Phase advance
plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_id)/mux).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_id)/muy).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_result)/mux).cpu().numpy(), fmt='-', marker='o', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_result)/muy).cpu().numpy(), fmt='-', marker='o', color='red', alpha=0.75)
plt.tight_layout()
plt.show()
print(100*(((mux - mux_id)/mux)**2).mean().sqrt())
print(100*(((muy - muy_id)/muy)**2).mean().sqrt())
print()
print(100*(((mux - mux_result)/mux)**2).mean().sqrt())
print(100*(((muy - muy_result)/muy)**2).mean().sqrt())
print()
tensor(8.7941, dtype=torch.float64)
tensor(1.7778, dtype=torch.float64)
tensor(1.0051, dtype=torch.float64)
tensor(4.1787, dtype=torch.float64)