Optimising Cheetah for speed

One of Cheetah’s standout features is its computational speed. This is achieved through some optimisations under the hood, which the user never needs to worry about. Often, however, there further optimisations that can be made when knowledge on how the model will be used is available. For example, in many cases, one might load a large lattice of an entire facility that has thousands of elements, but then only ever changes a handful of these elements for the experiments at hand. For this case, Cheetah offers some opt-in optimisation features that can help speed up simulations significantly by an order of magnitude or more in some cases.

import cheetah
import torch

incoming_beam = cheetah.ParameterBeam.from_astra(
    "../../tests/resources/ACHIP_EA1_2021.1351.001"
)

Let’s define a large lattice. With many quadrupole magnets and drift sections in the center and a pair of steerers at each end. We assume that the quadrupole magnets are at their design settings and will never be touched. Only the two steerers at each end are of interest to us, for example because we would like to train a neural network policy to steer the beam using these steerers. Furthermore, as many lattices do, there are a bunch of markers in this lattice. These markers may be helpful to mark certain positions along the beamline, but they don’t actually add anything to the physics of the simulation.

original_segment = cheetah.Segment(
    elements=[
        cheetah.HorizontalCorrector(
            length=torch.tensor([0.1]), angle=torch.tensor([0.0]), name="HCOR_1"
        ),
        cheetah.Drift(length=torch.tensor([0.3])),
        cheetah.VerticalCorrector(
            length=torch.tensor([0.1]), angle=torch.tensor([0.0]), name="VCOR_1"
        ),
        cheetah.Drift(length=torch.tensor([0.3])),
    ]
    + [
        cheetah.Quadrupole(length=torch.tensor([0.1]), k1=torch.tensor([4.2])),
        cheetah.Drift(length=torch.tensor([0.2])),
        cheetah.Quadrupole(length=torch.tensor([0.1]), k1=torch.tensor([-4.2])),
        cheetah.Drift(length=torch.tensor([0.2])),
        cheetah.Marker(),
        cheetah.Quadrupole(length=torch.tensor([0.1]), k1=torch.tensor([0.0])),
        cheetah.Drift(length=torch.tensor([0.2])),
    ]
    * 150
    + [
        cheetah.HorizontalCorrector(
            length=torch.tensor([0.1]), angle=torch.tensor([0.0]), name="HCOR_2"
        ),
        cheetah.Drift(length=torch.tensor([0.3])),
        cheetah.VerticalCorrector(
            length=torch.tensor([0.1]), angle=torch.tensor([0.0]), name="VCOR_2"
        ),
        cheetah.Drift(length=torch.tensor([0.3])),
    ]
)

len(original_segment.elements)

1058

First, we test how long it takes to track a beam through this segment without any optimisations beyond the ones automatically done under the hood.

%%timeit
original_segment.track(incoming_beam)

138 ms ± 431 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Just by removing unused markers, we already see a small performance improvement.

markers_removed_segment = original_segment.without_inactive_markers()

%%timeit
markers_removed_segment.track(incoming_beam)

136 ms ± 249 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

Drift sections tend to be the cheapest elements to compute. At the same time, many elements in a lattice may be switched off at any given time. When they are switched off, they behave almost exactly like drift sections, but they still require additional computations to arrive at this result. We can however safely replace them by actual Drift elements, which clearly speeds up computations.

inactive_to_drifts_segment = original_segment.inactive_elements_as_drifts(
    except_for=["HCOR_1", "VCOR_1", "HCOR_2", "VCOR_2"]
)
len(inactive_to_drifts_segment.elements)

1058

%%timeit
inactive_to_drifts_segment.track(incoming_beam)

112 ms ± 162 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

The most significant improvement can be made by merging elements that are not expected to be changed in the future. For this, Cheetah offers the transfer_maps_merged method. This will by default merge the transfer maps of all elements in the segment. In almost all realistic applications, however, there are some elements the settings of which we wish to change in the future. By passing a list of their names to except_for, we can instruct Cheetah to only merge elements in between the passed elements.

NOTE: Transfer map merging can only be done for a constant incoming beam energy, because the transfer maps need to be computed before they can be merged, and computing them might require the beam energy at the entrance of the element that the transfer map belongs to. If you want to try a different beam energy, you will need to reapply the optimisations to the original lattice while passing a beam with the desired energy.

