{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Optimising _Cheetah_ for speed\n", "\n", "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.\n", "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.\n" ] }, { "cell_type": "code", "execution_count": 1, "metadata": {}, "outputs": [], "source": [ "import torch\n", "\n", "import cheetah" ] }, { "cell_type": "code", "execution_count": 2, "metadata": { "jupyter": { "source_hidden": true } }, "outputs": [], "source": [ "incoming_beam = cheetah.ParameterBeam.from_astra(\n", " \"../../tests/resources/ACHIP_EA1_2021.1351.001\"\n", ")" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n" ] }, { "cell_type": "code", "execution_count": 3, "metadata": {}, "outputs": [], "source": [ "original_segment = cheetah.Segment(\n", " elements=[\n", " cheetah.HorizontalCorrector(\n", " length=torch.tensor(0.1), angle=torch.tensor(0.0), name=\"HCOR_1\"\n", " ),\n", " cheetah.Drift(length=torch.tensor(0.3)),\n", " cheetah.VerticalCorrector(\n", " length=torch.tensor(0.1), angle=torch.tensor(0.0), name=\"VCOR_1\"\n", " ),\n", " cheetah.Drift(length=torch.tensor(0.3)),\n", " ]\n", " + [\n", " cheetah.Quadrupole(length=torch.tensor(0.1), k1=torch.tensor(4.2)),\n", " cheetah.Drift(length=torch.tensor(0.2)),\n", " cheetah.Quadrupole(length=torch.tensor(0.1), k1=torch.tensor(-4.2)),\n", " cheetah.Drift(length=torch.tensor(0.2)),\n", " cheetah.Marker(),\n", " cheetah.Quadrupole(length=torch.tensor(0.1), k1=torch.tensor(0.0)),\n", " cheetah.Drift(length=torch.tensor(0.2)),\n", " ]\n", " * 150\n", " + [\n", " cheetah.HorizontalCorrector(\n", " length=torch.tensor(0.1), angle=torch.tensor(0.0), name=\"HCOR_2\"\n", " ),\n", " cheetah.Drift(length=torch.tensor(0.3)),\n", " cheetah.VerticalCorrector(\n", " length=torch.tensor(0.1), angle=torch.tensor(0.0), name=\"VCOR_2\"\n", " ),\n", " cheetah.Drift(length=torch.tensor(0.3)),\n", " ]\n", ")" ] }, { "cell_type": "code", "execution_count": 4, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "1058" ] }, "execution_count": 4, "metadata": {}, "output_type": "execute_result" } ], "source": [ "len(original_segment.elements)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n" ] }, { "cell_type": "code", "execution_count": 5, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "80.5 ms ± 572 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "_ = original_segment.track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Just by removing unused markers, we already see a small performance improvement.\n" ] }, { "cell_type": "code", "execution_count": 6, "metadata": {}, "outputs": [], "source": [ "markers_removed_segment = original_segment.without_inactive_markers()" ] }, { "cell_type": "code", "execution_count": 7, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "79.1 ms ± 368 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "%%timeit\n", "_ = markers_removed_segment.track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n" ] }, { "cell_type": "code", "execution_count": 8, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "1058" ] }, "execution_count": 8, "metadata": {}, "output_type": "execute_result" } ], "source": [ "inactive_to_drifts_segment = original_segment.inactive_elements_as_drifts(\n", " except_for=[\"HCOR_1\", \"VCOR_1\", \"HCOR_2\", \"VCOR_2\"]\n", ")\n", "len(inactive_to_drifts_segment.elements)" ] }, { "cell_type": "code", "execution_count": 9, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "65.3 ms ± 181 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)\n" ] } ], "source": [ "%%timeit\n", "_ = inactive_to_drifts_segment.track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n", "\n", "**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.\n" ] }, { "cell_type": "code", "execution_count": 10, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "8" ] }, "execution_count": 10, "metadata": {}, "output_type": "execute_result" } ], "source": [ "transfer_maps_merged_segment = original_segment.transfer_maps_merged(\n", " incoming_beam=incoming_beam, except_for=[\"HCOR_1\", \"VCOR_1\", \"HCOR_2\", \"VCOR_2\"]\n", ")\n", "len(transfer_maps_merged_segment.elements)" ] }, { "cell_type": "code", "execution_count": 11, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "318 µs ± 1.42 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n" ] } ], "source": [ "%%timeit\n", "_ = transfer_maps_merged_segment.track(incoming_beam)" ] }, { "cell_type": "code", "execution_count": 12, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Segment(elements=ModuleList(\n", " (0): HorizontalCorrector(length=tensor(0.1000), angle=tensor(0.), name='HCOR_1')\n", " (1): Drift(length=tensor(0.3000), tracking_method='cheetah', name='unnamed_element_0')\n", " (2): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_1')\n", " (3): CustomTransferMap(predefined_transfer_map=tensor([[ 1.3122e+00, -3.4576e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 1.8828e-01, 2.6597e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 3.0363e-01, -3.2558e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 1.8828e-01, 1.2746e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,\n", " -3.0678e-03, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 1.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 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')\n", " (4): HorizontalCorrector(length=tensor(0.1000), angle=tensor(0.), name='HCOR_2')\n", " (5): Drift(length=tensor(0.3000), tracking_method='cheetah', name='unnamed_element_9')\n", " (6): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_2')\n", " (7): CustomTransferMap(predefined_transfer_map=tensor([[ 1.0000e+00, 3.0000e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 1.0000e+00, 3.0000e-01, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,\n", " -6.8021e-06, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 1.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 1.0000e+00]]), length=tensor(0.3000), name='combined_unnamed_element_10')\n", "), name='unnamed_element_11')" ] }, "execution_count": 12, "metadata": {}, "output_type": "execute_result" } ], "source": [ "transfer_maps_merged_segment" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n" ] }, { "cell_type": "code", "execution_count": 13, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "8" ] }, "execution_count": 13, "metadata": {}, "output_type": "execute_result" } ], "source": [ "fully_optimized_segment = (\n", " original_segment.without_inactive_markers()\n", " .inactive_elements_as_drifts(except_for=[\"HCOR_1\", \"VCOR_1\", \"HCOR_2\", \"VCOR_2\"])\n", " .transfer_maps_merged(\n", " incoming_beam=incoming_beam, except_for=[\"HCOR_1\", \"VCOR_1\", \"HCOR_2\", \"VCOR_2\"]\n", " )\n", ")\n", "len(fully_optimized_segment.elements)" ] }, { "cell_type": "code", "execution_count": 14, "metadata": {}, "outputs": [ { "data": { "text/plain": [ "Segment(elements=ModuleList(\n", " (0): HorizontalCorrector(length=tensor(0.1000), angle=tensor(0.), name='HCOR_1')\n", " (1): Drift(length=tensor(0.3000), tracking_method='cheetah', name='unnamed_element_0')\n", " (2): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_1')\n", " (3): CustomTransferMap(predefined_transfer_map=tensor([[ 1.3122e+00, -3.4576e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 1.8828e-01, 2.6597e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 3.0363e-01, -3.2558e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 1.8828e-01, 1.2746e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,\n", " -3.0678e-03, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 1.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 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')\n", " (4): HorizontalCorrector(length=tensor(0.1000), angle=tensor(0.), name='HCOR_2')\n", " (5): Drift(length=tensor(0.3000), tracking_method='cheetah', name='unnamed_element_9')\n", " (6): VerticalCorrector(length=tensor(0.1000), angle=tensor(0.), name='VCOR_2')\n", " (7): CustomTransferMap(predefined_transfer_map=tensor([[ 1.0000e+00, 3.0000e-01, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 1.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 1.0000e+00, 3.0000e-01, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 1.0000e+00,\n", " -6.8021e-06, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 1.0000e+00, 0.0000e+00],\n", " [ 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n", " 0.0000e+00, 1.0000e+00]]), length=tensor(0.3000), name='combined_unnamed_element_10')\n", "), name='unnamed_element_11')" ] }, "execution_count": 14, "metadata": {}, "output_type": "execute_result" } ], "source": [ "fully_optimized_segment" ] }, { "cell_type": "code", "execution_count": 15, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "318 µs ± 2.23 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n" ] } ], "source": [ "%%timeit\n", "_ = fully_optimized_segment.track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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.\n", "\n", "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.\n" ] }, { "cell_type": "code", "execution_count": 16, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "349 ms ± 2.73 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "for angle in torch.linspace(-1e-4, 1e-4, 1_000):\n", " fully_optimized_segment.HCOR_1.angle = angle\n", " _ = fully_optimized_segment.track(incoming_beam)" ] }, { "cell_type": "code", "execution_count": 17, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "1.06 ms ± 12.2 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n" ] } ], "source": [ "%%timeit\n", "fully_optimized_segment.HCOR_1.angle = torch.linspace(-1e-4, 1e-4, 1_000)\n", "_ = fully_optimized_segment.track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Just-in-Time (JIT) Compiling Cheetah\n", "\n", "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]()].\n", "\n", "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." ] }, { "cell_type": "code", "execution_count": 18, "metadata": {}, "outputs": [], "source": [ "compiled_track = torch.compile(transfer_maps_merged_segment.track)" ] }, { "cell_type": "code", "execution_count": 19, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "The slowest run took 4.16 times longer than the fastest. This could mean that an intermediate result is being cached.\n", "1.31 ms ± 815 µs per loop (mean ± std. dev. of 7 runs, 1 loop each)\n" ] } ], "source": [ "%%timeit\n", "_ = compiled_track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "Any later call will be faster" ] }, { "cell_type": "code", "execution_count": 20, "metadata": {}, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "753 µs ± 30.4 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n" ] } ], "source": [ "%%timeit\n", "_ = compiled_track(incoming_beam)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "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:" ] }, { "cell_type": "code", "execution_count": 21, "metadata": {}, "outputs": [], "source": [ "# C++ backend (default: False)\n", "torch._inductor.config.cpp.enable_unsafe_math_opt_flag = False\n", "# CUDA backend (default: False)\n", "torch._inductor.config.cuda.use_fast_math = False\n", "# ROCm backend (default: True)\n", "torch._inductor.config.rocm.use_fast_math = False" ] } ], "metadata": { "kernelspec": { "display_name": "Python 3 (ipykernel)", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.10.18" } }, "nbformat": 4, "nbformat_minor": 4 }