Tune matching

In the following example some gradient errors will be introduced to quadrupoles and then the tunes will be rematched to their original values by varying the strength of horizontally focusing quadrupoles. To begin with we load the example MADX script:

[ ]:
with open('example.madx') as fh:
    script = fh.read()

As mentioned already, the script assigns gradient errors to all 18 quadrupoles and then varies the strength of the 12 horizontal quadrupoles in order to rematch the tune. Let’s inspect the results by running the script via dipas.madx.run_script:

[2]:
import os
from dipas.madx import run_script

result = run_script(
    script,
    {'twiss': True, 'twiss_error': True, 'twiss_matched': True, 'errors': False},
    madx=os.path.expanduser('~/bin/madx')
)
twiss = result['twiss']
twiss_error = result['twiss_error']
twiss_matched = result['twiss_matched']
errors = result['errors']
/home/dominik/Projects/DiPAS/dipas/madx/utils.py:247: UserWarning: MADX issued the following warnings: ['MTSIMP More variables than constraints seen. SIMPLEX may not converge to optimal solution.']
  warnings.warn(f'MADX issued the following warnings: {warnings_list}')

Here we specified True for the twiss files since we want to retrieve the “@”-prefixed meta data besides the actual data frame (since the meta data contains the tune values). For the errors file we’re not interested in meta data and so False will result in just the corresponding data frame. Let’s inspect the tune values:

[3]:
print('Tune values:')
print(f' - original: {        twiss[1]["Q1"]:.3f}, {        twiss[1]["Q2"]:.3f}')
print(f' - shifted : {  twiss_error[1]["Q1"]:.3f}, {  twiss_error[1]["Q2"]:.3f}')
print(f' - matched : {twiss_matched[1]["Q1"]:.3f}, {twiss_matched[1]["Q2"]:.3f}')
Tune values:
 - original: 2.420, 2.420
 - shifted : 2.362, 2.459
 - matched : 2.420, 2.420

Now let’s do the same thing with gradient-based optimization via DiPAS. First we load the script via dipas.build.from_script and assign the errors via the errors data frame:

[4]:
from dipas.build import from_script

lattice = from_script(script, errors=errors)

We select the quadrupoles by checking for k1 != 0 since the script defines drift spaces as k1 == 0 quadrupoles:

[5]:
from dipas.elements import Quadrupole, Parameter

quadrupoles = [q for q in lattice[Quadrupole] if q.k1 != 0]
for q in quadrupoles:
    print(f'{q.label}, k1 = {q.k1: .6f}, dk1 = {q.dk1: .6f}')
yr02qs1, k1 =  1.760051, dk1 = -0.013478
yr02qs2, k1 = -2.252993, dk1 =  0.012333
yr02qs3, k1 =  1.760051, dk1 = -0.077924
yr04qs1, k1 =  1.760051, dk1 = -0.003707
yr04qs2, k1 = -2.252993, dk1 =  0.047466
yr04qs3, k1 =  1.760051, dk1 =  0.035394
yr06qs1, k1 =  1.760051, dk1 = -0.144814
yr06qs2, k1 = -2.252993, dk1 = -0.069407
yr06qs3, k1 =  1.760051, dk1 = -0.069233
yr08qs1, k1 =  1.760051, dk1 =  0.055320
yr08qs2, k1 = -2.252993, dk1 = -0.008849
yr08qs3, k1 =  1.760051, dk1 =  0.001484
yr10qs1, k1 =  1.760051, dk1 = -0.086816
yr10qs2, k1 = -2.252993, dk1 = -0.001718
yr10qs3, k1 =  1.760051, dk1 = -0.015357
yr12qs1, k1 =  1.760051, dk1 = -0.096350
yr12qs2, k1 = -2.252993, dk1 = -0.068815
yr12qs3, k1 =  1.760051, dk1 =  0.135457

For the matching we will use the horizontally focusing quadrupoles and so we’ll select these and turn their k1 attributes into parameters (being varied during the optimization). We have to call update_transfer_map as well in order for the change to k1 to become effective. In general, after altering any (to-be-)parametrized attribute of a lattice element (also its value), we need to call the update_transfer_map method for bringing the change into effect.

[6]:
h_quadrupoles = [q for q in quadrupoles if q.k1 > 0]
for q in h_quadrupoles:
    q.k1 = Parameter(q.k1)
    q.update_transfer_map()
print(f'# Parameters: {len(list(lattice.parameters()))}')
# Parameters: 12

Next we’ll prepare the optimization by creating an optimizer and defining a cost (loss) function. The cost function indicates the distance to the optimization target(s). Before we can use the lattice’s transfer maps we need to convert Kicker elements to thin counterparts since the transfer map for a thick kicker doesn’t exist. We specify two slices placed at the edges of the original elements (this is the configuration used by MADX during TWISS computation).

We use the dipas.compute.twiss function for computing the tune values (alongside other lattice functions). When computing the gradients via cost.backward we specify retain_graph=True since at every iteration we’re optimizing against the same data and so retaining the graph is required (e.g. the transfer map tensors of lattice elements will be reused at every iteration so their memory buffers need to be retained). At the end of each iteration, after the optimizer has updated the k1 values, we need to call update_transfer_map again in order to activate the updates.

