Example-40: Tune (Computation of tunes and chromaticities)

# 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

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

# Set exact dipoles

for element in ring:
    if element.__class__.__name__ == 'Dipole':
        element.exact = True

# Set number of integration steps and integration order in sextupoles and dipoles

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

# 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 chormaticities

psiqx, psiqy = chromaticity(ring, [], alignment=False, matched=True, guess=guess, limit=1, epsilon=None)

print(torch.allclose(psi_qx, psiqx))
print(torch.allclose(psi_qy, psiqy))

# Note, since closed orbit depends on dp, matched flag should be True
# These chromaticity values depend on the number of intergation steps in sextupoles
# Without linear flag, dipoles will also have sextupole like component in the body due to cylindrical potential
# With exact flag in dipoles, fringes will also contribute to chromaticities
# These result match (converge to) MADX 5.09.01 TWISS (default settings in SBENDs) and ELEGNAT SVN revision: 30188M for CSBENDs with FINT=0.0 and EDGE_ORDER=0

True
True