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. As this notebook demonstrates, speed improvements of more than 5 orders of magnitude can be achieved.

import torch

import cheetah

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)

80.5 ms ± 572 µs per loop (mean ± std. dev. of 7 runs, 1 loop 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)

79.1 ms ± 368 µ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)

65.3 ms ± 181 µ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)

318 µs ± 1.42 µ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), tracking_method='cheetah', name='unnamed_element_0')
  (2): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_1')
  (3): CustomTransferMap(predefined_transfer_map=tensor([[ 1.3122e+00, -3.4576e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [ 1.8828e-01,  2.6597e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [ 0.0000e+00,  0.0000e+00,  3.0363e-01, -3.2558e+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), tracking_method='cheetah', name='unnamed_element_9')
  (6): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_2')
  (7): CustomTransferMap(predefined_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), tracking_method='cheetah', name='unnamed_element_0')
  (2): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_1')
  (3): CustomTransferMap(predefined_transfer_map=tensor([[ 1.3122e+00, -3.4576e+00,  0.0000e+00,  0.0000e+00,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [ 1.8828e-01,  2.6597e-01,  0.0000e+00,  0.0000e+00,  0.0000e+00,
            0.0000e+00,  0.0000e+00],
          [ 0.0000e+00,  0.0000e+00,  3.0363e-01, -3.2558e+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_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8_unnamed_element_2_unnamed_element_3_unnamed_element_4_unnamed_element_5_unnamed_element_7_unnamed_element_8')
  (4): HorizontalCorrector(length=tensor(0.1000), angle=tensor(0.), name='HCOR_2')
  (5): Drift(length=tensor(0.3000), tracking_method='cheetah', name='unnamed_element_9')
  (6): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_2')
  (7): CustomTransferMap(predefined_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')

%%timeit
_ = fully_optimized_segment.track(incoming_beam)

318 µs ± 2.23 µs 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 simulations with 1000 different corrector angles. Without Cheetah’s vectorisation, one would have to run 1000 simulations one after the other in a loop. With Cheetah’s vectorisation, we can run all 1000 simulations at the same time. This simplifies the code and, in the example below, speeds up the simulation by a factor of over 400. Note that this example is run on CPU, but the speedup should be even more significant on a GPU and with larger numbers of simulations.

%%timeit
for angle in torch.linspace(-1e-4, 1e-4, 1_000):
    fully_optimized_segment.HCOR_1.angle = angle
    _ = fully_optimized_segment.track(incoming_beam)

349 ms ± 2.73 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%%timeit
fully_optimized_segment.HCOR_1.angle = torch.linspace(-1e-4, 1e-4, 1_000)
_ = fully_optimized_segment.track(incoming_beam)

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

Just-in-Time (JIT) Compiling Cheetah

PyTorch supports to compile code to optimize performance. We observed that JIT compilation provides a 10-20% speedup on AMD CPUs, a 0.5-2x speedup on Intel CPUs, and up to 8x speedup on Nvidia GPUs [`HueblNAPAC25 <>`__].

JIT compilation will be performed on the first call to a function and can take around 4-20 seconds on CPU and more on GPU, depending on the lengths of a segment. Thus, it is most useful to JIT compile when calling a function multiple times.

compiled_track = torch.compile(transfer_maps_merged_segment.track)

%%timeit
_ = compiled_track(incoming_beam)

The slowest run took 4.16 times longer than the fastest. This could mean that an intermediate result is being cached.
1.31 ms ± 815 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)

Any later call will be faster

%%timeit
_ = compiled_track(incoming_beam)

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

Notes: By default, PyTorch compiles with fast-math optimizations for some platforms. If you observe issues with precision from compiling, it is worth experimenting with these options before compile:

# C++ backend (default: False)
torch._inductor.config.cpp.enable_unsafe_math_opt_flag = False
# CUDA backend (default: False)
torch._inductor.config.cuda.use_fast_math = False
# ROCm backend (default: True)
torch._inductor.config.rocm.use_fast_math = False