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:

>>> from importlib import resources
>>> from particleflow.build import from_file
>>> from particleflow.elements import Quadrupole, HKicker, Offset, Tilt
>>> import particleflow.test.sequences
>>>
>>> with resources.path(particleflow.test.sequences, 'hades.seq') as path:
...     lattice = from_file(path)  # Load a fresh lattice.
...
>>> lattice['gte1qd11']  # Returns a single element, good.
Quadrupole(l=tensor(0.6660), k1=Parameter containing: tensor(0.5668, requires_grad=True), label='gte1qd11')
>>> lattice['gte1qd11'] = Offset(lattice['gte1qd11'], dx=0.25, dy=0.50)
>>> lattice['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), label='gte1qd11'))
>>>
>>> lattice[HKicker, 1]  # Returns a single element, good.
HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=tensor(0.), label='gth1kx1')
>>> lattice[HKicker, 1] = Tilt(lattice[HKicker, 1], psi=1.0)
>>> lattice[HKicker, 1]
Tilt(psi=tensor(1.), target=HKicker(l=tensor(0.), hkick=tensor(0.), vkick=tensor(0.), kick=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:

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

Two things are worth noting here:

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.
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.