Example-41: Tune (Tune and chormaticity correction)

# Import

from random import random
from pprint import pprint

import torch
from torch import Tensor

from pathlib import Path

from model.library.line import Line

from model.command.external import load_sdds
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

# Load ELEGANT twiss

path = Path('ic.twiss')
parameters, columns = load_sdds(path)

nu_qx:Tensor = torch.tensor(parameters['nux'] % 1, dtype=torch.float64)
nu_qy:Tensor = torch.tensor(parameters['nuy'] % 1, dtype=torch.float64)

psi_qx:Tensor = torch.tensor(parameters['dnux/dp'], dtype=torch.float64)
psi_qy:Tensor = torch.tensor(parameters['dnuy/dp'], dtype=torch.float64)

# Build and setup lattice

# Note, sextupoles are turned off and dipoles are linear

# Load ELEGANT table

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

# Build ELEGANT table

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

# Merge drifts

ring.merge()

# Turn off sextupoles and set linear dipoles

for element in ring:
    if element.__class__.__name__ == 'Sextupole':
        element.ms = 0.0
    if element.__class__.__name__ == 'Dipole':
        element.linear = True

# Set number of elements of different kinds

nb = ring.describe['BPM']
nq = ring.describe['Quadrupole']
ns = ring.describe['Sextupole']

# Compute tunes (fractional part)

guess = torch.tensor(4*[0.0], dtype=torch.float64)
nuqx, nuqy = tune(ring, [], alignment=False, matched=True, guess=guess, limit=8, epsilon=1.0E-9)

# Compare with elegant

print(torch.allclose(nu_qx, nuqx))
print(torch.allclose(nu_qy, nuqy))

True
True

# Compute tunes derivative with respect to deviations in quadrupole settings

kn = torch.zeros(nq, dtype=torch.float64)

pprint(torch.func.jacrev(lambda kn: tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))(kn))
print()

tensor([[ 0.0187,  0.0217,  0.0461,  0.0601,  0.0484,  0.0297,  0.0299,  0.0488,
          0.0597,  0.0457,  0.0214,  0.0192,  0.0669,  0.0667,  0.0191,  0.0216,
          0.0461,  0.0597,  0.0483,  0.0296,  0.0297,  0.0485,  0.0595,  0.0459,
          0.0209,  0.0196,  0.0683,  0.0647],
        [-0.0668, -0.0283, -0.0136, -0.0078, -0.0141, -0.0205, -0.0202, -0.0137,
         -0.0079, -0.0137, -0.0282, -0.0661, -0.0250, -0.0251, -0.0663, -0.0282,
         -0.0137, -0.0079, -0.0138, -0.0203, -0.0205, -0.0140, -0.0078, -0.0137,
         -0.0288, -0.0659, -0.0250, -0.0258]], dtype=torch.float64)

# Compute chromaticities derivatives

ms = torch.zeros(ns, dtype=torch.float64)

pprint(torch.func.jacrev(lambda ms: chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))(ms))
print()

tensor([[ 0.0051,  0.0144,  0.0143,  0.0050,  0.0049,  0.0142,  0.0143,  0.0051,
          0.0050,  0.0143,  0.0142,  0.0050,  0.0049,  0.0141,  0.0142,  0.0050],
        [-0.0047, -0.0037, -0.0038, -0.0049, -0.0049, -0.0038, -0.0038, -0.0049,
         -0.0049, -0.0038, -0.0038, -0.0049, -0.0049, -0.0038, -0.0037, -0.0048]],
       dtype=torch.float64)

# Tune correction

# [nux, nuy] = M [..., kn_i, ..., ks_i, ...]

kn = torch.zeros(nq, dtype=torch.float64)
response = torch.func.jacrev(lambda kn: tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))(kn)
print(response.shape)
print()

# Set target tunes

nuqx_target = torch.tensor(0.75, dtype=torch.float64)
nuqy_target = torch.tensor(0.85, dtype=torch.float64)

# Perform correction (model to experiment)

lr = 1.0

kn = torch.zeros(nq, dtype=torch.float64)

for _ in range(4):
    nuqx, nuqy = tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None)
    dkn = (- lr*torch.linalg.lstsq(response, torch.stack([nuqx, nuqy]) - torch.stack([nuqx_target, nuqy_target]), driver='gels').solution)
    kn += dkn
    print((torch.stack([nuqx, nuqy]) - torch.stack([nuqx_target, nuqy_target])).norm())
print()

print(tune(ring, [kn], ('kn', ['Quadrupole'], None, None), matched=True, limit=1, epsilon=None))
print()

torch.Size([2, 28])

tensor(0.0471, dtype=torch.float64)
tensor(0.0007, dtype=torch.float64)
tensor(1.9547e-05, dtype=torch.float64)
tensor(5.4013e-07, dtype=torch.float64)

tensor([0.7500, 0.8500], dtype=torch.float64)

# Chromaticity correction
# Note, one iteration is sufficient for chromaticity correction

# [psix, psiy] = M [..., ms_i, ..., ms_i, ...]

ms = torch.zeros(ns, dtype=torch.float64)
response = torch.func.jacrev(lambda ms: chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))(ms)
print(response.shape)
print()

# Set target tunes

psiqx_target = torch.tensor(1.0, dtype=torch.float64)
psiqy_target = torch.tensor(1.0, dtype=torch.float64)

# Perform correction (model to experiment)

lr = 1.0

ks = torch.zeros(ns, dtype=torch.float64)

for _ in range(1):
    psiqx, psiqy = chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None)
    dms = (- lr*torch.linalg.lstsq(response, torch.stack([psiqx, psiqy]) - torch.stack([psiqx_target, psiqy_target]), driver='gels').solution)
    ms += dms
    print((torch.stack([psiqx, psiqy]) - torch.stack([psiqx_target, psiqy_target])).norm())
print()

print(chromaticity(ring, [ms], ('ms', ['Sextupole'], None, None), matched=True, limit=1, epsilon=None))
print()

torch.Size([2, 16])

tensor(10.1673, dtype=torch.float64)

tensor([1.0000, 1.0000], dtype=torch.float64)