Optics calculations¶
Optics calculations can be performed via the compute module.
Closed orbit search¶
Using compute.closed_orbit we can perform closed orbit search for a given lattice:
>>> from importlib import resources
>>> from particleflow.build import from_file
>>> from particleflow.compute import closed_orbit, linear_closed_orbit
>>> from particleflow.elements import Kicker, HKicker, VKicker
>>> import particleflow.test.sequences
>>>
>>> with resources.path(particleflow.test.sequences, 'cryring.seq') as path:
... lattice = from_file(path)
...
>>> thin = lattice.makethin({Kicker: 1})
>>> closed_orbit(thin, order=1)
tensor([[0.],
[0.],
[0.],
[0.],
[0.],
[0.]])
As can be seen from the above example we first need to convert all elements that don’t support thick tracking (the kicker magnets) to thin elements because the closed orbit search is performed by tracking through the lattice. Since all kickers are off (kick = 0) the closed orbit is just zero. Let’s add some kicks:
>>> import random
>>>
>>> for kicker in lattice[HKicker] + lattice[VKicker]:
... kicker.kick = random.uniform(-0.001, 0.001)
...
>>> thin = lattice.makethin({Kicker: 1})
>>> closed_orbit(thin, order=1).flatten()
tensor([-0.0002, -0.0020, -0.0028, -0.0004, 0.0000, 0.0000])
>>> linear_closed_orbit(thin).flatten()
tensor([-0.0002, -0.0020, -0.0028, -0.0004, 0.0000, 0.0000])
One important thing to note is that we need to assign the kicks to the original lattice, since the thin version doesn’t contain the original kickers anymore. Then for the new kicker values we need to makethin the lattice again before performing the closed orbit search. This seems somewhat repetitive but it is important in order to maintain the relation between thick and thin elements. Especially if optimization parameters are involved, it is important to always makethin the original lattice (which stores the optimization parameters) in order to always get the up-to-date values of the optimization parameters. linear_closed_orbit computes the closed orbit directly from first order transfer maps (using a slightly different algorithm).
Transfer maps (sector maps)¶
By using compute.sectormaps we can compute the transfer maps along the lattice. The parameters accumulate
let’s us specify
whether we want the local maps (accumulate=False
) or the cumulative transfer maps w.r.t. to the start of the lattice
(accumulate=True
).
>>> from particleflow.compute import sectormaps
>>>
>>> maps = dict(sectormaps(lattice))
>>> maps[lattice[HKicker, 0].label]
tensor([[1.0000, 0.4573, 0.0000, 0.0000, 0.0000, 0.0000],
[0.0000, 1.0000, 0.0000, 0.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 1.0000, 0.4573, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, 1.0000, 0.0000, 0.0000],
[0.0000, 0.0000, 0.0000, 0.0000, 1.0000, 3.3648],
[0.0000, 0.0000, 0.0000, 0.0000, 0.0000, 1.0000]])
Since by default the cumulative transfer maps are computed, the first order map for the first kicker is not the identity
matrix. If we want the local maps instead we can use accumulate=False
:
>>> maps = dict(sectormaps(lattice, accumulate=False))
>>> maps[lattice[HKicker, 0].label]
tensor([[1., 0., 0., 0., 0., 0.],
[0., 1., 0., 0., 0., 0.],
[0., 0., 1., 0., 0., 0.],
[0., 0., 0., 1., 0., 0.],
[0., 0., 0., 0., 1., 0.],
[0., 0., 0., 0., 0., 1.]])
Orbit Response Matrix¶
Using compute.orm we can compute the orbit response matrix for a given lattice. We need to specify the kickers and monitors to be used, which can be done in a similar way as for selecting lattice elements in general: either we can specify an identifier that selects multiple elements directly, such as a lattice element type or a regex, or we can specify a list of single element identifiers, such as unambiguous labels for example. Let’s compute the horizontal ORM for one of the example lattices:
>>> from importlib import resources
>>> from particleflow.build import from_file
>>> from particleflow.compute import orm
>>> from particleflow.elements import HKicker, HMonitor
>>> import particleflow.test.sequences
>>>
>>> with resources.path(particleflow.test.sequences, 'cryring.seq') as path:
... lattice = from_file(path)
...
>>> orm_x, orm_y = orm(lattice, kickers=HKicker, monitors=HMonitor)
Here we don’t need to call makethin beforehand because this will be done inside the orm function. This is necessary because the orm function will temporarily vary the kicker strengths and, as explained above, for each change to the original lattice we need to create a new thin version (i.e. changes to the original lattice are not automatically mapped to any thin versions that have been created before).
>>> orm_x
tensor([[ 1.6402, 0.8140, 0.4583, -0.3921, 1.1925, 1.4039, -1.9135, -0.4376,
1.7387],
[ 1.5428, 1.0098, 0.6885, -0.8058, 1.3288, 1.7722, -2.1046, -0.6838,
1.6842],
[ 0.8887, 1.7728, 1.6933, -2.7059, 1.7242, 3.2348, -2.6066, -1.7705,
1.0460],
[ 2.0276, 0.9366, 1.2629, -2.4109, 0.4483, 1.8029, -0.5499, -1.3686,
-0.8998],
[ 1.6111, 1.2233, 0.8969, -0.0308, -1.3899, -1.3615, 2.2668, 0.2506,
-2.3613],
[ 1.0921, 1.8082, 1.6400, 1.2010, -2.0740, -2.7341, 3.2892, 1.0372,
-2.6678],
[-2.9513, -1.3954, -0.7518, 1.1115, 1.5688, 0.9301, -2.6391, 0.2900,
3.4040],
[ 3.2348, 0.5891, -0.1673, 1.0460, 1.4899, 0.8887, 1.2091, -1.1510,
-2.7059],
[-2.7059, 0.4745, 1.1597, -2.6066, -0.4258, 1.0460, 0.8887, 1.8197,
1.2091],
[ 1.2091, -1.3684, -1.7372, 3.2348, -0.7906, -2.6066, 1.0460, 1.8726,
0.8887],
[ 2.1282, -0.1676, -0.6953, 1.6820, 0.5093, -0.4424, -0.9556, 0.7970,
2.0117],
[ 2.0422, 0.0053, -0.4921, 1.3166, 0.6296, -0.1171, -1.1243, 0.5795,
1.9636]])
>>> orm_y
tensor([[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0., 0., 0., 0., 0.]])
>>>
>>> orm_x.shape
torch.Size([12, 9])
>>> len(lattice[HKicker]), len(lattice[HMonitor])
(12, 9)
The ORMs rows correspond to kickers and the columns to monitors, as can be seen from the column and row count. Since we didn’t specify to include the vertical kickers and the lattice does not contain any coupling the vertical ORM is just zero.