Converting lattices from other simulation codes

In this example, we demonstrate how to convert lattices from other simulation codes. At the moment, Cheetah supports the conversion of lattices Ocelot and Bmad.

import ocelot
import torch

from cheetah import Segment

math_op.py: module Numba is not installed. Install it if you want speed up correlation calculations

[INFO    ] : : beam.py: module NUMBA is not installed. Install it to speed up calculation
[INFO    ] : : : : : : : : high_order.py: module NUMBA is not installed. Install it to speed up calculation
[INFO    ] radiation_py.py: module NUMBA is not installed. Install it to speed up calculation
[INFO    ] radiation_py.py: module NUMBA is not installed. Install it to speed up calculation
[INFO    ] csr.py: module NUMBA is not installed. Install it to speed up calculation
[INFO    ] csr.py: module PYFFTW is not installed. Install it to speed up calculation.
[INFO    ] csr.py: module NUMEXPR is not installed. Install it to speed up calculation
[INFO    ] wake3D.py: module NUMBA is not installed. Install it to speed up calculation

initializing ocelot...
import: module NUMBA is not installed. Install it to speed up calculation
import: module PYFFTW is not installed. Install it to speed up calculation
import: module NUMEXPR is not installed. Install it to speed up calculation

Lattice conversions can conveniently be done using class methods defined by the Segment class.

To convert an Ocelot cell that is stored as a Python variable, simply pass it to Segment.from_ocelot():

ocelot_cell = [
    ocelot.Drift(l=1.0),
    ocelot.Quadrupole(l=0.2, k1=4.2),
    ocelot.Drift(l=1.0),
]

ocelot_converted = Segment.from_ocelot(ocelot_cell)
ocelot_converted

Segment(elements=ModuleList(
  (0): Drift(length=tensor([1.]))
  (1): Quadrupole(length=tensor([0.2000]), k1=tensor([4.2000]), misalignment=tensor([[0., 0.]]), tilt=tensor([0.]), name='ID_42769541_')
  (2): Drift(length=tensor([1.]))
), name='unnamed_element_0')

Bmad on the other are read from .bmad files. To convert the following Bmad lattice

!cat ../../tests/resources/bmad_tutorial_lattice.bmad

! Lattice file: simple.bmad
beginning[beta_a] = 10. ! m a-mode beta function
beginning[beta_b] = 10. ! m b-mode beta function
beginning[e_tot] = 10e6 ! eV    Or can set beginning[p0c]

parameter[geometry] = open  ! Or closed
parameter[particle] = electron  ! Reference particle.

d: drift, L = 0.5
b: sbend, L = 0.5, g = 1, e1 = 0.1, dg = 0.001 ! g = 1/design_radius
q: quadrupole, L = 0.6, k1 = 0.23

lat: line = (d, b, q) ! List of lattice elements
use, lat ! Line used to construct the lattice

, pass the file path to Segment.from_bmad().

bmad_converted = Segment.from_bmad("../../tests/resources/bmad_tutorial_lattice.bmad")
bmad_converted

Segment(elements=ModuleList(
  (0): Drift(length=tensor([0.5000]))
  (1): Dipole(length=tensor([0.5000]), angle=tensor([0.]), e1=tensor([0.1000]),e2=tensor([0.]),tilt=tensor([0.]),fringe_integral=tensor([0.]),fringe_integral_exit=tensor([0.]),gap=tensor([0.]),name='b')
  (2): Quadrupole(length=tensor([0.6000]), k1=tensor([0.2300]), misalignment=tensor([[0., 0.]]), tilt=tensor([0.]), name='q')
), name='lat')