Modifying lattices

We can modify single lattice elements or the lattice itself. In the previous section there was already a hint about how to replace specific lattice elements. This can be done via lattice[identifier] = ... where identifier must unambiguously identify a lattice element. That is lattice[identifier] (not setting, but getting the element) should return a single element, not a list of elements. Note that identifer can be a tuple as well, in order to narrow down the selection. For example, let’s offset the Quadrupole with label "gte1qd11" and tilt the second Kicker:

[1]:

from importlib import resources
import warnings
from dipas.build import from_file
from dipas.elements import HKicker, Offset, Tilt
import dipas.test.sequences

warnings.simplefilter('ignore')

with resources.path(dipas.test.sequences, 'hades.seq') as path:
    lattice = from_file(path)  # Load a fresh lattice.

print(lattice['gte1qd11'])  # Returns a single element, good.
lattice['gte1qd11'] = Offset(lattice['gte1qd11'], dx=0.25, dy=0.50)
print(lattice['gte1qd11'], end='\n\n')

print(lattice[HKicker, 1])  # Returns a single element, good.
lattice[HKicker, 1] = Tilt(lattice[HKicker, 1], psi=1.0)
print(lattice[HKicker, 1])

Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), dk1=tensor(0.), label='gte1qd11')
Offset(dx=tensor(0.2500), dy=tensor(0.5000),
       target=Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), dk1=tensor(0.), label='gte1qd11'))

HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1')
Tilt(psi=tensor(1.),
     target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), dkh=tensor(0.), dkv=tensor(0.), label='gth1kx1'))

We can of course also modify attributes of single elements. For example let’s introduce some random errors to the quadrupole gradient strengths:

[2]:

import numpy as np
from dipas.elements import Quadrupole

for quad in lattice[Quadrupole]:
    quad.element.k1.data += np.random.normal(scale=0.1)

Two things are worth noting here:

  1. We used quad.element.k1 instead of just quad.k1. This is because lattice[Quadrupole] returns a list of all Quadrupole elements, potentially wrapped by alignment error classes. Because we applied an offset to the first quadrupole beforehand, the first quad is actually an Offset object. By using quad.element we ensure that we always get the underlying Quadrupole object. Using element on a Quadrupole itself will just return the same object.

  2. We used k1.data instead of just k1. This is because the MADX sequence file that we used to parse the lattice from actually contained optimization parameter definition (see the next section for more details) and so we need to use .data to modify the actual number of the tensor.