{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "556562f3-8ece-4517-8c93-ee5e2fc29131",
   "metadata": {},
   "source": [
    "# Example-01: Optimizaion (import/export API)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "baa39129-4d73-44c2-bf1e-88996ea48c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading API facilitates (suboptimal) interface with different optimization libraries\n",
    "# In this example, quadrupole gradients are used to fit beta functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7aa2bbf6-cfe5-49c1-b9a7-de7cce585bbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "from model.command.external import load_sdds\n",
    "from model.command.external import load_lattice\n",
    "from model.command.external import text_lattice\n",
    "\n",
    "import numpy\n",
    "from numpy import ndarray as Array\n",
    "\n",
    "from scipy.optimize import minimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "540bbb11-ef4b-4f97-9d9a-179af31e0de8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set quadrupole gradient and compute and return twiss parameters\n",
    "\n",
    "def evaluate(knobs:Array) -> Array:\n",
    "    kf, kd = knobs\n",
    "    path:Path = Path('optimize.lte')\n",
    "    lattice:dict[str, dict[str, str | int | float | dict]] = load_lattice(path)\n",
    "    lattice['QF']['K1'] = float(kf)\n",
    "    lattice['QD']['K1'] = float(kd)\n",
    "    with path.open('w') as stream:\n",
    "        stream.write(text_lattice('LTE', lattice))\n",
    "    !elegant 'optimize.ele' > /dev/null\n",
    "    !sddsconvert -ascii 'binary.twiss' 'optimize.twiss'\n",
    "    path:Path = Path('optimize.twiss')\n",
    "    _, columns = load_sdds(path)\n",
    "    return numpy.asarray([[data['betax'], data['betay']] for location, data in columns.items()]).T"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "f7c9dc5a-7d3e-4b74-a219-82cf714abf3c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set target beta functions\n",
    "\n",
    "target:Array = numpy.asarray([+0.21, -0.19])\n",
    "result:Array = evaluate(target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "ffdb81b2-1cde-4b44-8133-f442d9c3065c",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "0.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Set objetive function to minimize\n",
    "\n",
    "def objective(knobs:Array) -> Array:\n",
    "    return numpy.sum((evaluate(knobs) - result)**2)\n",
    "\n",
    "objective(target)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3dda343b-b675-45c5-9ba3-e2b4f408c0b1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "       message: Optimization terminated successfully.\n",
       "       success: True\n",
       "        status: 0\n",
       "           fun: 1.3489521479279302e-28\n",
       "             x: [ 2.100e-01 -1.900e-01]\n",
       "           nit: 20\n",
       "          nfev: 40\n",
       " final_simplex: (array([[ 2.100e-01, -1.900e-01],\n",
       "                       [ 2.100e-01, -1.900e-01],\n",
       "                       [ 2.100e-01, -1.900e-01]]), array([ 1.349e-28,  2.891e-05,  6.010e-05]))"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Optimize \n",
    "\n",
    "knobs:Array = numpy.asarray([+0.20, -0.20])\n",
    "\n",
    "minimize(objective, knobs, method='Nelder-Mead')"
   ]
  }
 ],
 "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
}