transfer_maps_merged_segment = original_segment.transfer_maps_merged(
    incoming_beam=incoming_beam, except_for=["HCOR_1", "VCOR_1", "HCOR_2", "VCOR_2"]
)
len(transfer_maps_merged_segment.elements)

8

%%timeit
transfer_maps_merged_segment.track(incoming_beam)

440 µs ± 1.9 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

transfer_maps_merged_segment

Segment(elements=ModuleList(
  (0): HorizontalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='HCOR_1')
  (1): Drift(length=tensor([0.3000]))
  (2): VerticalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='VCOR_1')
  (3): CustomTransferMap(transfer_map=tensor([[[ 1.3122e+00, -3.4577e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 1.8828e-01,  2.6594e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  3.0360e-01, -3.2559e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  1.8828e-01,  1.2746e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,
            -3.0678e-03,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             1.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  1.0000e+00]]]), length=tensor([135.2991]), name='combined_unnamed_element_1_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_6_unnamed_element_7_unnamed_element_8')
  (4): HorizontalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='HCOR_2')
  (5): Drift(length=tensor([0.3000]))
  (6): VerticalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='VCOR_2')
  (7): CustomTransferMap(transfer_map=tensor([[[ 1.0000e+00,  3.0000e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  3.0000e-01,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,
            -6.8021e-06,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             1.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  1.0000e+00]]]), length=tensor([0.3000]), name='combined_unnamed_element_10')
), name='unnamed_element_11')

It is also possible and often advisable to combine optimisations. However, note that this might not always yield as much of an improvement as one may have hoped looking at the improvements delivered by each optimisation on its own. This is usually because these optimisations share some of their effects, i.e. if the first optimisation has already performed a change on the lattice that the second optimisation would have done as well, the second optimisation will not lead to a further speed improvement.

fully_optimized_segment = (
    original_segment.without_inactive_markers()
    .inactive_elements_as_drifts(except_for=["HCOR_1", "VCOR_1", "HCOR_2", "VCOR_2"])
    .transfer_maps_merged(
        incoming_beam=incoming_beam, except_for=["HCOR_1", "VCOR_1", "HCOR_2", "VCOR_2"]
    )
)
len(fully_optimized_segment.elements)

8

fully_optimized_segment

