{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8665bce9-d57c-445a-821a-a596efa3f4b3",
   "metadata": {},
   "source": [
    "## Utility Functions"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5a9fc409-a0e1-4c6b-9544-2c9494f40e22",
   "metadata": {},
   "source": [
    "`SIMBA` provides various utility functions for inspecting, modifying and saving lattices and lattice elements. This notebook will demonstrate some of this functionality.\n",
    "\n",
    "For general tips on setting up `SIMBA`, see [getting_started](./getting_started.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ba2a922d-f9fd-4b80-b6a7-7248b20d1601",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import simba.Framework as fw  # noqa E402\n",
    "\n",
    "# Define a new framework instance, in directory 'getting_started'.\n",
    "#       \"clean\" will empty (delete everything!) in the directory if true\n",
    "#       \"verbose\" will print a progressbar if true\n",
    "framework = fw.Framework(\n",
    "        directory=\"./getting_started\",\n",
    "        master_lattice=os.environ[\"LATTICE_LOCATION\"],\n",
    "        generator_defaults=\"clara.yaml\",\n",
    "        clean=False,\n",
    "        verbose=False\n",
    "    )\n",
    "\n",
    "scaling = 4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9dd5ffc7-c315-4cb2-b5c1-14fafe6f655c",
   "metadata": {},
   "source": [
    "Load the settings from a definitions file and check which codes correspond to each line."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1c3e2a3f-bc8e-4ddf-a890-fd30188423df",
   "metadata": {},
   "outputs": [],
   "source": [
    "framework.loadSettings(\"Lattices/clara400_v13_combined.def\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "35af4ed8-0a7a-41b0-931b-9d091c5bd3af",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['generator: ASTRA', 'injector400: astra', 'Linac: elegant', 'FEBE: elegant']\n"
     ]
    }
   ],
   "source": [
    "print([f\"{line}: {framework[line].code}\" for line in framework.lines])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d394e42-af4b-419e-a706-ad1cdc21aef6",
   "metadata": {},
   "source": [
    "Check the global parameters accessible to this instance of SIMBA and all of the `frameworkLattice` objects."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "d58a7024-55fb-4ca6-9ce3-8d34b2d490eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'beam': {'filename': None, 'code': None}, 'GPTLICENSE': '', 'delete_tracking_files': False, 'astra_use_wsl': 1, 'master_subdir': '/home/xkc85723/Documents/simba/docs/source/examples/notebooks/getting_started', 'master_lattice': '/home/xkc85723/Documents/laura-lattices/CLARA/./', 'simcodes_location': '/home/xkc85723/Documents/venv/lib/python3.12/site-packages/SimCodes/'}\n"
     ]
    }
   ],
   "source": [
    "print(framework.global_parameters)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "22d9e304-ef7e-4f29-982e-c3f74b006aae",
   "metadata": {},
   "outputs": [],
   "source": [
    "framework.generator.load_defaults(\"clara_400_2ps_Gaussian\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "31b3d59d-571f-4162-a972-c099dcc5f20c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FrameworkSettings({'settingsFilename': '/home/xkc85723/Documents/laura-lattices/CLARA/./Lattices/clara400_v13_combined.def', 'global': {}, 'generator': {'code': 'astra', 'default': 'clara_400_3ps'}, 'section': '/home/xkc85723/Documents/laura-lattices/CLARA/sections.yaml', 'element_list': '/home/xkc85723/Documents/laura-lattices/CLARA/YAML/', 'groups': {'bunch_compressor': {'type': 'chicane', 'elements': ['CLA-VBC-MAG-DIP-01', 'CLA-VBC-MAG-DIP-02', 'CLA-VBC-MAG-DIP-03', 'CLA-VBC-MAG-DIP-04']}}, 'elements': {}, 'files': {'injector400': {'code': 'astra', 'charge': {'cathode': True, 'space_charge_mode': '2D', 'mirror_charge': True}, 'input': {'particle_definition': 'initial_distribution'}, 'output': {'zstart': 0, 'end_element': 'CLA-S02-SIM-APER-01'}}, 'Linac': {'code': 'elegant', 'output': {'start_element': 'CLA-S02-SIM-APER-01', 'end_element': 'CLA-FEA-SIM-START-01'}}, 'FEBE': {'code': 'elegant', 'charge': {'cathode': False, 'space_charge_mode': '3D'}, 'input': {}, 'output': {'start_element': 'CLA-FEA-SIM-START-01', 'end_element': 'CLA-FED-SIM-DUMP-01-START'}}}, 'layout': '/home/xkc85723/Documents/laura-lattices/CLARA/layouts.yaml'})\n"
     ]
    }
   ],
   "source": [
    "print(framework.settings)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e727cf61-40e4-4046-9929-2770923c1d2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "framework.change_subdirectory('./utility_functions')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2ae0777-71a6-4f5a-a451-f01069af06b2",
   "metadata": {},
   "source": [
    "### Accessing and modifying elements"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "707a8cbe-9cea-4fbd-b013-7f9f9097e35f",
   "metadata": {},
   "source": [
    "Elements and their parameters can be retrieved either using `getElement` or as a key in the `framework` object itself."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "3bac6e0d-b3da-4a78-8269-91eb0a79ba16",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "name='CLA-S02-MAG-QUAD-01' hardware_class='Magnet' hardware_type='Quadrupole' hardware_model='Generic' machine_area='S02' virtual_name='' alias=[] subelement=False CASCADING_RULES={} physical=PhysicalElement(middle=Position(x=0.0, y=0.0, z=3.52715239) datum=Position(x=0.0, y=0.0, z=3.59149478) length=0.12868478212775) electrical=ElectricalElement(minI=-53.0, maxI=53.0, read_tolerance=0.1) manufacturer=ManufacturerElement(manufacturer='Quadrupole Type 1', serial_number='13248', hardware_class='Quadrupole') simulation=MagnetSimulationElement(field_definition=None, wakefield_definition=None, field_reference_position=None, scale_field=False, n_kicks=7, smooth=2, edge_field_integral=0.5, edge1_effects=1, edge2_effects=1, sr_enable=True, integration_order=4, nonlinear=1, smoothing_half_width=1, edge_order=2, csr_bins=100, deltaL=0.0, csr_enable=True, isr_enable=True, field_amplitude=0.0) controls=ControlsInformation(variables={'readback': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:RBV', dtype=<class 'float'>, protocol='CA', units='N/A', description='Readback of CLA-S02-MAG-QUAD-01', read_only=True), 'setpoint': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:SP', dtype=<class 'float'>, protocol='CA', units='N/A', description='Setpoint of CLA-S02-MAG-QUAD-01', read_only=True), 'state': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:STATE', dtype=<class 'int'>, protocol='CA', units='N/A', description='State of CLA-S02-MAG-QUAD-01', read_only=False)}) degauss=DegaussableElement(tolerance=0.5, values=[52.89, -52.89, 31.8, -31.8, 21.2, -21.2, 10.6, -10.6, 5.3, -5.3, 0.0], steps=11) magnetic=Quadrupole_Magnet(skew=False, length=0.12868478212775, multipoles=Multipoles(K1L=Multipole(order=1 normal=5 skew=0.0 radius=0.0)), systematic_multipoles=Multipoles(), random_multipoles=Multipoles(), field_integral_coefficients=FieldIntegral(coefficients=[-2.23133410405682e-10, 4.5196171252132e-08, -3.46208258004659e-06, 0.000111195870210961, 0.0238129337415767, 0.00981229429460256]), linear_saturation_coefficients=LinearSaturationFit(m=0.025500836847330708, I_max=15.575435287034225, f=0.0, a=-1.3755963213967208e-05, I0=942.4766071889139, d=12.215564413215553, L=128.68477053001592), settle_time=45.0, entrance_edge_angle=0.0, exit_edge_angle=0.0, gap=0.032, bore=0.037, plane='horizontal', width=0.2, tilt=0.0, edge_field_integral=0.5, fringe_field_coefficient=0.0, half_gap=0.016)\n",
      "name='CLA-S02-MAG-QUAD-01' hardware_class='Magnet' hardware_type='Quadrupole' hardware_model='Generic' machine_area='S02' virtual_name='' alias=[] subelement=False CASCADING_RULES={} physical=PhysicalElement(middle=Position(x=0.0, y=0.0, z=3.52715239) datum=Position(x=0.0, y=0.0, z=3.59149478) length=0.12868478212775) electrical=ElectricalElement(minI=-53.0, maxI=53.0, read_tolerance=0.1) manufacturer=ManufacturerElement(manufacturer='Quadrupole Type 1', serial_number='13248', hardware_class='Quadrupole') simulation=MagnetSimulationElement(field_definition=None, wakefield_definition=None, field_reference_position=None, scale_field=False, n_kicks=7, smooth=2, edge_field_integral=0.5, edge1_effects=1, edge2_effects=1, sr_enable=True, integration_order=4, nonlinear=1, smoothing_half_width=1, edge_order=2, csr_bins=100, deltaL=0.0, csr_enable=True, isr_enable=True, field_amplitude=0.0) controls=ControlsInformation(variables={'readback': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:RBV', dtype=<class 'float'>, protocol='CA', units='N/A', description='Readback of CLA-S02-MAG-QUAD-01', read_only=True), 'setpoint': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:SP', dtype=<class 'float'>, protocol='CA', units='N/A', description='Setpoint of CLA-S02-MAG-QUAD-01', read_only=True), 'state': ControlVariable(identifier='CLA-S02-MAG-QUAD-01:STATE', dtype=<class 'int'>, protocol='CA', units='N/A', description='State of CLA-S02-MAG-QUAD-01', read_only=False)}) degauss=DegaussableElement(tolerance=0.5, values=[52.89, -52.89, 31.8, -31.8, 21.2, -21.2, 10.6, -10.6, 5.3, -5.3, 0.0], steps=11) magnetic=Quadrupole_Magnet(skew=False, length=0.12868478212775, multipoles=Multipoles(K1L=Multipole(order=1 normal=5 skew=0.0 radius=0.0)), systematic_multipoles=Multipoles(), random_multipoles=Multipoles(), field_integral_coefficients=FieldIntegral(coefficients=[-2.23133410405682e-10, 4.5196171252132e-08, -3.46208258004659e-06, 0.000111195870210961, 0.0238129337415767, 0.00981229429460256]), linear_saturation_coefficients=LinearSaturationFit(m=0.025500836847330708, I_max=15.575435287034225, f=0.0, a=-1.3755963213967208e-05, I0=942.4766071889139, d=12.215564413215553, L=128.68477053001592), settle_time=45.0, entrance_edge_angle=0.0, exit_edge_angle=0.0, gap=0.032, bore=0.037, plane='horizontal', width=0.2, tilt=0.0, edge_field_integral=0.5, fringe_field_coefficient=0.0, half_gap=0.016)\n"
     ]
    }
   ],
   "source": [
    "print(framework.getElement(\"CLA-S02-MAG-QUAD-01\"))\n",
    "print(framework[\"CLA-S02-MAG-QUAD-01\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "fc044a1a-0ecc-4abb-9fcb-2e596926213d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "5\n",
      "5\n"
     ]
    }
   ],
   "source": [
    "print(framework.getElement(\"CLA-S02-MAG-QUAD-01\", \"k1l\"))\n",
    "print(framework[\"CLA-S02-MAG-QUAD-01\"].k1l)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8074391-a6fc-4a9d-87ae-215bb8dceac8",
   "metadata": {},
   "source": [
    "All objects of a given type defined in the entire lattice can be accessed as follows.\n",
    "\n",
    "(NB here cavity on-crest phase is zero and ASTRA convention is followed for the phase sign)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "c7554ba3-8c96-46d9-9ddf-3abab3894d3c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'CLA-L02-LIN-CAV-01': 23.0, 'CLA-L4H-LIN-CAV-01': -184.0, 'CLA-L01-LIN-CAV-01': 16.0, 'CLA-L03-LIN-CAV-01': 8.0, 'CLA-L04-LIN-CAV-01': 45.0, 'CLA-HRG1-GUN-CAV-01': 9.0}\n"
     ]
    }
   ],
   "source": [
    "cavnames = [c[\"name\"] for c in framework.getElementType(\"rfcavity\")]\n",
    "print({k: v for k, v in zip(cavnames, framework.getElementType(\"rfcavity\", \"phase\"))})"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b33487a-a78b-4e6c-aac2-aeaf0316dff4",
   "metadata": {},
   "source": [
    "Various options are available for changing a lattice element, or multiple elements of a given type simultaneously."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "02bb21f7-b70e-42db-ade7-fcb91a9ea7bb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'CLA-FMS-MAG-DIP-01': 0, 'CLA-VBC-MAG-DIP-01': 0, 'CLA-FMS-MAG-PMD-03': 0, 'CLA-FMS-MAG-DIP-05': 0, 'CLA-FMS-MAG-PMD-02': 0, 'CLA-FMS-MAG-DIP-06': 0, 'CLA-FMS-MAG-DIP-07': 0, 'CLA-FEA-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-03': 0, 'CLA-FEA-MAG-DIP-02': 0, 'CLA-VBC-MAG-DIP-03': 0, 'CLA-FED-MAG-DIP-01': 0, 'CLA-FMS-MAG-PMD-01': 0, 'CLA-VBC-MAG-DIP-02': 0, 'CLA-SP2-MAG-DIP-01': 0, 'CLA-VBC-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-08': 0, 'CLA-FMS-MAG-DIP-04': 0, 'CLA-FMS-MAG-DIP-02': 0, 'CLA-SP3-MAG-DIP-01': 0, 'CLA-FEA-MAG-DIP-01': 0, 'CLA-FEA-MAG-DIP-03': 0, 'CLA-C2V-MAG-DIP-01': 0}\n",
      "{'CLA-FMS-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-01': 0.1, 'CLA-FMS-MAG-PMD-03': 0.1, 'CLA-FMS-MAG-DIP-05': 0.1, 'CLA-FMS-MAG-PMD-02': 0.1, 'CLA-FMS-MAG-DIP-06': 0.1, 'CLA-FMS-MAG-DIP-07': 0.1, 'CLA-FEA-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-03': 0.1, 'CLA-FEA-MAG-DIP-02': 0.1, 'CLA-VBC-MAG-DIP-03': 0.1, 'CLA-FED-MAG-DIP-01': 0.1, 'CLA-FMS-MAG-PMD-01': 0.1, 'CLA-VBC-MAG-DIP-02': 0.1, 'CLA-SP2-MAG-DIP-01': 0.1, 'CLA-VBC-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-08': 0.1, 'CLA-FMS-MAG-DIP-04': 0.1, 'CLA-FMS-MAG-DIP-02': 0.1, 'CLA-SP3-MAG-DIP-01': 0.1, 'CLA-FEA-MAG-DIP-01': 0.1, 'CLA-FEA-MAG-DIP-03': 0.1, 'CLA-C2V-MAG-DIP-01': 0.1}\n"
     ]
    }
   ],
   "source": [
    "# Set all dipole angles in the lattice to zero and use `setElementType`\n",
    "values = [0 for _ in framework.getElementType(\"dipole\")]\n",
    "dipnames = [d[\"name\"] for d in framework.getElementType(\"dipole\")]\n",
    "framework.setElementType(\"dipole\", \"angle\", values)\n",
    "print({k: v for k, v in zip(dipnames, framework.getElementType(\"dipole\", \"angle\"))})\n",
    "\n",
    "# Set all dipole angles to 0.1 and use modifyElementType\n",
    "value = 0.1\n",
    "framework.modifyElementType(\"dipole\", \"angle\", 0.1)\n",
    "print({k: v for k, v in zip(dipnames, framework.getElementType(\"dipole\", \"angle\"))})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "ee6d7c47-d443-45f4-b93a-0d191563c23c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CLA-S02-MAG-QUAD-01 k1l = 1\n",
      "CLA-S02-MAG-QUAD-01 k1l = 2\n",
      "CLA-S02-MAG-QUAD-01 k1l = 3\n",
      "CLA-S02-MAG-QUAD-01 k1l = 3, n_kicks = 5\n",
      "CLA-S02-MAG-QUAD-01 k1l = 4, n_kicks = 6\n",
      "CLA-S02-MAG-QUAD-01 k1l = 5, n_kicks = 7,CLA-S02-MAG-QUAD-02 k1l = 5, n_kicks = 7\n"
     ]
    }
   ],
   "source": [
    "s02q1 = \"CLA-S02-MAG-QUAD-01\"\n",
    "s02q2 = \"CLA-S02-MAG-QUAD-02\"\n",
    "setattr(framework.elementObjects[s02q1], 'k1l', 1)\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n",
    "framework[s02q1].k1l = 2\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n",
    "framework.modifyElement(s02q1, \"k1l\", 3)\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}\")\n",
    "framework.modifyElement(s02q1, [\"k1l\", \"n_kicks\"], [3, 5])\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}\")\n",
    "framework.modifyElement(s02q1, {\"k1l\": 4, \"n_kicks\": 6})\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks}\")\n",
    "framework.modifyElements([s02q1, s02q2], {\"k1l\": 5, \"n_kicks\": 7})\n",
    "print(f\"{s02q1} k1l = {framework.getElement(s02q1, 'k1l')}, n_kicks = {framework[s02q1].n_kicks},\\\n",
    "{s02q2} k1l = {framework.getElement(s02q2, 'k1l')}, n_kicks = {framework[s02q2].n_kicks}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 33,
   "id": "ab02d5cd-d699-4da5-a8c7-a70a1017e38a",
   "metadata": {},
   "outputs": [],
   "source": [
    "import shutil\n",
    "shutil.rmtree(\"./utility_functions/\")\n",
    "os.remove('dipoles.yaml')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31c9844a-e615-4194-95e1-04479e0e9954",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}