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.