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 cheetah

[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

initializing ocelot...
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 cheetah.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 = cheetah.Segment.from_ocelot(ocelot_cell)
ocelot_converted

Segment(elements=ModuleList(
  (0): Drift(length=tensor(1.), tracking_method='cheetah', name='ID_22162625_')
  (1): Quadrupole(length=tensor(0.2000), k1=tensor(4.2000), misalignment=tensor([0., 0.]), tilt=tensor(0.), num_steps=1, tracking_method='cheetah', name='ID_41789612_')
  (2): Drift(length=tensor(1.), tracking_method='cheetah', name='ID_55514304_')
), 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.

! Bmad lattices treat variable and function names separately. This is to test that.
abs = -0.6

d: drift, L = 0.5 * (0.3 + 0.7)
! The two spaces in the expression for L are on purpose to test the parser's ability to
! handle this
b: sbend, L = 0.6  -0.1, g = 1, e1 = 0.1, dg = sqrt(0.000001) ! g = 1/design_radius
q: quadrupole, L = abs(abs), k1 = 0.23
s: sextupole, tilt = -0.1, L = 0.3, k2 = 0.42

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

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

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

/Users/jankaiser/Documents/DESY/cheetah/cheetah/converters/bmad.py:113: NotUnderstoodPropertyWarning: Property g with value 1 for element type sbend is currently not understood.
  validate_understood_properties(
/Users/jankaiser/Documents/DESY/cheetah/cheetah/converters/bmad.py:113: NotUnderstoodPropertyWarning: Property dg with value 0.001 for element type sbend is currently not understood.
  validate_understood_properties(

Segment(elements=ModuleList(
  (0): Drift(length=tensor(0.5000), tracking_method='cheetah', name='d')
  (1): Dipole(length=tensor(0.5000), angle=tensor(0.), k1=tensor(0.), dipole_e1=tensor(0.1000),dipole_e2=tensor(0.),tilt=tensor(0.),gap=tensor(0.),gap_exit=tensor(0.),fringe_integral=tensor(0.),fringe_integral_exit=tensor(0.),fringe_at='both',fringe_type='linear_edge',tracking_method='cheetah', name='b')
  (2): Quadrupole(length=tensor(0.6000), k1=tensor(0.2300), misalignment=tensor([0., 0.]), tilt=tensor(0.), num_steps=1, tracking_method='cheetah', name='q')
  (3): Sextupole(length=tensor(0.3000), k2=tensor(0.4200), misalignment=tensor([0., 0.]), tilt=tensor(-0.1000), name='s')
), name='lat')