ELETTRA-47: EU100 (local/global correction & short section IDs)
[1]:
# 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.tune import chromaticity
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
[2]:
# Set data type and device
Element.dtype = dtype = torch.float64
Element.device = device = torch.device('cpu')
[3]:
# Load lattice (ELEGANT table)
# Note, lattice is allowed to have repeated elements
path = Path('elettra.lte')
data = load_lattice(path)
[4]:
# 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'])
# 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
[4]:
{'BPM': 168, 'Drift': 720, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 24}
[5]:
# Compute tunes (fractional part)
nux, nuy = tune(ring, [], matched=True, limit=1)
[6]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx, etapx, etaqy, etapy = dispersion(ring, orbit, [], limit=1)
[7]:
# Compute twiss parameters
ax, bx, ay, by = twiss(ring, [], matched=True, advance=True, full=False).T
[8]:
# Compute phase advances
mux, muy = advance(ring, [], alignment=False, matched=True).T
[9]:
# Compute coupling
c = coupling(ring, [])
[10]:
# Compute chromaticity
psi = chromaticity(ring, [])
[11]:
# Quadrupole names for global tune correction
QF = [f'QF_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]
QD = [f'QD_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]
[12]:
# Global tune responce matrix
def global_observable(knobs):
kf, kd = knobs
kn = torch.stack(len(QF)*[kf] + len(QD)*[kd])
return tune(ring, [kn], ('kn', None, QF + QD, None), matched=True, limit=1)
knobs = torch.tensor([0.0, 0.0], dtype=dtype)
global_target = global_observable(knobs)
global_matrix = torch.func.jacfwd(global_observable)(knobs)
print(global_target)
print(global_matrix)
tensor([0.2994, 0.1608], dtype=torch.float64)
tensor([[ 5.8543, 2.0964],
[-2.9918, -1.2602]], dtype=torch.float64)
[13]:
# 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']
[14]:
# Define knobs to magnets mixing matrices
Sn = torch.tensor([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0], [0.0, 0.0, 1.0], [0.0, 1.0, 0.0], [1.0, 0.0, 0.0]], dtype=dtype)
print(Sn)
print()
Ss = torch.tensor([[+1.0, 0.0], [0.0, +1.0], [0.0, -1.0], [-1.0, 0.0]], dtype=dtype)
print(Ss)
print()
tensor([[1., 0., 0.],
[0., 1., 0.],
[0., 0., 1.],
[0., 0., 1.],
[0., 1., 0.],
[1., 0., 0.]], dtype=torch.float64)
tensor([[ 1., 0.],
[ 0., 1.],
[ 0., -1.],
[-1., 0.]], dtype=torch.float64)
[15]:
# Define twiss observable (full coupled twiss)
def observable_twiss(kn, ks):
return twiss(ring, [kn, ks], ('kn', None, nkn, None), ('ks', None, nks, None), matched=True, advance=True, full=False, convert=False)
[16]:
# 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
[17]:
# Construct full target observable vector and corresponding responce matrix
def observable(knobs):
kn, ks = torch.split(knobs, [3, 2])
kn = Sn @ kn
ks = Ss @ ks
betas = observable_twiss(kn, ks)
etas = observable_dispersion(kn, ks)
return torch.cat([betas.flatten(), etas.flatten()])
knobs = torch.tensor((3 + 2)*[0.0], dtype=dtype)
print((target := observable(knobs)).shape)
print((matrix := torch.func.jacfwd(observable)(knobs)).shape)
torch.Size([5712])
torch.Size([5712, 5])
[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]:
# W96_S05
ARK = torch.tensor([[-0.000248802, 0.0, 0.0, 0.0],
[ 0.0, -9.53447e-6, 0.0, 0.0],
[ 0.0, 0.0, 0.0176392,0.0],
[ 0.0, 0.0, 0.0, 0.000856368]], dtype=dtype)
W96_POS = 0.0
W96_S05 = Matrix('W96_S05', length=0.0, A=ARK[torch.triu(torch.ones_like(ARK, dtype=torch.bool))].tolist())
[20]:
# W96_S06
ARK = torch.tensor([[-0.000248802, 0.0, 0.0, 0.0],
[ 0.0, -9.53447e-6, 0.0, 0.0],
[ 0.0, 0.0, 0.0176392,0.0],
[ 0.0, 0.0, 0.0, 0.000856368]], dtype=dtype)
W96_POS = 0.0
W96_S06 = Matrix('W96_S06', length=0.0, A=ARK[torch.triu(torch.ones_like(ARK, dtype=torch.bool))].tolist())
[21]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(W96_S05, error.next('MSS_S05').name, position=W96_POS*1.0E-3)
error.splice()
error.describe
[21]:
{'BPM': 168,
'Drift': 721,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 1}
[22]:
# Compute tunes (fractional part)
nux_id_a, nuy_id_a = tune(error, [], matched=True, limit=1)
[23]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_a, etapx_id_a, etaqy_id_a, etapy_id_a = dispersion(error, orbit, [], limit=1)
[24]:
# Compute twiss parameters
ax_id_a, bx_id_a, ay_id_a, by_id_a = twiss(error, [], matched=True, advance=True, full=False).T
[25]:
# Compute phase advances
mux_id_a, muy_id_a = advance(error, [], alignment=False, matched=True).T
[26]:
# Compute coupling
c_id_a = coupling(error, [])
[27]:
# Compute hromaticity
psi_id_a = chromaticity(error, [])
[28]:
# Tune shifts
print((nux - nux_id_a))
print((nuy - nuy_id_a))
tensor(7.8934e-05, dtype=torch.float64)
tensor(-0.0028, dtype=torch.float64)
[29]:
# Coupling (minimal tune distance)
print(c)
print(c_id_a)
tensor(0., dtype=torch.float64)
tensor(0., dtype=torch.float64)
[30]:
# Chromaticity
print(psi)
print(psi_id_a)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.2073, -66.6162], dtype=torch.float64)
[31]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=0.0)
error.splice()
error.describe
[31]:
{'BPM': 168,
'Drift': 721,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 1}
[32]:
# Compute tunes (fractional part)
nux_id_b, nuy_id_b = tune(error, [], matched=True, limit=1)
[33]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_b, etapx_id_b, etaqy_id_b, etapy_id_b = dispersion(error, orbit, [], limit=1)
[34]:
# Compute twiss parameters
ax_id_b, bx_id_b, ay_id_b, by_id_b = twiss(error, [], matched=True, advance=True, full=False).T
[35]:
# Compute phase advances
mux_id_b, muy_id_b = advance(error, [], alignment=False, matched=True).T
[36]:
# Compute coupling
c_id_b = coupling(error, [])
[37]:
# Compute hromaticity
psi_id_b = chromaticity(error, [])
[38]:
# Tune shifts
print((nux - nux_id_b))
print((nuy - nuy_id_b))
tensor(0.0260, dtype=torch.float64)
tensor(-0.0114, dtype=torch.float64)
[39]:
# Coupling (minimal tune distance)
print(c)
print(c_id_b)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
[40]:
# Chromaticity
print(psi)
print(psi_id_b)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-72.2597, -66.6681], dtype=torch.float64)
[41]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=0.0)
error.insert(W96_S05, error.next('MSS_S05').name, position=W96_POS*1.0E-3)
error.splice()
error.describe
[41]:
{'BPM': 168,
'Drift': 722,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 2}
[42]:
# Compute tunes (fractional part)
nux_id_c, nuy_id_c = tune(error, [], matched=True, limit=1)
[43]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_c, etapx_id_c, etaqy_id_c, etapy_id_c = dispersion(error, orbit, [], limit=1)
[44]:
# Compute twiss parameters
ax_id_c, bx_id_c, ay_id_c, by_id_c = twiss(error, [], matched=True, advance=True, full=False).T
[45]:
# Compute phase advances
mux_id_c, muy_id_c = advance(error, [], alignment=False, matched=True).T
[46]:
# Compute coupling
c_id_c = coupling(error, [])
[47]:
# Compute hromaticity
psi_id_c = chromaticity(error, [])
[48]:
# Tune shifts
print((nux - nux_id_c))
print((nuy - nuy_id_c))
tensor(0.0260, dtype=torch.float64)
tensor(-0.0142, dtype=torch.float64)
[49]:
# Coupling (minimal tune distance)
print(c)
print(c_id_c)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
[50]:
# Chromaticity
print(psi)
print(psi_id_c)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-72.2592, -66.6015], dtype=torch.float64)
[51]:
# Chromaticity
(psi_x, psi_y) = psi
(psi_id_a_x, psi_id_a_y) = psi_id_a
(psi_id_b_x, psi_id_b_y) = psi_id_b
(psi_id_c_x, psi_id_c_y) = psi_id_c
print(psi_x, psi_y)
print(psi_id_a_x, psi_id_a_y)
print(psi_id_b_x, psi_id_b_y)
print(psi_id_c_x, psi_id_c_y)
tensor(-71.2093, dtype=torch.float64) tensor(-66.6787, dtype=torch.float64)
tensor(-71.2073, dtype=torch.float64) tensor(-66.6162, dtype=torch.float64)
tensor(-72.2597, dtype=torch.float64) tensor(-66.6681, dtype=torch.float64)
tensor(-72.2592, dtype=torch.float64) tensor(-66.6015, dtype=torch.float64)
[52]:
# Beta-beating
bx_a_bb = 100.0*(bx - bx_id_a) / bx
by_a_bb = 100.0*(by - by_id_a) / by
bx_b_bb = 100.0*(bx - bx_id_b) / bx
by_b_bb = 100.0*(by - by_id_b) / by
bx_c_bb = 100.0*(bx - bx_id_c) / bx
by_c_bb = 100.0*(by - by_id_c) / by
def rms(x):
return (x**2).mean().sqrt()
rms_x_a = rms(bx_a_bb).item()
ptp_x_a = (bx_a_bb.max() - bx_a_bb.min()).item()
rms_y_a = rms(by_a_bb).item()
ptp_y_a = (by_a_bb.max() - by_a_bb.min()).item()
rms_x_b = rms(bx_b_bb).item()
ptp_x_b = (bx_b_bb.max() - bx_b_bb.min()).item()
rms_y_b = rms(by_b_bb).item()
ptp_y_b = (by_b_bb.max() - by_b_bb.min()).item()
rms_x_c = rms(bx_c_bb).item()
ptp_x_c = (bx_c_bb.max() - bx_c_bb.min()).item()
rms_y_c = rms(by_c_bb).item()
ptp_y_c = (by_c_bb.max() - by_c_bb.min()).item()
s = ring.locations().cpu().numpy()
bx_a_np = bx_a_bb.cpu().numpy()
by_a_np = by_a_bb.cpu().numpy()
bx_b_np = bx_b_bb.cpu().numpy()
by_b_np = by_b_bb.cpu().numpy()
bx_c_np = bx_c_bb.cpu().numpy()
by_c_np = by_c_bb.cpu().numpy()
fig, ax = plt.subplots(
1, 1, figsize=(16, 6),
sharex=True,
gridspec_kw={'hspace': 0.3}
)
fig.subplots_adjust(top=0.85, bottom=0.35)
ax.errorbar(s, bx_a_np, fmt='-', marker='x', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (W96)')
ax.errorbar(s, by_a_np, fmt='-', marker='x', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (W96)')
ax.errorbar(s, bx_b_np, fmt='-', marker='o', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100)')
ax.errorbar(s, by_b_np, fmt='-', marker='o', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100)')
ax.errorbar(s, bx_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100 \& W96)')
ax.errorbar(s, by_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100 \& W96)')
ax.set_xlabel('s [m]', fontsize=18)
ax.set_ylabel(r'$\Delta \beta / \beta$ [\%]', fontsize=18)
ax.tick_params(width=2, labelsize=16)
ax.tick_params(axis='x', length=8, direction='in')
ax.tick_params(axis='y', length=8, direction='in')
title = (
r'W96~~~~~~~~~~~~~:~'
rf'RMS$_x$={rms_x_a:05.2f}\% \quad RMS$_y$={rms_y_a:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_a:05.2f}\% \quad PTP$_y$={ptp_y_a:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_a)}{(nux - nux_id_a).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_a)}{(nuy - nuy_id_a).abs().item():.4f}'
rf'\quad C={c_id_a.item():.6f}'
)
ax.text(0.0, 1.25, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
r'EU100~~~~~~~~~~~:~'
rf'RMS$_x$={rms_x_b:05.2f}\% \quad RMS$_y$={rms_y_b:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_b:05.2f}\% \quad PTP$_y$={ptp_y_b:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_b)}{(nux - nux_id_b).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_b)}{(nuy - nuy_id_b).abs().item():.4f}'
rf'\quad C={c_id_b.item():.6f}'
)
ax.text(0.0, 1.15, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
r'EU100 \& W96~:~'
rf'RMS$_x$={rms_x_c:05.2f}\% \quad RMS$_y$={rms_y_c:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_c:05.2f}\% \quad PTP$_y$={ptp_y_c:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_c)}{(nux - nux_id_c).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_c)}{(nuy - nuy_id_c).abs().item():.4f}'
rf'\quad C={c_id_c.item():.6f}'
)
ax.text(0.0, 1.05, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
error.flatten()
for position in error.locations()[[error.position(name) for name in ['ID', 'W96_S05']]]:
ax.axvline(position, -1, 1, linestyle='dashed', color='black')
error.splice()
ax.legend(loc='upper right', frameon=False, fontsize=14, ncol=6)
ax.set_ylim(-20, 20)
plt.setp(ax.spines.values(), linewidth=2.0)
plt.show()
[53]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=0.0)
error.splice()
error.describe
[53]:
{'BPM': 168,
'Drift': 721,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 1}
[54]:
# Compute tunes (fractional part)
nux_id, nuy_id = tune(error, [], matched=True, limit=1)
[55]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id, etapx_id, etaqy_id, etapy_id = dispersion(error, orbit, [], limit=1)
[56]:
# Compute twiss parameters
ax_id, bx_id, ay_id, by_id = twiss(error, [], matched=True, advance=True, full=False).T
[57]:
# Compute phase advances
mux_id, muy_id = advance(error, [], alignment=False, matched=True).T
[58]:
# Compute coupling
c_id = coupling(error, [])
[59]:
# Compute hromaticity
psi_id = chromaticity(error, [])
[60]:
# Tune shifts
print((nux - nux_id))
print((nuy - nuy_id))
tensor(0.0260, dtype=torch.float64)
tensor(-0.0114, dtype=torch.float64)
[61]:
# Coupling (minimal tune distance)
print(c)
print(c_id)
tensor(0., dtype=torch.float64)
tensor(0.0004, dtype=torch.float64)
[62]:
# Chromaticity
print(psi)
print(psi_id)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-72.2597, -66.6681], dtype=torch.float64)
[63]:
# Define parametric observable vector
def global_observable(knobs):
kf, kd = knobs
kn = torch.stack(len(QF)*[kf] + len(QD)*[kd])
return tune(error, [kn], ('kn', None, QF + QD, None), matched=True, limit=1)
def observable_twiss(kn, ks):
return twiss(error, [kn, ks], ('kn', None, nkn, None), ('ks', None, nks, None), matched=True, advance=True, full=False, convert=False)
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, [3, 2])
kn = Sn @ kn
ks = Ss @ ks
betas = observable_twiss(kn, ks)
etas = observable_dispersion(kn, ks)
return torch.cat([betas.flatten(), etas.flatten()])
[64]:
# Check the residual vector norm
global_knobs = torch.tensor(2*[0.0], dtype=dtype)
knobs = torch.tensor((3 + 2)*[0.0], dtype=dtype)
print(((global_observable(global_knobs) - global_target)**2).sum())
print(((observable(knobs) - target)**2).sum())
tensor(0.0008, dtype=torch.float64)
tensor(212.9162, dtype=torch.float64)
[65]:
# Optimization loop (local)
# 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((3 + 2)*[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(212.9162, dtype=torch.float64)
tensor(12.5009, dtype=torch.float64)
tensor(1.1080, dtype=torch.float64)
tensor(0.3135, dtype=torch.float64)
tensor(0.2637, dtype=torch.float64)
tensor(0.2607, dtype=torch.float64)
tensor(0.2605, dtype=torch.float64)
tensor(0.2605, dtype=torch.float64)
[66]:
# Knob values
kn, ks = torch.split(knobs, [3, 2])
kn = Sn @ kn
ks = Ss @ ks
print(kn.numpy())
print(ks.numpy())
[-0.06037695 -0.00614317 0.20562899 0.20562899 -0.00614317 -0.06037695]
[-0.00051896 -0.00251065 0.00251065 0.00051896]
[67]:
# Apply final corrections
error.flatten()
for name, knob in zip(nkn, kn):
error[name].kn = (error[name].kn + knob).item()
for name, knob in zip(nks, ks):
error[name].ks = (error[name].ks + knob).item()
error.splice()
[68]:
# Check
print(((global_observable(0.0*global_knobs) - global_target)**2).sum())
print(((observable(0.0*knobs) - target)**2).sum())
tensor(0.0001, dtype=torch.float64)
tensor(0.2605, dtype=torch.float64)
[69]:
# Compute tunes (fractional part)
nux_id_a, nuy_id_a = tune(error, [], matched=True, limit=1)
[70]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_a, etapx_id_a, etaqy_id_a, etapy_id_a = dispersion(error, orbit, [], limit=1)
[71]:
# Compute twiss parameters
ax_id_a, bx_id_a, ay_id_a, by_id_a = twiss(error, [], matched=True, advance=True, full=False).T
[72]:
# Compute phase advances
mux_id_a, muy_id_a = advance(error, [], alignment=False, matched=True).T
[73]:
# Compute coupling
c_id_a = coupling(error, [])
[74]:
# Compute hromaticity
psi_id_a = chromaticity(error, [])
[75]:
# Tune shifts
print((nux - nux_id_a))
print((nuy - nuy_id_a))
tensor(-0.0031, dtype=torch.float64)
tensor(-0.0106, dtype=torch.float64)
[76]:
# Coupling (minimal tune distance)
print(c)
print(c_id_a)
tensor(0., dtype=torch.float64)
tensor(6.0686e-06, dtype=torch.float64)
[77]:
# Chromaticity
print(psi)
print(psi_id_a)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1952, -66.7068], dtype=torch.float64)
[78]:
# Optimization loop (global)
# Responce matrix (jacobian)
M = global_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
global_knobs = torch.tensor(2*[0.0], dtype=dtype)
# Correction loop
for _ in range(1 + 1):
value = global_observable(global_knobs)
residual = global_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
global_knobs += lr*delta
print(((value - global_target)**2).sum())
print()
tensor(0.0001, dtype=torch.float64)
tensor(9.0535e-06, dtype=torch.float64)
[79]:
# Knob values
kf, kd = global_knobs
print(kd.numpy())
print(kf.numpy())
0.061527614111642676
-0.022554709001744343
[80]:
# Apply final corrections
error.flatten()
for name in QD:
error[name].kn = (error[name].kn + kd).item()
for name in QF:
error[name].kn = (error[name].kn + kf).item()
error.splice()
[81]:
# Compute tunes (fractional part)
nux_id_b, nuy_id_b = tune(error, [], matched=True, limit=1)
[82]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_b, etapx_id_b, etaqy_id_b, etapy_id_b = dispersion(error, orbit, [], limit=1)
[83]:
# Compute twiss parameters
ax_id_b, bx_id_b, ay_id_b, by_id_b = twiss(error, [], matched=True, advance=True, full=False).T
[84]:
# Compute phase advances
mux_id_b, muy_id_b = advance(error, [], alignment=False, matched=True).T
[85]:
# Compute coupling
c_id_b = coupling(error, [])
[86]:
# Compute chromaticity
psi_id_b = chromaticity(error, [])
[87]:
# Tune shifts
print((nux - nux_id_b))
print((nuy - nuy_id_b))
tensor(-0.0004, dtype=torch.float64)
tensor(-0.0008, dtype=torch.float64)
[88]:
# Coupling (minimal tune distance)
print(c)
print(c_id_b)
tensor(0., dtype=torch.float64)
tensor(6.5885e-06, dtype=torch.float64)
[89]:
# Chromaticity
print(psi)
print(psi_id_b)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1487, -66.7459], dtype=torch.float64)
[90]:
# Beta-beating
bx_a_bb = 100.0*(bx - bx_id_a) / bx
by_a_bb = 100.0*(by - by_id_a) / by
bx_b_bb = 100.0*(bx - bx_id_b) / bx
by_b_bb = 100.0*(by - by_id_b) / by
def rms(x):
return (x**2).mean().sqrt()
rms_x_a = rms(bx_a_bb).item()
ptp_x_a = (bx_a_bb.max() - bx_a_bb.min()).item()
rms_y_a = rms(by_a_bb).item()
ptp_y_a = (by_a_bb.max() - by_a_bb.min()).item()
rms_x_b = rms(bx_b_bb).item()
ptp_x_b = (bx_b_bb.max() - bx_b_bb.min()).item()
rms_y_b = rms(by_b_bb).item()
ptp_y_b = (by_b_bb.max() - by_b_bb.min()).item()
s = ring.locations().cpu().numpy()
bx_a_np = bx_a_bb.cpu().numpy()
by_a_np = by_a_bb.cpu().numpy()
bx_b_np = bx_b_bb.cpu().numpy()
by_b_np = by_b_bb.cpu().numpy()
fig, ax = plt.subplots(
1, 1, figsize=(16, 6),
sharex=True,
gridspec_kw={'hspace': 0.3}
)
fig.subplots_adjust(top=0.85, bottom=0.35)
ax.errorbar(s, bx_a_np, fmt='-', marker='x', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (local)')
ax.errorbar(s, by_a_np, fmt='-', marker='x', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (local)')
ax.errorbar(s, bx_b_np, fmt='-', marker='o', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (local + global)')
ax.errorbar(s, by_b_np, fmt='-', marker='o', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (local + global)')
ax.set_xlabel('s [m]', fontsize=18)
ax.set_ylabel(r'$\Delta \beta / \beta$ [\%]', fontsize=18)
ax.tick_params(width=2, labelsize=16)
ax.tick_params(axis='x', length=8, direction='in')
ax.tick_params(axis='y', length=8, direction='in')
title = (
rf'RMS$_x$={rms_x_a:05.2f}\% \quad RMS$_y$={rms_y_a:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_a:05.2f}\% \quad PTP$_y$={ptp_y_a:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_a)}{(nux - nux_id_a).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_a)}{(nuy - nuy_id_a).abs().item():.4f}'
rf'\quad C={c_id_a.item():.6f}'
)
ax.text(0.0, 1.15, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
rf'RMS$_x$={rms_x_b:05.2f}\% \quad RMS$_y$={rms_y_b:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_b:05.2f}\% \quad PTP$_y$={ptp_y_b:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_b)}{(nux - nux_id_b).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_b)}{(nuy - nuy_id_b).abs().item():.4f}'
rf'\quad C={c_id_b.item():.6f}'
)
ax.text(0.0, 1.05, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
error.flatten()
for position in error.locations()[[error.position(name) for name in ['ID']]]:
ax.axvline(position, -1, 1, linestyle='dashed', color='black')
error.splice()
ax.legend(loc='upper right', frameon=False, fontsize=14, ncol=6)
ax.set_ylim(-3, 5)
plt.setp(ax.spines.values(), linewidth=2.0)
plt.show()
[91]:
# Correction
QF = [f'QF_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]
QD = [f'QD_S{i:02}_{j:02}' for j in [2, 3] for i in range(1, 12 + 1)]
nkn = ['OCT_S01_02', 'QF_S01_02', 'QD_S01_02', 'QD_S01_03', 'QF_S01_03', 'OCT_S01_03']
nks = ['SD_S01_05', 'SH_S01_02', 'SH_S01_03', 'SD_S01_06']
ring.flatten()
error.flatten()
dkn = {name: (error[name].kn - ring[name].kn).item() for name in nkn}
dks = {name: (error[name].ks - ring[name].ks).item() for name in nks}
dkf = {name: (error[name].kn - ring[name].kn).item() for name in QF}
dkd = {name: (error[name].kn - ring[name].kn).item() for name in QD}
ring.splice()
error.splice()
[92]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=0.0)
error.insert(W96_S05, error.next('MSS_S05').name, position=W96_POS*1.0E-3)
error.splice()
error.describe
[92]:
{'BPM': 168,
'Drift': 722,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 2}
[93]:
# Apply correction
error.flatten()
for name in nkn:
error[name].kn = error[name].kn.item() + dkn[name]
for name in nks:
error[name].ks = error[name].ks.item() + dks[name]
for name in QD:
if name not in nkn:
error[name].kn = error[name].kn.item() + dkd[name]
for name in QF:
if name not in nkn:
error[name].kn = error[name].kn.item() + dkf[name]
error.splice()
[94]:
# Compute tunes (fractional part)
nux_id_b, nuy_id_b = tune(error, [], matched=True, limit=1)
[95]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_b, etapx_id_b, etaqy_id_b, etapy_id_b = dispersion(error, orbit, [], limit=1)
[96]:
# Compute twiss parameters
ax_id_b, bx_id_b, ay_id_b, by_id_b = twiss(error, [], matched=True, advance=True, full=False).T
[97]:
# Compute phase advances
mux_id_b, muy_id_b = advance(error, [], alignment=False, matched=True).T
[98]:
# Compute coupling
c_id_b = coupling(error, [])
[99]:
# Compute chromaticity
psi_id_b = chromaticity(error, [])
[100]:
# Tune shifts
print((nux - nux_id_b))
print((nuy - nuy_id_b))
tensor(-0.0003, dtype=torch.float64)
tensor(-0.0036, dtype=torch.float64)
[101]:
# Coupling (minimal tune distance)
print(c)
print(c_id_b)
tensor(0., dtype=torch.float64)
tensor(6.3666e-06, dtype=torch.float64)
[102]:
# Chromaticity
print(psi)
print(psi_id_b)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1466, -66.6850], dtype=torch.float64)
[103]:
# Optimization loop (global)
# Responce matrix (jacobian)
M = global_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
global_knobs = torch.tensor(2*[0.0], dtype=dtype)
# Correction loop
for _ in range(1 + 1):
value = global_observable(global_knobs)
residual = global_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
global_knobs += lr*delta
print(((value - global_target)**2).sum())
print()
tensor(1.3087e-05, dtype=torch.float64)
tensor(1.1039e-06, dtype=torch.float64)
[104]:
# Knob values
kf, kd = global_knobs
print(kd.numpy())
print(kf.numpy())
0.019468586966671107
-0.007040292415693757
[105]:
# Apply final corrections
error.flatten()
for name in QD:
error[name].kn = (error[name].kn + kd).item()
for name in QF:
error[name].kn = (error[name].kn + kf).item()
error.splice()
[106]:
# Compute tunes (fractional part)
nux_id_c, nuy_id_c = tune(error, [], matched=True, limit=1)
[107]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_c, etapx_id_c, etaqy_id_c, etapy_id_c = dispersion(error, orbit, [], limit=1)
[108]:
# Compute twiss parameters
ax_id_c, bx_id_c, ay_id_c, by_id_c = twiss(error, [], matched=True, advance=True, full=False).T
[109]:
# Compute phase advances
mux_id_c, muy_id_c = advance(error, [], alignment=False, matched=True).T
[110]:
# Compute coupling
c_id_c = coupling(error, [])
[111]:
# Compute hromaticity
psi_id_c = chromaticity(error, [])
[112]:
# Tune shifts
print((nux - nux_id_c))
print((nuy - nuy_id_c))
tensor(-0.0001, dtype=torch.float64)
tensor(-0.0003, dtype=torch.float64)
[113]:
# Coupling (minimal tune distance)
print(c)
print(c_id_c)
tensor(0., dtype=torch.float64)
tensor(6.6028e-06, dtype=torch.float64)
[114]:
# Chromaticity
print(psi)
print(psi_id_c)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1294, -66.6959], dtype=torch.float64)
[115]:
# Chromaticity
(psi_x, psi_y) = psi
(psi_id_b_x, psi_id_b_y) = psi_id_b
(psi_id_c_x, psi_id_c_y) = psi_id_c
print(psi_x, psi_y)
print(psi_id_b_x, psi_id_b_y)
print(psi_id_c_x, psi_id_c_y)
tensor(-71.2093, dtype=torch.float64) tensor(-66.6787, dtype=torch.float64)
tensor(-71.1466, dtype=torch.float64) tensor(-66.6850, dtype=torch.float64)
tensor(-71.1294, dtype=torch.float64) tensor(-66.6959, dtype=torch.float64)
[116]:
# Beta-beating
bx_b_bb = 100.0*(bx - bx_id_b) / bx
by_b_bb = 100.0*(by - by_id_b) / by
bx_c_bb = 100.0*(bx - bx_id_c) / bx
by_c_bb = 100.0*(by - by_id_c) / by
def rms(x):
return (x**2).mean().sqrt()
rms_x_b = rms(bx_b_bb).item()
ptp_x_b = (bx_b_bb.max() - bx_b_bb.min()).item()
rms_y_b = rms(by_b_bb).item()
ptp_y_b = (by_b_bb.max() - by_b_bb.min()).item()
rms_x_c = rms(bx_c_bb).item()
ptp_x_c = (bx_c_bb.max() - bx_c_bb.min()).item()
rms_y_c = rms(by_c_bb).item()
ptp_y_c = (by_c_bb.max() - by_c_bb.min()).item()
s = ring.locations().cpu().numpy()
bx_b_np = bx_b_bb.cpu().numpy()
by_b_np = by_b_bb.cpu().numpy()
bx_c_np = bx_c_bb.cpu().numpy()
by_c_np = by_c_bb.cpu().numpy()
fig, ax = plt.subplots(
1, 1, figsize=(16, 6),
sharex=True,
gridspec_kw={'hspace': 0.3}
)
fig.subplots_adjust(top=0.85, bottom=0.35)
ax.errorbar(s, bx_b_np, fmt='-', marker='o', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100)')
ax.errorbar(s, by_b_np, fmt='-', marker='o', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100)')
ax.errorbar(s, bx_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100 \& TUNE)')
ax.errorbar(s, by_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100 \& TUNE)')
ax.set_xlabel('s [m]', fontsize=18)
ax.set_ylabel(r'$\Delta \beta / \beta$ [\%]', fontsize=18)
ax.tick_params(width=2, labelsize=16)
ax.tick_params(axis='x', length=8, direction='in')
ax.tick_params(axis='y', length=8, direction='in')
title = (
rf'RMS$_x$={rms_x_b:05.2f}\% \quad RMS$_y$={rms_y_b:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_b:05.2f}\% \quad PTP$_y$={ptp_y_b:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_b)}{(nux - nux_id_b).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_b)}{(nuy - nuy_id_b).abs().item():.4f}'
rf'\quad C={c_id_b.item():.6f}'
)
ax.text(0.0, 1.15, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
rf'RMS$_x$={rms_x_c:05.2f}\% \quad RMS$_y$={rms_y_c:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_c:05.2f}\% \quad PTP$_y$={ptp_y_c:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_c)}{(nux - nux_id_c).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_c)}{(nuy - nuy_id_c).abs().item():.4f}'
rf'\quad C={c_id_c.item():.6f}'
)
ax.text(0.0, 1.05, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
error.flatten()
for position in error.locations()[[error.position(name) for name in ['ID', 'W96_S05']]]:
ax.axvline(position, -1, 1, linestyle='dashed', color='black')
error.splice()
ax.legend(loc='upper right', frameon=False, fontsize=14, ncol=6)
ax.set_ylim(-3, 5)
plt.setp(ax.spines.values(), linewidth=2.0)
plt.show()
[117]:
# Insert ID(s)
error = ring.clone()
error.flatten()
error.insert(ID, error.next('MLL_S01').name, position=0.0)
error.insert(W96_S05, error.next('MSS_S05').name, position=W96_POS*1.0E-3)
error.insert(W96_S06, error.next('MSS_S06').name, position=W96_POS*1.0E-3)
error.splice()
error.describe
[117]:
{'BPM': 168,
'Drift': 723,
'Dipole': 156,
'Quadrupole': 360,
'Marker': 24,
'Matrix': 3}
[118]:
# Apply correction
error.flatten()
for name in nkn:
error[name].kn = error[name].kn.item() + dkn[name]
for name in nks:
error[name].ks = error[name].ks.item() + dks[name]
for name in QD:
if name not in nkn:
error[name].kn = error[name].kn.item() + dkd[name]
for name in QF:
if name not in nkn:
error[name].kn = error[name].kn.item() + dkf[name]
error.splice()
[119]:
# Compute tunes (fractional part)
nux_id_b, nuy_id_b = tune(error, [], matched=True, limit=1)
[120]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_b, etapx_id_b, etaqy_id_b, etapy_id_b = dispersion(error, orbit, [], limit=1)
[121]:
# Compute twiss parameters
ax_id_b, bx_id_b, ay_id_b, by_id_b = twiss(error, [], matched=True, advance=True, full=False).T
[122]:
# Compute phase advances
mux_id_b, muy_id_b = advance(error, [], alignment=False, matched=True).T
[123]:
# Compute coupling
c_id_b = coupling(error, [])
[124]:
# Compute chromaticity
psi_id_b = chromaticity(error, [])
[125]:
# Tune shifts
print((nux - nux_id_b))
print((nuy - nuy_id_b))
tensor(-0.0002, dtype=torch.float64)
tensor(-0.0064, dtype=torch.float64)
[126]:
# Coupling (minimal tune distance)
print(c)
print(c_id_b)
tensor(0., dtype=torch.float64)
tensor(6.1968e-06, dtype=torch.float64)
[127]:
# Chromaticity
print(psi)
print(psi_id_b)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1445, -66.5964], dtype=torch.float64)
[128]:
# Optimization loop (global)
# Responce matrix (jacobian)
M = global_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
global_knobs = torch.tensor(2*[0.0], dtype=dtype)
# Correction loop
for _ in range(1 + 1):
value = global_observable(global_knobs)
residual = global_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
global_knobs += lr*delta
print(((value - global_target)**2).sum())
print()
tensor(4.1593e-05, dtype=torch.float64)
tensor(3.5005e-06, dtype=torch.float64)
[129]:
# Knob values
kf, kd = global_knobs
print(kd.numpy())
print(kf.numpy())
0.03407559744334108
-0.012277688972451763
[130]:
# Apply final corrections
error.flatten()
for name in QD:
error[name].kn = (error[name].kn + kd).item()
for name in QF:
error[name].kn = (error[name].kn + kf).item()
error.splice()
[131]:
# Compute tunes (fractional part)
nux_id_c, nuy_id_c = tune(error, [], matched=True, limit=1)
[132]:
# Compute dispersion
orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id_c, etapx_id_c, etaqy_id_c, etapy_id_c = dispersion(error, orbit, [], limit=1)
[133]:
# Compute twiss parameters
ax_id_c, bx_id_c, ay_id_c, by_id_c = twiss(error, [], matched=True, advance=True, full=False).T
[134]:
# Compute phase advances
mux_id_c, muy_id_c = advance(error, [], alignment=False, matched=True).T
[135]:
# Compute coupling
c_id_c = coupling(error, [])
[136]:
# Compute hromaticity
psi_id_c = chromaticity(error, [])
[137]:
# Tune shifts
print((nux - nux_id_c))
print((nuy - nuy_id_c))
tensor(-0.0002, dtype=torch.float64)
tensor(-0.0005, dtype=torch.float64)
[138]:
# Coupling (minimal tune distance)
print(c)
print(c_id_c)
tensor(0., dtype=torch.float64)
tensor(6.6421e-06, dtype=torch.float64)
[139]:
# Chromaticity
print(psi)
print(psi_id_c)
tensor([-71.2093, -66.6787], dtype=torch.float64)
tensor([-71.1133, -66.6140], dtype=torch.float64)
[140]:
# Chromaticity
(psi_x, psi_y) = psi
(psi_id_b_x, psi_id_b_y) = psi_id_b
(psi_id_c_x, psi_id_c_y) = psi_id_c
print(psi_x, psi_y)
print(psi_id_b_x, psi_id_b_y)
print(psi_id_c_x, psi_id_c_y)
tensor(-71.2093, dtype=torch.float64) tensor(-66.6787, dtype=torch.float64)
tensor(-71.1445, dtype=torch.float64) tensor(-66.5964, dtype=torch.float64)
tensor(-71.1133, dtype=torch.float64) tensor(-66.6140, dtype=torch.float64)
[141]:
# Beta-beating
bx_b_bb = 100.0*(bx - bx_id_b) / bx
by_b_bb = 100.0*(by - by_id_b) / by
bx_c_bb = 100.0*(bx - bx_id_c) / bx
by_c_bb = 100.0*(by - by_id_c) / by
def rms(x):
return (x**2).mean().sqrt()
rms_x_b = rms(bx_b_bb).item()
ptp_x_b = (bx_b_bb.max() - bx_b_bb.min()).item()
rms_y_b = rms(by_b_bb).item()
ptp_y_b = (by_b_bb.max() - by_b_bb.min()).item()
rms_x_c = rms(bx_c_bb).item()
ptp_x_c = (bx_c_bb.max() - bx_c_bb.min()).item()
rms_y_c = rms(by_c_bb).item()
ptp_y_c = (by_c_bb.max() - by_c_bb.min()).item()
s = ring.locations().cpu().numpy()
bx_b_np = bx_b_bb.cpu().numpy()
by_b_np = by_b_bb.cpu().numpy()
bx_c_np = bx_c_bb.cpu().numpy()
by_c_np = by_c_bb.cpu().numpy()
fig, ax = plt.subplots(
1, 1, figsize=(16, 6),
sharex=True,
gridspec_kw={'hspace': 0.3}
)
fig.subplots_adjust(top=0.85, bottom=0.35)
ax.errorbar(s, bx_b_np, fmt='-', marker='o', ms=4, color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100)')
ax.errorbar(s, by_b_np, fmt='-', marker='o', ms=4, color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100)')
ax.errorbar(s, bx_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='blue', alpha=0.75, lw=1.5, label=r'$\beta_x$ (EU100 \& TUNE)')
ax.errorbar(s, by_c_np, fmt='-', marker='s', ms=6, markerfacecolor='none', color='red', alpha=0.75, lw=1.5, label=r'$\beta_y$ (EU100 \& TUNE)')
ax.set_xlabel('s [m]', fontsize=18)
ax.set_ylabel(r'$\Delta \beta / \beta$ [\%]', fontsize=18)
ax.tick_params(width=2, labelsize=16)
ax.tick_params(axis='x', length=8, direction='in')
ax.tick_params(axis='y', length=8, direction='in')
title = (
rf'RMS$_x$={rms_x_b:05.2f}\% \quad RMS$_y$={rms_y_b:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_b:05.2f}\% \quad PTP$_y$={ptp_y_b:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_b)}{(nux - nux_id_b).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_b)}{(nuy - nuy_id_b).abs().item():.4f}'
rf'\quad C={c_id_b.item():.6f}'
)
ax.text(0.0, 1.15, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
rf'RMS$_x$={rms_x_c:05.2f}\% \quad RMS$_y$={rms_y_c:05.2f}\% \quad '
rf'PTP$_x$={ptp_x_c:05.2f}\% \quad PTP$_y$={ptp_y_c:05.2f}\% \quad '
rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id_c)}{(nux - nux_id_c).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id_c)}{(nuy - nuy_id_c).abs().item():.4f}'
rf'\quad C={c_id_c.item():.6f}'
)
ax.text(0.0, 1.05, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
error.flatten()
for position in error.locations()[[error.position(name) for name in ['ID', 'W96_S05']]]:
ax.axvline(position, -1, 1, linestyle='dashed', color='black')
error.splice()
ax.legend(loc='upper right', frameon=False, fontsize=14, ncol=6)
ax.set_ylim(-3, 5)
plt.setp(ax.spines.values(), linewidth=2.0)
plt.show()
