{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "8ae07586-3342-4386-9f7c-2a32c60e510e",
   "metadata": {},
   "source": [
    "# Example-21: Build MADX and ELEGANT style lattice"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2f5561be-7da6-499a-aac4-371570ceae6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# In this example it is demonstraited how to load MADX and ELEGANT lattice files\n",
    "\n",
    "# Supported elements (other elemet types will be casted to drifts)\n",
    "\n",
    "# Drift\n",
    "# Quadrupole\n",
    "# Sextupole\n",
    "# Octupole\n",
    "# Dipole\n",
    "# BPM\n",
    "# Marker\n",
    "# Line\n",
    "\n",
    "# Note, only basic options are translated (not translated options are ignored)\n",
    "# Note, reference orbit slicing is not correct with negative bending"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "bb620c55-6c18-4bb4-a5a8-32fab36bd695",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from model.command.external import load_lattice\n",
    "from model.command.build import build\n",
    "from model.command.layout import Layout\n",
    "\n",
    "from plotly import graph_objects"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "959741ab-6371-4fa0-b71e-3aa2707a697d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Marker(name=\"HEAD\")\n",
      "BPM(name=\"M\", direction=forward)\n",
      "Quadrupole(name=\"QD\", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QF\", length=1.0, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD\", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Marker(name=\"TAIL\")\n"
     ]
    }
   ],
   "source": [
    "# Load and build simple MADX lattice\n",
    "\n",
    "# Set lattice file path\n",
    "\n",
    "file = Path('../../../config/initial.madx')\n",
    "\n",
    "# Load lattice to dictionary\n",
    "\n",
    "table = load_lattice(file)\n",
    "\n",
    "# Build lattice\n",
    "\n",
    "FODO = build('FODO', 'MADX', table)\n",
    "print(FODO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "7093404a-e240-491a-a614-d0198a2175f3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Marker(name=\"HEAD\")\n",
      "BPM(name=\"M\", direction=forward)\n",
      "Quadrupole(name=\"QD\", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QF\", length=1.0, kn=0.200000000000001, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2, dp=0.0, exact=False, ns=1, order=0)\n",
      "Dipole(name=\"BM\", length=1.0, angle=0.17453292519943395, e1=0.0, e2=0.0, kn=1e-15, ks=0.0, ms=0.0, mo=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Drift(name=\"DR\", length=2, dp=0.0, exact=False, ns=1, order=0)\n",
      "Quadrupole(name=\"QD\", length=0.5, kn=-0.199999999999999, ks=0.0, dp=0.0, exact=False, ns=1, order=0)\n",
      "Marker(name=\"TAIL\")\n"
     ]
    }
   ],
   "source": [
    "# Load and build simple ELEGANT lattice\n",
    "\n",
    "# Set lattice file path\n",
    "\n",
    "file = Path('../../../config/initial.lte')\n",
    "\n",
    "# Load lattice to dictionary\n",
    "\n",
    "table = load_lattice(file)\n",
    "\n",
    "# Build lattice\n",
    "\n",
    "FODO = build('FODO', 'ELEGANT', table)\n",
    "print(FODO)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "521decd2-a036-4167-a317-b69c9d05dd8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build vepp4 lattice\n",
    "\n",
    "file = Path('../../../config/vepp4m.lte')\n",
    "data = load_lattice(file)\n",
    "vepp = build('VEPP4M', 'ELEGANT', data)\n",
    "\n",
    "# Slice dipoles\n",
    "\n",
    "vepp.ns = (('Dipole', 0.1), )\n",
    "\n",
    "# Profile\n",
    "\n",
    "layout = Layout(vepp)\n",
    "\n",
    "# Generate reference orbit\n",
    "\n",
    "x, y, z = layout.orbit(flat=False, step=None, start=(0, 0))\n",
    "\n",
    "# Generate layout (can be saved as html with write_html method)\n",
    "\n",
    "blocks = layout.profile_3d(scale=2.5, exclude=['Drift', 'Marker'])\n",
    "\n",
    "# # Plot\n",
    "\n",
    "# figure = graph_objects.Figure(\n",
    "#     data=[\n",
    "#         graph_objects.Scatter3d(\n",
    "#             x=x.numpy(),\n",
    "#             y=y.numpy(),\n",
    "#             z=z.numpy(), \n",
    "#             mode='lines',\n",
    "#             name='Orbit',\n",
    "#             line=dict(color='black',width=2.0,dash='solid'),\n",
    "#             opacity=0.75,\n",
    "#             showlegend=True\n",
    "#         ),   \n",
    "#         *[graph_objects.Mesh3d(block) for block in blocks]\n",
    "#     ]\n",
    "# )\n",
    "# figure.update_layout(\n",
    "#     scene=dict(\n",
    "#         xaxis=dict(visible=False, range=[-100,100]),\n",
    "#         yaxis=dict(visible=False, range=[-100,100]),\n",
    "#         zaxis=dict(visible=False, range=[-25,25]),\n",
    "#         aspectratio=dict(x=1, y=1, z=1/4),\n",
    "#         annotations=[]\n",
    "#     ),\n",
    "#     margin=dict(l=0, r=0, t=0, b=0),\n",
    "#     legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),\n",
    "#     hoverlabel=dict(font_size=12, font_family=\"Rockwell\", font_color='white'),\n",
    "#     legend_groupclick='toggleitem'\n",
    "# )\n",
    "# figure.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "960fb194-99e2-426d-9488-a4c4c20fe540",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Build skif lattice\n",
    "\n",
    "file = Path('../../../config/skif.lte')\n",
    "data = load_lattice(file)\n",
    "skif = build('SKIF', 'ELEGANT', data)\n",
    "\n",
    "# Slice dipoles\n",
    "\n",
    "skif.ns = (('Dipole', 0.1), )\n",
    "\n",
    "# Profile\n",
    "\n",
    "layout = Layout(skif)\n",
    "\n",
    "# Generate reference orbit\n",
    "\n",
    "x, y, z = layout.orbit(flat=False, step=None, start=(0, 0))\n",
    "\n",
    "# Generate layout (can be saved as html with write_html method)\n",
    "\n",
    "blocks = layout.profile_3d(scale=2.5, exclude=['Drift', 'Marker'])\n",
    "\n",
    "# # Plot\n",
    "\n",
    "# figure = graph_objects.Figure(\n",
    "#     data=[\n",
    "#         graph_objects.Scatter3d(\n",
    "#             x=x.numpy(),\n",
    "#             y=y.numpy(),\n",
    "#             z=z.numpy(), \n",
    "#             mode='lines',\n",
    "#             name='Orbit',\n",
    "#             line=dict(color='black',width=2.0,dash='solid'),\n",
    "#             opacity=0.75,\n",
    "#             showlegend=True\n",
    "#         ),   \n",
    "#         *[graph_objects.Mesh3d(block) for block in blocks]\n",
    "#     ]\n",
    "# )\n",
    "# figure.update_layout(\n",
    "#     scene=dict(\n",
    "#         xaxis=dict(visible=False, range=[-200,200]),\n",
    "#         yaxis=dict(visible=False, range=[-200,200]),\n",
    "#         zaxis=dict(visible=False, range=[-50,50]),\n",
    "#         aspectratio=dict(x=1, y=1, z=1/4),\n",
    "#         annotations=[]\n",
    "#     ),\n",
    "#     margin=dict(l=0, r=0, t=0, b=0),\n",
    "#     legend=dict(orientation='v', x=0., y=1., xanchor='left', yanchor='top'),\n",
    "#     hoverlabel=dict(font_size=12, font_family=\"Rockwell\", font_color='white'),\n",
    "#     legend_groupclick='toggleitem'\n",
    "# )\n",
    "# figure.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [
    "myt0_gMIOq7b",
    "5d97819c"
   ],
   "name": "03_frequency.ipynb",
   "provenance": []
  },
  "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.12.1"
  },
  "latex_envs": {
   "LaTeX_envs_menu_present": true,
   "autoclose": false,
   "autocomplete": true,
   "bibliofile": "biblio.bib",
   "cite_by": "apalike",
   "current_citInitial": 1,
   "eqLabelWithNumbers": true,
   "eqNumInitial": 1,
   "hotkeys": {
    "equation": "Ctrl-E",
    "itemize": "Ctrl-I"
   },
   "labels_anchors": false,
   "latex_user_defs": false,
   "report_style_numbering": false,
   "user_envs_cfg": false
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}