Segment(elements=ModuleList(
  (0): HorizontalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='HCOR_1')
  (1): Drift(length=tensor([0.3000]))
  (2): VerticalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='VCOR_1')
  (3): CustomTransferMap(transfer_map=tensor([[[ 1.3122e+00, -3.4577e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 1.8828e-01,  2.6594e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  3.0360e-01, -3.2559e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  1.8828e-01,  1.2746e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,
            -3.0678e-03,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             1.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  1.0000e+00]]]), length=tensor([135.2991]), name='combined_unnamed_element_617_unnamed_element_2_unnamed_element_618_unnamed_element_4_unnamed_element_619_unnamed_element_620_unnamed_element_621_unnamed_element_2_unnamed_element_622_unnamed_element_4_unnamed_element_623_unnamed_element_624_unnamed_element_625_unnamed_element_2_unnamed_element_626_unnamed_element_4_unnamed_element_627_unnamed_element_628_unnamed_element_629_unnamed_element_2_unnamed_element_630_unnamed_element_4_unnamed_element_631_unnamed_element_632_unnamed_element_633_unnamed_element_2_unnamed_element_634_unnamed_element_4_unnamed_element_635_unnamed_element_636_unnamed_element_637_unnamed_element_2_unnamed_element_638_unnamed_element_4_unnamed_element_639_unnamed_element_640_unnamed_element_641_unnamed_element_2_unnamed_element_642_unnamed_element_4_unnamed_element_643_unnamed_element_644_unnamed_element_645_unnamed_element_2_unnamed_element_646_unnamed_element_4_unnamed_element_647_unnamed_element_648_unnamed_element_649_unnamed_element_2_unnamed_element_650_unnamed_element_4_unnamed_element_651_unnamed_element_652_unnamed_element_653_unnamed_element_2_unnamed_element_654_unnamed_element_4_unnamed_element_655_unnamed_element_656_unnamed_element_657_unnamed_element_2_unnamed_element_658_unnamed_element_4_unnamed_element_659_unnamed_element_660_unnamed_element_661_unnamed_element_2_unnamed_element_662_unnamed_element_4_unnamed_element_663_unnamed_element_664_unnamed_element_665_unnamed_element_2_unnamed_element_666_unnamed_element_4_unnamed_element_667_unnamed_element_668_unnamed_element_669_unnamed_element_2_unnamed_element_670_unnamed_element_4_unnamed_element_671_unnamed_element_672_unnamed_element_673_unnamed_element_2_unnamed_element_674_unnamed_element_4_unnamed_element_675_unnamed_element_676_unnamed_element_677_unnamed_element_2_unnamed_element_678_unnamed_element_4_unnamed_element_679_unnamed_element_680_unnamed_element_681_unnamed_element_2_unnamed_element_682_unnamed_element_4_unnamed_element_683_unnamed_element_684_unnamed_element_685_unnamed_element_2_unnamed_element_686_unnamed_element_4_unnamed_element_687_unnamed_element_688_unnamed_element_689_unnamed_element_2_unnamed_element_690_unnamed_element_4_unnamed_element_691_unnamed_element_692_unnamed_element_693_unnamed_element_2_unnamed_element_694_unnamed_element_4_unnamed_element_695_unnamed_element_696_unnamed_element_697_unnamed_element_2_unnamed_element_698_unnamed_element_4_unnamed_element_699_unnamed_element_700_unnamed_element_701_unnamed_element_2_unnamed_element_702_unnamed_element_4_unnamed_element_703_unnamed_element_704_unnamed_element_705_unnamed_element_2_unnamed_element_706_unnamed_element_4_unnamed_element_707_unnamed_element_708_unnamed_element_709_unnamed_element_2_unnamed_element_710_unnamed_element_4_unnamed_element_711_unnamed_element_712_unnamed_element_713_unnamed_element_2_unnamed_element_714_unnamed_element_4_unnamed_element_715_unnamed_element_716_unnamed_element_717_unnamed_element_2_unnamed_element_718_unnamed_element_4_unnamed_element_719_unnamed_element_720_unnamed_element_721_unnamed_element_2_unnamed_element_722_unnamed_element_4_unnamed_element_723_unnamed_element_724_unnamed_element_725_unnamed_element_2_unnamed_element_726_unnamed_element_4_unnamed_element_727_unnamed_element_728_unnamed_element_729_unnamed_element_2_unnamed_element_730_unnamed_element_4_unnamed_element_731_unnamed_element_732_unnamed_element_733_unnamed_element_2_unnamed_element_734_unnamed_element_4_unnamed_element_735_unnamed_element_736_unnamed_element_737_unnamed_element_2_unnamed_element_738_unnamed_element_4_unnamed_element_739_unnamed_element_740_unnamed_element_741_unnamed_element_2_unnamed_element_742_unnamed_element_4_unnamed_element_743_unnamed_element_744_unnamed_element_745_unnamed_element_2_unnamed_element_746_unnamed_element_4_unnamed_element_747_unnamed_element_748_unnamed_element_749_unnamed_element_2_unnamed_element_750_unnamed_element_4_unnamed_element_751_unnamed_element_752_unnamed_element_753_unnamed_element_2_unnamed_element_754_unnamed_element_4_unnamed_element_755_unnamed_element_756_unnamed_element_757_unnamed_element_2_unnamed_element_758_unnamed_element_4_unnamed_element_759_unnamed_element_760_unnamed_element_761_unnamed_element_2_unnamed_element_762_unnamed_element_4_unnamed_element_763_unnamed_element_764_unnamed_element_765_unnamed_element_2_unnamed_element_766_unnamed_element_4_unnamed_element_767_unnamed_element_768_unnamed_element_769_unnamed_element_2_unnamed_element_770_unnamed_element_4_unnamed_element_771_unnamed_element_772_unnamed_element_773_unnamed_element_2_unnamed_element_774_unnamed_element_4_unnamed_element_775_unnamed_element_776_unnamed_element_777_unnamed_element_2_unnamed_element_778_unnamed_element_4_unnamed_element_779_unnamed_element_780_unnamed_element_781_unnamed_element_2_unnamed_element_782_unnamed_element_4_unnamed_element_783_unnamed_element_784_unnamed_element_785_unnamed_element_2_unnamed_element_786_unnamed_element_4_unnamed_element_787_unnamed_element_788_unnamed_element_789_unnamed_element_2_unnamed_element_790_unnamed_element_4_unnamed_element_791_unnamed_element_792_unnamed_element_793_unnamed_element_2_unnamed_element_794_unnamed_element_4_unnamed_element_795_unnamed_element_796_unnamed_element_797_unnamed_element_2_unnamed_element_798_unnamed_element_4_unnamed_element_799_unnamed_element_800_unnamed_element_801_unnamed_element_2_unnamed_element_802_unnamed_element_4_unnamed_element_803_unnamed_element_804_unnamed_element_805_unnamed_element_2_unnamed_element_806_unnamed_element_4_unnamed_element_807_unnamed_element_808_unnamed_element_809_unnamed_element_2_unnamed_element_810_unnamed_element_4_unnamed_element_811_unnamed_element_812_unnamed_element_813_unnamed_element_2_unnamed_element_814_unnamed_element_4_unnamed_element_815_unnamed_element_816_unnamed_element_817_unnamed_element_2_unnamed_element_818_unnamed_element_4_unnamed_element_819_unnamed_element_820_unnamed_element_821_unnamed_element_2_unnamed_element_822_unnamed_element_4_unnamed_element_823_unnamed_element_824_unnamed_element_825_unnamed_element_2_unnamed_element_826_unnamed_element_4_unnamed_element_827_unnamed_element_828_unnamed_element_829_unnamed_element_2_unnamed_element_830_unnamed_element_4_unnamed_element_831_unnamed_element_832_unnamed_element_833_unnamed_element_2_unnamed_element_834_unnamed_element_4_unnamed_element_835_unnamed_element_836_unnamed_element_837_unnamed_element_2_unnamed_element_838_unnamed_element_4_unnamed_element_839_unnamed_element_840_unnamed_element_841_unnamed_element_2_unnamed_element_842_unnamed_element_4_unnamed_element_843_unnamed_element_844_unnamed_element_845_unnamed_element_2_unnamed_element_846_unnamed_element_4_unnamed_element_847_unnamed_element_848_unnamed_element_849_unnamed_element_2_unnamed_element_850_unnamed_element_4_unnamed_element_851_unnamed_element_852_unnamed_element_853_unnamed_element_2_unnamed_element_854_unnamed_element_4_unnamed_element_855_unnamed_element_856_unnamed_element_857_unnamed_element_2_unnamed_element_858_unnamed_element_4_unnamed_element_859_unnamed_element_860_unnamed_element_861_unnamed_element_2_unnamed_element_862_unnamed_element_4_unnamed_element_863_unnamed_element_864_unnamed_element_865_unnamed_element_2_unnamed_element_866_unnamed_element_4_unnamed_element_867_unnamed_element_868_unnamed_element_869_unnamed_element_2_unnamed_element_870_unnamed_element_4_unnamed_element_871_unnamed_element_872_unnamed_element_873_unnamed_element_2_unnamed_element_874_unnamed_element_4_unnamed_element_875_unnamed_element_876_unnamed_element_877_unnamed_element_2_unnamed_element_878_unnamed_element_4_unnamed_element_879_unnamed_element_880_unnamed_element_881_unnamed_element_2_unnamed_element_882_unnamed_element_4_unnamed_element_883_unnamed_element_884_unnamed_element_885_unnamed_element_2_unnamed_element_886_unnamed_element_4_unnamed_element_887_unnamed_element_888_unnamed_element_889_unnamed_element_2_unnamed_element_890_unnamed_element_4_unnamed_element_891_unnamed_element_892_unnamed_element_893_unnamed_element_2_unnamed_element_894_unnamed_element_4_unnamed_element_895_unnamed_element_896_unnamed_element_897_unnamed_element_2_unnamed_element_898_unnamed_element_4_unnamed_element_899_unnamed_element_900_unnamed_element_901_unnamed_element_2_unnamed_element_902_unnamed_element_4_unnamed_element_903_unnamed_element_904_unnamed_element_905_unnamed_element_2_unnamed_element_906_unnamed_element_4_unnamed_element_907_unnamed_element_908_unnamed_element_909_unnamed_element_2_unnamed_element_910_unnamed_element_4_unnamed_element_911_unnamed_element_912_unnamed_element_913_unnamed_element_2_unnamed_element_914_unnamed_element_4_unnamed_element_915_unnamed_element_916_unnamed_element_917_unnamed_element_2_unnamed_element_918_unnamed_element_4_unnamed_element_919_unnamed_element_920_unnamed_element_921_unnamed_element_2_unnamed_element_922_unnamed_element_4_unnamed_element_923_unnamed_element_924_unnamed_element_925_unnamed_element_2_unnamed_element_926_unnamed_element_4_unnamed_element_927_unnamed_element_928_unnamed_element_929_unnamed_element_2_unnamed_element_930_unnamed_element_4_unnamed_element_931_unnamed_element_932_unnamed_element_933_unnamed_element_2_unnamed_element_934_unnamed_element_4_unnamed_element_935_unnamed_element_936_unnamed_element_937_unnamed_element_2_unnamed_element_938_unnamed_element_4_unnamed_element_939_unnamed_element_940_unnamed_element_941_unnamed_element_2_unnamed_element_942_unnamed_element_4_unnamed_element_943_unnamed_element_944_unnamed_element_945_unnamed_element_2_unnamed_element_946_unnamed_element_4_unnamed_element_947_unnamed_element_948_unnamed_element_949_unnamed_element_2_unnamed_element_950_unnamed_element_4_unnamed_element_951_unnamed_element_952_unnamed_element_953_unnamed_element_2_unnamed_element_954_unnamed_element_4_unnamed_element_955_unnamed_element_956_unnamed_element_957_unnamed_element_2_unnamed_element_958_unnamed_element_4_unnamed_element_959_unnamed_element_960_unnamed_element_961_unnamed_element_2_unnamed_element_962_unnamed_element_4_unnamed_element_963_unnamed_element_964_unnamed_element_965_unnamed_element_2_unnamed_element_966_unnamed_element_4_unnamed_element_967_unnamed_element_968_unnamed_element_969_unnamed_element_2_unnamed_element_970_unnamed_element_4_unnamed_element_971_unnamed_element_972_unnamed_element_973_unnamed_element_2_unnamed_element_974_unnamed_element_4_unnamed_element_975_unnamed_element_976_unnamed_element_977_unnamed_element_2_unnamed_element_978_unnamed_element_4_unnamed_element_979_unnamed_element_980_unnamed_element_981_unnamed_element_2_unnamed_element_982_unnamed_element_4_unnamed_element_983_unnamed_element_984_unnamed_element_985_unnamed_element_2_unnamed_element_986_unnamed_element_4_unnamed_element_987_unnamed_element_988_unnamed_element_989_unnamed_element_2_unnamed_element_990_unnamed_element_4_unnamed_element_991_unnamed_element_992_unnamed_element_993_unnamed_element_2_unnamed_element_994_unnamed_element_4_unnamed_element_995_unnamed_element_996_unnamed_element_997_unnamed_element_2_unnamed_element_998_unnamed_element_4_unnamed_element_999_unnamed_element_1000_unnamed_element_1001_unnamed_element_2_unnamed_element_1002_unnamed_element_4_unnamed_element_1003_unnamed_element_1004_unnamed_element_1005_unnamed_element_2_unnamed_element_1006_unnamed_element_4_unnamed_element_1007_unnamed_element_1008_unnamed_element_1009_unnamed_element_2_unnamed_element_1010_unnamed_element_4_unnamed_element_1011_unnamed_element_1012_unnamed_element_1013_unnamed_element_2_unnamed_element_1014_unnamed_element_4_unnamed_element_1015_unnamed_element_1016_unnamed_element_1017_unnamed_element_2_unnamed_element_1018_unnamed_element_4_unnamed_element_1019_unnamed_element_1020_unnamed_element_1021_unnamed_element_2_unnamed_element_1022_unnamed_element_4_unnamed_element_1023_unnamed_element_1024_unnamed_element_1025_unnamed_element_2_unnamed_element_1026_unnamed_element_4_unnamed_element_1027_unnamed_element_1028_unnamed_element_1029_unnamed_element_2_unnamed_element_1030_unnamed_element_4_unnamed_element_1031_unnamed_element_1032_unnamed_element_1033_unnamed_element_2_unnamed_element_1034_unnamed_element_4_unnamed_element_1035_unnamed_element_1036_unnamed_element_1037_unnamed_element_2_unnamed_element_1038_unnamed_element_4_unnamed_element_1039_unnamed_element_1040_unnamed_element_1041_unnamed_element_2_unnamed_element_1042_unnamed_element_4_unnamed_element_1043_unnamed_element_1044_unnamed_element_1045_unnamed_element_2_unnamed_element_1046_unnamed_element_4_unnamed_element_1047_unnamed_element_1048_unnamed_element_1049_unnamed_element_2_unnamed_element_1050_unnamed_element_4_unnamed_element_1051_unnamed_element_1052_unnamed_element_1053_unnamed_element_2_unnamed_element_1054_unnamed_element_4_unnamed_element_1055_unnamed_element_1056_unnamed_element_1057_unnamed_element_2_unnamed_element_1058_unnamed_element_4_unnamed_element_1059_unnamed_element_1060_unnamed_element_1061_unnamed_element_2_unnamed_element_1062_unnamed_element_4_unnamed_element_1063_unnamed_element_1064_unnamed_element_1065_unnamed_element_2_unnamed_element_1066_unnamed_element_4_unnamed_element_1067_unnamed_element_1068_unnamed_element_1069_unnamed_element_2_unnamed_element_1070_unnamed_element_4_unnamed_element_1071_unnamed_element_1072_unnamed_element_1073_unnamed_element_2_unnamed_element_1074_unnamed_element_4_unnamed_element_1075_unnamed_element_1076_unnamed_element_1077_unnamed_element_2_unnamed_element_1078_unnamed_element_4_unnamed_element_1079_unnamed_element_1080_unnamed_element_1081_unnamed_element_2_unnamed_element_1082_unnamed_element_4_unnamed_element_1083_unnamed_element_1084_unnamed_element_1085_unnamed_element_2_unnamed_element_1086_unnamed_element_4_unnamed_element_1087_unnamed_element_1088_unnamed_element_1089_unnamed_element_2_unnamed_element_1090_unnamed_element_4_unnamed_element_1091_unnamed_element_1092_unnamed_element_1093_unnamed_element_2_unnamed_element_1094_unnamed_element_4_unnamed_element_1095_unnamed_element_1096_unnamed_element_1097_unnamed_element_2_unnamed_element_1098_unnamed_element_4_unnamed_element_1099_unnamed_element_1100_unnamed_element_1101_unnamed_element_2_unnamed_element_1102_unnamed_element_4_unnamed_element_1103_unnamed_element_1104_unnamed_element_1105_unnamed_element_2_unnamed_element_1106_unnamed_element_4_unnamed_element_1107_unnamed_element_1108_unnamed_element_1109_unnamed_element_2_unnamed_element_1110_unnamed_element_4_unnamed_element_1111_unnamed_element_1112_unnamed_element_1113_unnamed_element_2_unnamed_element_1114_unnamed_element_4_unnamed_element_1115_unnamed_element_1116_unnamed_element_1117_unnamed_element_2_unnamed_element_1118_unnamed_element_4_unnamed_element_1119_unnamed_element_1120_unnamed_element_1121_unnamed_element_2_unnamed_element_1122_unnamed_element_4_unnamed_element_1123_unnamed_element_1124_unnamed_element_1125_unnamed_element_2_unnamed_element_1126_unnamed_element_4_unnamed_element_1127_unnamed_element_1128_unnamed_element_1129_unnamed_element_2_unnamed_element_1130_unnamed_element_4_unnamed_element_1131_unnamed_element_1132_unnamed_element_1133_unnamed_element_2_unnamed_element_1134_unnamed_element_4_unnamed_element_1135_unnamed_element_1136_unnamed_element_1137_unnamed_element_2_unnamed_element_1138_unnamed_element_4_unnamed_element_1139_unnamed_element_1140_unnamed_element_1141_unnamed_element_2_unnamed_element_1142_unnamed_element_4_unnamed_element_1143_unnamed_element_1144_unnamed_element_1145_unnamed_element_2_unnamed_element_1146_unnamed_element_4_unnamed_element_1147_unnamed_element_1148_unnamed_element_1149_unnamed_element_2_unnamed_element_1150_unnamed_element_4_unnamed_element_1151_unnamed_element_1152_unnamed_element_1153_unnamed_element_2_unnamed_element_1154_unnamed_element_4_unnamed_element_1155_unnamed_element_1156_unnamed_element_1157_unnamed_element_2_unnamed_element_1158_unnamed_element_4_unnamed_element_1159_unnamed_element_1160_unnamed_element_1161_unnamed_element_2_unnamed_element_1162_unnamed_element_4_unnamed_element_1163_unnamed_element_1164_unnamed_element_1165_unnamed_element_2_unnamed_element_1166_unnamed_element_4_unnamed_element_1167_unnamed_element_1168_unnamed_element_1169_unnamed_element_2_unnamed_element_1170_unnamed_element_4_unnamed_element_1171_unnamed_element_1172_unnamed_element_1173_unnamed_element_2_unnamed_element_1174_unnamed_element_4_unnamed_element_1175_unnamed_element_1176_unnamed_element_1177_unnamed_element_2_unnamed_element_1178_unnamed_element_4_unnamed_element_1179_unnamed_element_1180_unnamed_element_1181_unnamed_element_2_unnamed_element_1182_unnamed_element_4_unnamed_element_1183_unnamed_element_1184_unnamed_element_1185_unnamed_element_2_unnamed_element_1186_unnamed_element_4_unnamed_element_1187_unnamed_element_1188_unnamed_element_1189_unnamed_element_2_unnamed_element_1190_unnamed_element_4_unnamed_element_1191_unnamed_element_1192_unnamed_element_1193_unnamed_element_2_unnamed_element_1194_unnamed_element_4_unnamed_element_1195_unnamed_element_1196_unnamed_element_1197_unnamed_element_2_unnamed_element_1198_unnamed_element_4_unnamed_element_1199_unnamed_element_1200_unnamed_element_1201_unnamed_element_2_unnamed_element_1202_unnamed_element_4_unnamed_element_1203_unnamed_element_1204_unnamed_element_1205_unnamed_element_2_unnamed_element_1206_unnamed_element_4_unnamed_element_1207_unnamed_element_1208_unnamed_element_1209_unnamed_element_2_unnamed_element_1210_unnamed_element_4_unnamed_element_1211_unnamed_element_1212_unnamed_element_1213_unnamed_element_2_unnamed_element_1214_unnamed_element_4_unnamed_element_1215_unnamed_element_1216_unnamed_element_1217')
  (4): HorizontalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='HCOR_2')
  (5): Drift(length=tensor([0.3000]))
  (6): VerticalCorrector(length=tensor([0.1000]), angle=tensor([0.]), name='VCOR_2')
  (7): CustomTransferMap(transfer_map=tensor([[[ 1.0000e+00,  3.0000e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  1.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  1.0000e+00,  3.0000e-01,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,  0.0000e+00,
             0.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  1.0000e+00,
            -6.8021e-06,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             1.0000e+00,  0.0000e+00],
           [ 0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
             0.0000e+00,  1.0000e+00]]]), length=tensor([0.3000]), name='combined_unnamed_element_1219')
), name='unnamed_element_11')

%%timeit
fully_optimized_segment.track(incoming_beam)

443 µs ± 805 ns per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Especially in the context of machine learning, it is quite common to need to run multiple simulations over the same lattice with slight changed lattice or beam parameters. Cheetah is ideal for this use case because in can run vectorised simulations, meaning multiple simulations are run concurrently. This can lead to significant further speed improvements.

Let’s assume we want to run 1000 simulations rather than just one on the fully optimised lattice. We can use the broadcast method of both the lattice and the incoming beam to vectorise them to any number of concurrent simulations (1000 in our case). When running a simulation with the vectorised objects, we see that the time Cheetah takes to run the simulation increases. But keep in mind that we are doing 1000-times the work in that time and that for each sample the computation time has actually decreased significantly.

broadcast_fully_optimized_segment = fully_optimized_segment.broadcast((1_000,))
broadcast_incoming_beam = incoming_beam.broadcast((1_000,))

%%timeit
broadcast_fully_optimized_segment.track(broadcast_incoming_beam)

1.9 ms ± 18.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)