{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Converting lattices from other simulation codes\n",
    "\n",
    "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_.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "math_op.py: module Numba is not installed. Install it if you want speed up correlation calculations\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[INFO    ] : : \u001b[0mbeam.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] : : : : : : : : \u001b[0mhigh_order.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] \u001b[0mradiation_py.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] \u001b[0mradiation_py.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] \u001b[0mcsr.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] \u001b[0mcsr.py: module PYFFTW is not installed. Install it to speed up calculation.\u001b[0m\n",
      "[INFO    ] \u001b[0mcsr.py: module NUMEXPR is not installed. Install it to speed up calculation\u001b[0m\n",
      "[INFO    ] \u001b[0mwake3D.py: module NUMBA is not installed. Install it to speed up calculation\u001b[0m\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "initializing ocelot...\n",
      "import: module NUMBA is not installed. Install it to speed up calculation\n",
      "import: module PYFFTW is not installed. Install it to speed up calculation\n",
      "import: module NUMEXPR is not installed. Install it to speed up calculation\n"
     ]
    }
   ],
   "source": [
    "import ocelot\n",
    "\n",
    "import cheetah"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lattice conversions can conveniently be done using class methods defined by the\n",
    "`Segment` class.\n",
    "\n",
    "To convert an _Ocelot_ cell that is stored as a Python variable, simply pass it to\n",
    "`cheetah.Segment.from_ocelot()`:\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Segment(elements=ModuleList(\n",
       "  (0): Drift(name='ID_22845104_', tracking_method='linear', length=tensor(1.))\n",
       "  (1): Quadrupole(name='ID_45190360_', tracking_method='linear', length=tensor(0.2000), k1=tensor(4.2000), misalignment=tensor([0., 0.]), tilt=tensor(0.), num_steps=1)\n",
       "  (2): Drift(name='ID_44014220_', tracking_method='linear', length=tensor(1.))\n",
       "), name='unnamed_element_0')"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ocelot_cell = [\n",
    "    ocelot.Drift(l=1.0),\n",
    "    ocelot.Quadrupole(l=0.2, k1=4.2),\n",
    "    ocelot.Drift(l=1.0),\n",
    "]\n",
    "\n",
    "ocelot_converted = cheetah.Segment.from_ocelot(ocelot_cell)\n",
    "ocelot_converted"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "_Bmad_ on the other are read from `.bmad` files. To convert the following _Bmad_ lattice\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "! Lattice file: simple.bmad\n",
      "beginning[beta_a] = 10. ! m a-mode beta function\n",
      "beginning[beta_b] = 10. ! m b-mode beta function\n",
      "beginning[e_tot] = 10e6 ! eV    Or can set beginning[p0c]\n",
      "\n",
      "parameter[geometry] = open  ! Or closed\n",
      "parameter[particle] = electron  ! Reference particle.\n",
      "\n",
      "! Bmad lattices treat variable and function names separately. This is to test that.\n",
      "abs = -0.6\n",
      "\n",
      "d: drift, L = 0.5 * (0.3 + 0.7)\n",
      "! The two spaces in the expression for L are on purpose to test the parser's ability to\n",
      "! handle this\n",
      "b: sbend, L = 0.6  -0.1, g = 1, e1 = 0.1, dg = sqrt(0.000001) ! g = 1/design_radius\n",
      "n: drift, L=-0.4\n",
      "q: quadrupole, L = abs(abs), k1 = 0.23\n",
      "s: sextupole, tilt = -0.1, L = 0.3, k2 = 0.42\n",
      "v: drift, l = -q[l]\n",
      "\n",
      "lat: line = (d, b, n, q, s, v) ! List of lattice elements\n",
      "use, lat ! Line used to construct the lattice\n"
     ]
    }
   ],
   "source": [
    "!cat ../../tests/resources/bmad_tutorial_lattice.bmad"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ", pass the file path to `cheetah.Segment.from_bmad()`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/jankaiser/Documents/DESY/cheetah/cheetah/converters/bmad.py:113: NotUnderstoodPropertyWarning: Property g with value 1 for element type sbend is currently not understood.\n",
      "  validate_understood_properties(\n",
      "/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.\n",
      "  validate_understood_properties(\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Segment(elements=ModuleList(\n",
       "  (0): Drift(name='d', tracking_method='linear', length=tensor(0.5000))\n",
       "  (1): Dipole(name='b', tracking_method='linear', 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')\n",
       "   ⋮\n",
       "  (4): Sextupole(name='s', tracking_method='second_order', length=tensor(0.3000), k2=tensor(0.4200), misalignment=tensor([0., 0.]), tilt=tensor(-0.1000))\n",
       "  (5): Drift(name='v', tracking_method='linear', length=tensor(-0.6000))\n",
       "), name='lat')"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bmad_converted = cheetah.Segment.from_bmad(\n",
    "    \"../../tests/resources/bmad_tutorial_lattice.bmad\"\n",
    ")\n",
    "bmad_converted"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cheetah-dev",
   "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.14.0"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}