Specifying optimization parameters

Optimization parameters can be specified either directly in the MADX files that are parsed or they can set on the lattice elements after parsing (or building in general), similar to modifying lattice elements as seen in the previous section.

Inside MADX scripts

For details please consider the documentation part about “Compatibility with MADX”.

Optimization parameters (“flow variables”) can be indicates by placing dedicated comments in MADX scripts. There comments should be of the form // <some optional text> [flow] variable and should either precede or conclude the line of a variable definition or attribute assignment. Let’s peek into one of the example scripts:

[1]:
from importlib import resources
from pprint import pprint
import dipas.test.sequences

script = resources.read_text(dipas.test.sequences, 'hades.seq')
script = script.splitlines()
pprint(script[:4])
['beam, particle=ion, charge=6, energy=28.5779291448, mass=11.1779291448;',
 '',
 'k1l_GTE1QD11 := 0.3774561583995819;       // [flow] variable',
 'k1l_GTE1QD12 := -0.35923901200294495;     // [flow] variable']

Here we can see that the two variables k1l_GTE1QD11 and k1l_GTE1QD12 have been declared as optimization parameters. Now let’s check the lattice element after parsing the script:

[2]:
import warnings
from dipas.build import from_script

warnings.simplefilter('ignore')

lattice = from_script('\n'.join(script))
print(lattice['gte1qd11'])
print(lattice['gte1qd12'])
print(type(lattice['gte1qd11'].k1))
Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), dk1=tensor(0.), label='gte1qd11')
Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(-0.5394, requires_grad=True), dk1=tensor(0.), label='gte1qd12')
<class 'torch.nn.parameter.Parameter'>

Here we can see that the k1 parameters are indicates as Parameter which can be optimized for using PyTorch’s optimization machinery.

Using the API

Now let’s modify the script so that the first quadrupole’s k1 attribute won’t be parsed to a parameter:

[3]:
script[2] = script[2].split(';')[0] + ';'
pprint(script[:4])
lattice = from_script('\n'.join(script))
print(lattice['gte1qd11'])
print(lattice['gte1qd12'])
['beam, particle=ion, charge=6, energy=28.5779291448, mass=11.1779291448;',
 '',
 'k1l_GTE1QD11 := 0.3774561583995819;',
 'k1l_GTE1QD12 := -0.35923901200294495;     // [flow] variable']
Quadrupole(l=tensor(0.6660), k1=tensor(0.5668), dk1=tensor(0.), label='gte1qd11')
Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(-0.5394, requires_grad=True), dk1=tensor(0.), label='gte1qd12')

Now the lattice['gte1qd11'].k1 attribute is set as a tensor. If we want to optimize for that value nevertheless we can simply convert it to a parameter manually:

[4]:
import torch

lattice['gte1qd11'].k1 = torch.nn.Parameter(lattice['gte1qd11'].k1)
print(lattice['gte1qd11'])
Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), dk1=tensor(0.), label='gte1qd11')

Similarly we could convert the k1-values of all quadrupoles to parameters:

[5]:
from dipas.elements import Quadrupole

for q in lattice[Quadrupole]:
    q.k1 = torch.nn.Parameter(q.k1)

For the example lattice however the k1-values are already parameters.