[7]:
import itertools as it
import dipas.compute as compute
from dipas.elements import Kicker
import torch

lattice = lattice.makethin({Kicker: 2}, style={Kicker: 'edge'})
print(f'# Parameters: {len(list(lattice.parameters()))}')

targets = {'Q1': torch.tensor(twiss[1]['Q1']), 'Q2': torch.tensor(twiss[1]['Q2'])}
optimizer = torch.optim.LBFGS(lattice.parameters())
cost_fn = torch.nn.MSELoss()

for step in it.count():
    def closure():
        optimizer.zero_grad()
        data = compute.twiss(lattice)
        Q1, Q2 = data['Q1'], data['Q2']
        cost = cost_fn(Q1, targets['Q1']) + cost_fn(Q2, targets['Q2'])
        print(f'Step {step:03d}: Q1 = {Q1:.3f}, Q2 = {Q2:.3f}, cost = {cost:.2e}')
        if cost < 1e-6:
            raise RuntimeError
        cost.backward(retain_graph=True)
        return cost

    try:
        optimizer.step(closure)
    except RuntimeError:
        break

    for q in h_quadrupoles:
        q.update_transfer_map()
# Parameters: 12
Step 000: Q1 = 2.362, Q2 = 2.459, cost = 4.88e-03
Step 000: Q1 = 2.362, Q2 = 2.459, cost = 4.88e-03
Step 001: Q1 = 2.386, Q2 = 2.446, cost = 1.87e-03
Step 001: Q1 = 2.386, Q2 = 2.446, cost = 1.87e-03
Step 002: Q1 = 2.408, Q2 = 2.435, cost = 3.69e-04
Step 002: Q1 = 2.408, Q2 = 2.435, cost = 3.69e-04
Step 003: Q1 = 2.418, Q2 = 2.430, cost = 1.10e-04
Step 003: Q1 = 2.418, Q2 = 2.430, cost = 1.10e-04
Step 004: Q1 = 2.422, Q2 = 2.428, cost = 7.26e-05
Step 004: Q1 = 2.422, Q2 = 2.428, cost = 7.26e-05
Step 005: Q1 = 2.423, Q2 = 2.428, cost = 6.80e-05
Step 005: Q1 = 2.423, Q2 = 2.428, cost = 6.80e-05
Step 006: Q1 = 2.424, Q2 = 2.427, cost = 6.25e-05
Step 006: Q1 = 2.424, Q2 = 2.427, cost = 6.25e-05
Step 007: Q1 = 2.425, Q2 = 2.425, cost = 5.13e-05
Step 007: Q1 = 2.425, Q2 = 2.425, cost = 5.13e-05
Step 008: Q1 = 2.425, Q2 = 2.423, cost = 3.71e-05
Step 008: Q1 = 2.425, Q2 = 2.423, cost = 3.71e-05
Step 009: Q1 = 2.424, Q2 = 2.422, cost = 1.66e-05
Step 009: Q1 = 2.424, Q2 = 2.422, cost = 1.66e-05
Step 010: Q1 = 2.420, Q2 = 2.420, cost = 3.00e-07

Now let’s crosscheck the solution by running it through MADX. We can use dipas.build.create_script in order to convert the lattice object to a corresponding MADX script.

[8]:
from dipas.build import create_script

twiss_check = run_script(
    create_script(sequence=lattice, errors=True, beam={'particle': 'proton', 'energy': 1}),
    twiss=True,
    madx=os.path.expanduser('~/bin/madx')
)['twiss']
print(f'Tune values: {twiss_check[1]["Q1"]:.3f}, {twiss_check[1]["Q2"]:.3f}')
Tune values: 2.420, 2.420

Finally let’s compare the results to the solution which MADX originally computed:

[9]:
import pandas as pd

q_names = [q.label.upper() for q in h_quadrupoles]
results = twiss_matched[0].set_index('NAME').loc[q_names, ['K1L']]
results = results.assign(DP=pd.Series([q.k1.item()*q.l.item() for q in h_quadrupoles], index=q_names))
results.columns = ['K1L_MADX', 'K1L_PF']
print(results)
         K1L_MADX    K1L_PF
NAME
YR02QS1  0.521336  0.517564
YR02QS3  0.513679  0.509166
YR04QS1  0.538908  0.530332
YR04QS3  0.517392  0.527857
YR06QS1  0.514691  0.513474
YR06QS3  0.513379  0.513102
YR08QS1  0.552487  0.526501
YR08QS3  0.533606  0.526681
YR10QS1  0.506585  0.515837
YR10QS3  0.507948  0.518586
YR12QS1  0.500134  0.513845
YR12QS3  0.500342  0.506482
[10]:
%matplotlib inline

ax = results.plot(kind='bar', figsize=(9, 4))
ax.set_ylabel('K1L [1/m]')
ax.plot([0, len(results)], [0.5086546699]*2, '--', color='red', label='initial')
ax.legend()
[10]:
<matplotlib.legend.Legend at 0x7f6063591910>
../_images/examples_tune_matching_18_1.png
[ ]: