ELETTRA-33: ID linear optics distortion (short wiggler)

# In this example effects of a short ID represented by a linear 4x4 symplectic matrix are presented
# ID is inserted replacing a selected marker as a thin insertion

# 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

# 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 or MSS (long and short 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

{'BPM': 168, 'Drift': 720, 'Dipole': 156, 'Quadrupole': 360, 'Marker': 24}

# Compute tunes (fractional part)

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

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
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, [])

# Define ID model (single kick)


A = torch.tensor([[-0.000152246, 0.0, 0.0, 0.0],
                  [0.0, 0.0, 0.0, 0.0],
                  [0.0, 0.0, 0.0156226, 0.0],
                  [0.0, 0.0, 0.0, 0.0]], dtype=dtype)

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

# Insert ID into the existing lattice
# This will replace the target marker

ring.flatten()
ring.replace('MSS_S01', ID)
ring.splice()

# Describe

ring.describe

{'BPM': 168,
 'Drift': 720,
 'Dipole': 156,
 'Quadrupole': 360,
 'Marker': 23,
 'Matrix': 1}

# Compute tunes (fractional part)

nux_id, nuy_id = tune(ring, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id, etapx_id, etaqy_id, etapy_id = dispersion(ring, orbit, [], limit=1)

# Compute twiss parameters

ax_id, bx_id, ay_id, by_id = twiss(ring, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux_id, muy_id = advance(ring, [], alignment=False, matched=True).T

# Compute coupling

c_id = coupling(ring, [])

# Tune shifts

print((nux - nux_id))
print((nuy - nuy_id))

tensor(4.8186e-05, dtype=torch.float64)
tensor(-0.0025, dtype=torch.float64)

# Coupling (minimal tune distance)

print(c)
print(c_id)

tensor(0., dtype=torch.float64)
tensor(0., dtype=torch.float64)

# 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()

# 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(0.0223, dtype=torch.float64)
tensor(1.2003, dtype=torch.float64)

# 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(0.0171, dtype=torch.float64)
tensor(1.1374, dtype=torch.float64)

# Define ID model (one kick for each period)
# Note, since ID is 'short' new diagonal elements are relatively small

A = torch.tensor([[-0.000152248, 0.0, 0.0, 0.0],
                  [0.0, -5.65919E-6, 0.0, 0.0],
                  [0.0, 0.0, 0.0155958, 0.0],
                  [0.0, 0.0, 0.0, 0.000581346]], dtype=dtype)

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

# Replace ID

ring.flatten()
ring.replace("ID", ID)
ring.splice()
ring.describe

{'BPM': 168,
 'Drift': 720,
 'Dipole': 156,
 'Quadrupole': 360,
 'Marker': 23,
 'Matrix': 1}

# Compute tunes (fractional part)

nux_id, nuy_id = tune(ring, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_id, etapx_id, etaqy_id, etapy_id = dispersion(ring, orbit, [], limit=1)

# Compute twiss parameters

ax_id, bx_id, ay_id, by_id = twiss(ring, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux_id, muy_id = advance(ring, [], alignment=False, matched=True).T

# Compute coupling

c_id = coupling(ring, [])

# Tune shifts

print((nux - nux_id))
print((nuy - nuy_id))

tensor(4.8299e-05, dtype=torch.float64)
tensor(-0.0025, dtype=torch.float64)

# Coupling (minimal tune distance)

print(c)
print(c_id)

tensor(0., dtype=torch.float64)
tensor(0., dtype=torch.float64)

# 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()

# 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(0.0223, dtype=torch.float64)
tensor(1.1868, dtype=torch.float64)

# 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(0.0171, dtype=torch.float64)
tensor(1.1249, dtype=torch.float64)

# Define ID model (full model)


A = torch.tensor([[-0.000249627, -1.0568E-9, 1.6222E-7,     -1.17708E-7],
                  [-1.0568E-9, -9.5497E-6, -1.38412E-7,     -5.44323E-8],
                  [1.6222E-7, -1.38412E-7, 0.0176404, -7.44269E-8],
                  [-1.17708E-7,     -5.44323E-8, -7.44269E-8, 0.000856394]], dtype=dtype)

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

# Replace ID

ring.flatten()
ring.replace("ID", ID)
ring.splice()
ring.describe

{'BPM': 168,
 'Drift': 720,
 'Dipole': 156,
 'Quadrupole': 360,
 'Marker': 23,
 'Matrix': 1}

# Compute tunes (fractional part)

nux_rk, nuy_rk = tune(ring, [], matched=True, limit=1)

# Compute dispersion

orbit = torch.tensor(4*[0.0], dtype=dtype)
etaqx_rk, etapx_rk, etaqy_rk, etapy_rk = dispersion(ring, orbit, [], limit=1)

# Compute twiss parameters

ax_rk, bx_rk, ay_rk, by_rk = twiss(ring, [], matched=True, advance=True, full=False).T

# Compute phase advances

mux_rk, muy_rk = advance(ring, [], alignment=False, matched=True).T

# Compute coupling

c_rk = coupling(ring, [])

# Tune shifts

print((nux - nux_rk))
print((nuy - nuy_rk))

tensor(7.9195e-05, dtype=torch.float64)
tensor(-0.0028, dtype=torch.float64)

# Coupling (minimal tune distance)

print(c)
print(c_rk)

tensor(0., dtype=torch.float64)
tensor(6.5965e-08, dtype=torch.float64)

# Dispersion

plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), (etaqx - etaqx_rk).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), (etaqy - etaqy_rk).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.tight_layout()
plt.show()

# Beta-beating

plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((bx - bx_rk)/bx).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((by - by_rk)/by).cpu().numpy(), fmt='-', marker='x', color='red', alpha=0.75)
plt.tight_layout()
plt.show()

print(100*(((bx - bx_rk)/bx)**2).mean().sqrt())
print(100*(((by - by_rk)/by)**2).mean().sqrt())

tensor(0.0365, dtype=torch.float64)
tensor(1.3367, dtype=torch.float64)

# Phase advance

plt.figure(figsize=(12, 4))
plt.errorbar(ring.locations().cpu().numpy(), 100*((mux - mux_rk)/mux).cpu().numpy(), fmt='-', marker='x', color='blue', alpha=0.75)
plt.errorbar(ring.locations().cpu().numpy(), 100*((muy - muy_rk)/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(0.0171, dtype=torch.float64)
tensor(1.1249, dtype=torch.float64)

# Beta-beating and dispersion

bx_id_bb = 100.0*(bx - bx_id)/bx
by_id_bb = 100.0*(by - by_id)/by
bx_rk_bb = 100.0*(bx - bx_rk)/ bx
by_rk_bb = 100.0*(by - by_rk)/ by

def rms(x):
    return (x**2).mean().sqrt()

rms_x_id = rms(bx_id_bb).item()
ptp_x_id = (bx_id_bb.max() - bx_id_bb.min()).item()
rms_y_id = rms(by_id_bb).item()
ptp_y_id = (by_id_bb.max() - by_id_bb.min()).item()

rms_x_rk = rms(bx_rk_bb).item()
ptp_x_rk = (bx_rk_bb.max() - bx_rk_bb.min()).item()
rms_y_rk = rms(by_rk_bb).item()
ptp_y_rk = (by_rk_bb.max() - by_rk_bb.min()).item()

s = ring.locations().cpu().numpy()
bx_id_np = bx_id_bb.cpu().numpy()
by_id_np = by_id_bb.cpu().numpy()
bx_rk_np = bx_rk_bb.cpu().numpy()
by_rk_np = by_rk_bb.cpu().numpy()

etax_id = etaqx - etaqx_id
etay_id = etaqy - etaqy_id
etax_rk = etaqx - etaqx_rk
etay_rk = etaqy - etaqy_rk

rms_etax_id = rms(etax_id).item()
ptp_etax_id = (etax_id.max() - etax_id.min()).item()
rms_etay_id = rms(etay_id).item()
ptp_etay_id = (etay_id.max() - etay_id.min()).item()

rms_etax_rk = rms(etax_rk).item()
ptp_etax_rk = (etax_rk.max() - etax_rk.min()).item()
rms_etay_rk = rms(etay_rk).item()
ptp_etay_rk = (etay_rk.max() - etay_rk.min()).item()

etax_id_np = etax_id.cpu().numpy()
etay_id_np = etay_id.cpu().numpy()
etax_rk_np = etax_rk.cpu().numpy()
etay_rk_np = etay_rk.cpu().numpy()

fig, (ax, ay) = plt.subplots(
    2, 1, figsize=(16, 10),
    sharex=True,
    gridspec_kw={'hspace': 0.3}
)

ax.errorbar(s, bx_id_np, fmt='-', marker='x', color='blue', alpha=0.75, lw=2.0, label=r'nk, $\beta_x$')
ax.errorbar(s, by_id_np, fmt='-', marker='x', color='red',  alpha=0.75, lw=2.0, label=r'nk, $\beta_y$')
ax.errorbar(s, bx_rk_np, fmt='-', marker='o', color='blue', alpha=0.75, lw=2.0, label=r'rk, $\beta_x$')
ax.errorbar(s, by_rk_np, fmt='-', marker='o', color='red',  alpha=0.75, lw=2.0, label=r'rk, $\beta_y$')
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_id:05.2f}\% \quad RMS$_y$={rms_y_id:05.2f}\% \quad '
    rf'PTP$_x$={ptp_x_id:05.2f}\% \quad PTP$_y$={ptp_y_id:05.2f}\% \quad '
    rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_id)}{(nux - nux_id).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_id)}{(nuy - nuy_id).abs().item():.4f}'
    rf'\quad C={c_id.item():.6f}'
)
ax.text(0.0, 1.10, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
    rf'RMS$_x$={rms_x_rk:05.2f}\% \quad RMS$_y$={rms_y_rk:05.2f}\% \quad '
    rf'PTP$_x$={ptp_x_rk:05.2f}\% \quad PTP$_y$={ptp_y_rk:05.2f}\% \quad '
    rf'$\Delta \nu_x$={(lambda x: '-' if x < 0 else '~')(nux - nux_rk)}{(nux - nux_rk).abs().item():.4f} \quad $\Delta \nu_y$={(lambda x: '-' if x < 0 else '~')(nuy - nuy_rk)}{(nuy - nuy_rk).abs().item():.4f}'
    rf'\quad C={c_rk.item():.6f}'
)
ax.text(0.0, 1.025, title, transform=ax.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
ax.legend(loc='upper right', frameon=False, fontsize=14, ncol=4)
ax.set_ylim(-5, 5)

ay.errorbar(s, etax_id_np, fmt='-', marker='x', color='blue', alpha=0.75, lw=2.0, label=r'initial, $\eta_x$')
ay.errorbar(s, etay_id_np, fmt='-', marker='x', color='red', alpha=0.75, lw=2.0, label=r'initial, $\eta_y$')
ay.errorbar(s, etax_rk_np, fmt='-', marker='o', color='blue',alpha=0.75, lw=2.0, label=r'final, $\eta_x$')
ay.errorbar(s, etay_rk_np, fmt='-', marker='o', color='red', alpha=0.75, lw=2.0, label=r'final, $\eta_y$')
ay.set_xlabel('s [m]', fontsize=18)
ay.set_ylabel(r'$\Delta \eta$ [m]', fontsize=18)
ay.tick_params(width=2, labelsize=16)
ay.tick_params(axis='x', length=8, direction='in')
ay.tick_params(axis='y', length=8, direction='in')
title = (
    rf'RMS$_x$={rms_etax_id:.4E} m \quad RMS$_y$={rms_etay_id:.4E} m  \quad '
    rf'PTP$_x$={ptp_etax_id:.4E} m \quad PTP$_y$={ptp_etay_id:.4E} m  \quad '
)
ay.text(0.0, 1.125, title, transform=ay.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')
title = (
    rf'RMS$_x$={rms_etax_rk:.4E} m  \quad RMS$_y$={rms_etay_rk:.4E} m  \quad '
    rf'PTP$_x$={ptp_etax_rk:.4E} m  \quad PTP$_y$={ptp_etay_rk:.4E} m  \quad '
)
ay.text(0.0, 1.05, title, transform=ay.transAxes, ha='left', va='bottom', fontsize=16, fontfamily='monospace')

plt.setp(ax.spines.values(), linewidth=2.0)
plt.setp(ay.spines.values(), linewidth=2.0)

plt.show()