tracking-parametrisation-tuner/parameterisations/notebooks/magnet_kink_position.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "incomplete input (4052887367.py, line 38)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  Cell \u001b[0;32mIn[16], line 38\u001b[0;36m\u001b[0m\n\u001b[0;31m    #   & (array[\"match_chi2\"] < 5)]\u001b[0m\n\u001b[0m                                    ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m incomplete input\n"
     ]
    }
   ],
   "source": [
    "import uproot\n",
    "import awkward as ak\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import mplhep\n",
    "\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "input_tree = uproot.open({\n",
    "    \"/work/cetin/LHCb/reco_tuner/data/tracking_losses_ntuple_B_BJpsi_def_selected.root\":\n",
    "    \"Selected\"\n",
    "})\n",
    "array = input_tree.arrays()\n",
    "\n",
    "array[\"dSlope_yEndT\"] = array[\"ideal_state_9410_ty\"] - array[\n",
    "    \"ideal_state_770_ty\"]\n",
    "array[\"dSlope_yEndT_abs\"] = abs(array[\"dSlope_yEndT\"])\n",
    "\n",
    "array[\"dSlope_xEndT\"] = array[\"ideal_state_9410_tx\"] - array[\n",
    "    \"ideal_state_770_tx\"]\n",
    "array[\"dSlope_xEndT_abs\"] = abs(array[\"dSlope_xEndT\"])\n",
    "array[\"x_EndT_abs\"] = abs(array[\"ideal_state_9410_x\"])\n",
    "array[\"x_EndVelo_abs\"] = abs(array[\"ideal_state_770_x\"])\n",
    "\n",
    "array[\"y_EndT_abs\"] = abs(array[\"ideal_state_9410_y\"])\n",
    "array[\"y_EndVelo_abs\"] = abs(array[\"ideal_state_770_y\"])\n",
    "\n",
    "array[\"z_mag_xEndT\"] = (\n",
    "    array[\"ideal_state_770_x\"] - array[\"ideal_state_9410_x\"] -\n",
    "    array[\"ideal_state_770_tx\"] * array[\"ideal_state_770_z\"] +\n",
    "    array[\"ideal_state_9410_tx\"] *\n",
    "    array[\"ideal_state_9410_z\"]) / array[\"dSlope_xEndT\"]\n",
    "\n",
    "array[\"z_mag_corr\"] = array[\"z_mag_xEndT\"] - array[\"match_zmag\"]\n",
    "\n",
    "sel_array = array[(array[\"z_mag_xEndT\"] < 5800)\n",
    "                  & (array[\"z_mag_xEndT\"] > 5000)]\n",
    "                #   & (array[\"match_chi2\"] < 5)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC2MklEQVR4nOz9T2wbd57n/79KyWGBxUgleTCAIAEdFeM9/HJYi5L7vhaZ7uvEpJxzOiLtOQWatmjlvrYpd/TLaduknJzHIu1ce0zKua9FOn3IKWFJDUggsLuRymwssJewvgeBHFF/KIr/qig+HwAxovTh+/MuUZ4ZvvL5fMpwXdcVAAAAAAAAcMWNeN0AAAAAAAAA0A8EYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGAoEYQAAAAAAABgKBGEAAAAAAAAYCgRhAAAAAAAAGArve90A0G3/+T//Z/2///f/9N577+mf/umfvG4HAAAAAAAc87/+1//Sr7/+qv/0n/6T/u///b99ndtwXdft64xAj7333nuqVqtetwEAAAAAAJoYGRnRr7/+2tc5WRGGK6cWhI2MjGhyctLrdhqUy2Xf9tZNw3Cdw3CN0nBc5zBcozQc1zkM1yhxnVfJMFyjNBzXOQzXKHGdV8kwXKM0HNc5qNdY6/u9997r+9wDF4Tt7u6qWCzq4OBAjuPol19+kSRdu3ZNpmnKsizNz89rdHS0r32tra0pl8tpe3tbjuPIsiwFg0HF43GFQqGBm9vL6+nUP/3TP2l/f1+Tk5Pa29vzup0G09PTvu2tm4bhOofhGqXhuM5huEZpOK5zGK5R4jqvkmG4Rmk4rnMYrlHiOq+SYbhGaTiuc1Cvsda3F8cZDUQQ9vLlS+VyOW1ubspxnJZeY1mWQqGQotGobt261bPe8vm8otGoHMdRKBRSJpORZVkqFotKJBIKh8P175um6fu5u1lzfHy85ffrpFgsplQq1dZrAQAAAAAAzuLrIOzZs2dKJpOybVuSdPw4M8MwznxNbUypVJJt20qn07IsSw8ePNAf/vCHrvaXz+cVDoclnQ5uLMtSJBLR3Nyc8vm85ubmVCgUuhaG9WLubtbMZrNth2CS6n0AAAAAAAB0y4jXDZzl9evXun79uuLx+JkhWO35WY+zxti2rVgspv/yX/6L/vrXv3alR8dxFI1GJR2FROetXspkMpIk27br4/04d7drdrqaKxKJdPR6AAAAAACAk3y3Imx1dVVra2unQq1QKKRgMKhr167Jsqz6SqSJiQlJ0sHBgaSjQMe2bf3yyy+ybVv5fL6+Munnn39WMBhUIpHQw4cPO+qztn1QkhKJxLnjaiupstms8vm80um0YrGY7+buZs3a792yLCUSCYVCofr71Mzc3Jxs2/b9GWQAAAAAAGAw+SYIq1QqWlhYUKFQkHQUfNXOo5qdne2o9rt37/T8+XOl02kVi0Ulk0nl83nl8/m2DtWvBT01i4uLTcffuXNH2WxW0lHI1EkQ1ou5u10zlUrJsqxLbQW1bbu++q9bK+cAAAAAAACO88XWyB9++EEzMzMqFAqKxWIqlUp69eqV7t+/33EIJkljY2OKxWLa3t7Wzz//rFu3bml7e1vz8/P6+9//ful6yWSy/nUoFLow7Dm+zc9xnHqI1I5ezN3tmul0+tI3Bzhe46IgDgAAAAAAoB2eB2Fv375VMBjUzMyMSqWSnj59qpmZmZ7NZ1mWcrmcXr16pf/9v/+35ubmVKlULlUjnU7Xvw4Ggy3PW/P8+fNLzdfrubtdM5PJtFznZI1WgjgAAAAAAIB2eLo1cmdnRwsLC4pEItrc3Ozr3KFQSDs7O5qbm9Pi4qL+8pe/tPS6YrHY8PzmzZstvS4YDNa3/rW7IqwXc/ei5mXP+HIcp94H2yKvhuXlZVUqlba2Hg+KYbhGaXiucxgMw3s5DNcoDc91DoNheS+H4TqH4RqHyTC8n8NwjdLwXCcux9MgLBwOa3FxUU+fPvVkftM0VSgUND8/r6+++kr/+q//euFrjp+lJTWujGrm5LhisXjpVVO9mNvL66k5viKNbZFXw/Lystct9NwwXKM0PNc5DIbhvRyGa5SG5zqHwbC8l8NwncNwjcNkGN7PYbhGaXiuE5fj2dbIBw8eyLIsz0KwGtM0tb29rf/+3/+7fvjhhwvHv3nz5tTrW3Ht2rWG59vb26222NO5vbyemtq2yGAwyLZIAAAAAADQM54FYdeuXVMmk/Fq+gamaWpra6u+1a+Zk2PaXUFVKpVab7CHc3t5PVLjtsg7d+60VQMAAAAAAKAVnm2NvH//vldTn2l2dralO1S2Epa1wnGcS7+mF3N7eT2SGs6GO343SgAAAAAAgG7z9Iywfrp3757+/Oc/d1yn3cDn5Ja/g4MDX8zt5fVIUiqVknS0LbLV1WitKpfLmp6e7rjO8vIye8sBAAAAAENvfX1d6+vrHdcpl8td6KY9QxOEpdNpJRIJffDBB163Iqn9AMqvc7dTs9fbIqvVqvb39zuuU6lUutANAAAAAACDrVKpdOVztpeGIgjb2tqS67p68eJFS3eGbMY0za4ESe0cCt+Lub28nl5vixwZGdHk5GTHdbjVLgAAAAAAR5+Pp6amOq5TLpdVrVa70NHleRKE3b17Vzs7Oz2f5+DgQI7j1M/Bevr0acdB2MTERFeCo4mJCV/M7eX11LZFWpbV9W2RkjQ5Oam9vb2u1wUAAAAAYBh16+ig6elpz1aWeRKEmaapXC4nwzB6Oo/rug3PbdvWDz/8oBs3brRds52VT9LprYPtrgjr9txeXc/xbZHDdEj+8vKyKpUKq8wAH+HfJeBP/NsE/Id/l4D/8O/y8jwJwu7cuaO1tbWWx7uu21Zodvw1tVBsc3OzoyBsfn6+HuBIR4FOKyHQycPkA4GAL+b26nqOb4vsxflgfsWh+4D/8O8S8Cf+bQL+w79LwH/4d3l5I15MOjs7W98K57ruhQ/DMFoad96jxnVdZTKZjnqfm5treF7bdnmRUqnU8DwUCvlibq+up/Y+mKapYDB4qdcCAAAAAAC0w5MgTJIWFhYkSfl8XtVq9dxHMpmU67oKBoPK5XI6PDxsOv7kY3t7WzMzM4rH46pWq/rpp5866nt+fr7heavB0fGthKZptnUmVi/m9uJ6HMdRPp+XJMVisZZfBwAAAAAA0AnPgrDFxUUFg0HdunXr3DE7Ozt68OCB4vG4tre3tbCwoLGxsUvNEwwG9erVK6VSKf3Lv/xLp20rGAw2bB188+ZNS6/b3t6uf30yfPJybi+uZ1i3RQIAAAAAAG95uiJsfHy86ZhkMinLsvTnP/+5o7ksy9LS0pJSqZS+//77jmpJRyFezfHztZo5Pi6RSPhq7n5fD9siAQAAAACAFzwLwiTp1atXTX++vb2teDzelbnm5+fluq5SqVTHtY73VNvi18zxMZZltXU+WC/n7uf1HN8WeTyAAwAAAAAA6DVPg7CLFIvFrq0YmpiYkCTlcrmOawWDwYbwJ5vNNh1//ID+i1ZPpdNpJRKJc8/q6sXcvbyek45vi4xGo5d6LQAAAAAAQCcM9/htFX1mYmJCT5480R/+8IeOa929e1fpdFqGYejXX3/tuJ5t2woEApKOgqRCoXDmOMdx6ltAQ6FQ0yAuHA43rLY6PDxsOL+rl3P3ouZZatdomqYODw8v9dpWTU9Pa39/X1NTU9rb2+vJHAAAAAAAoD1efm739Yqw+fl5PX36tOM67969UzqdlqQzg6V2WJZVXxlVLBa1trZ25rja3TFN02xYSXWWk9s
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(\n",
    "    sel_array[\"z_mag_xEndT\"],\n",
    "    bins=100,\n",
    "    # range=[5100, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    density=True,\n",
    ")\n",
    "plt.xlabel(r\"z$_{Mag}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACjlElEQVR4nOz9T4xj953fe39Y3RZajxU2WRoEIIrATJOjQLPKNNk9C99VpsjHGwl2RmR3YD2SgIyLlBXfLJiYdFm4GM9CLrPsIXCRG6PJ1gR42leLbtIeCdLCUbGd3QWSrmJ7dkJsHmmAKhAIoqrTHAluyBJ5FxXSZP0hD8nDc/jn/QII159v/c73FKus5qd+fzztdrstAAAAAAAAYMGtuN0AAAAAAAAA4ASCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALAWCMAAAAAAAACwFgjAAAAAAAAAsBYIwAAAAAAAALIWLbjcAjOvLX/6yHj9+rAsXLuif/tN/6nY7AAAAAADAgv/xP/6HvvjiC126dEmffvqpo9f2tNvttqNXBGxy4cIFtVott9sAAAAAAABjWFlZ0RdffOHoNZkRhrnVCcJWVlYUCAS6H280Gmd+fBqcuhbXmf1rLdp1nLzWol3HyWtxndm/1qJdx8lrcZ3Zv9aiXcfJay3adZy8FteZ/Wst2nWcvNYyXafzuQsXLkzt+udqA3NqbW2tLam9trZm6eNO9sB1ZuM6Tl5r0a7j5LUW7TpOXovrzP61Fu06Tl6L68z+tRbtOk5ea9Gu4+S1uM7sX2vRruPktZbpOk4+fye5OiPsmWeecfPyfTwej/77f//vbrcBAAAAAACAKXE1CKvX6/J4PGq7uE1Z5/oej8e1HgAAAAAAADB9K2434DY3QzgAAAAAAAA4x/XN8tvtthKJhEKhkCvXN01TpVLJlWsDAAAAAADAOa4HYaVSSd/85jdd7SGRSOirX/2qqz0AAAAAAABgulxfGhmLxdxuQdevX3e7BQAAAAAAAEyZ60HY6uqq2y0AAAAAAABgCbi6NJKTGjENmUxGzWZTXq93Ya61aNdx0qJ97xbx59spi/i9W7TrOGnRvneL+PPtFJ6j2b+Okxbte7eIP3dO4Tma/es4adG+d4t2nVF52i4em7iysiLTNF3/pjx69Eh+v1+tVsvVPjCaYDCog4MDra2taX9/3+12MAKeO/Ti5wEd/CygFz8P6OBnAb34eUAHPwvzzc3nz9Wlkfl83vUQTJIuX76sfD7vdhsAAAAAAACYIleDsO985zuOXu+jjz4693NO9wIAAAAAAABnub5ZvpPS6bTbLQAAAAAAAMAlSxWE7e7uut0CAAAAAAAAXOLqqZHj+uijj2SapuV6wzBULBZH+hoAAAAAAAAslrkIwj766CPl83lVq1UZhjHWGO12Wx6Px+bOAAAAAAAAMC9mPgjb3NzU9va2pOMwCwAAAAAAABjHTAdhP/vZz5TP5yVJHo9HHo+HMAynNBoNBYPBgTWZTEaZTMahjjBMJpNRs9mU1+t1uxXMAH4e0MHPAnrx84AOfhbQi58HdPCzMLsKhYIKhcLAmkaj4VA3p3naM5wsXbt2TbVarRuAhUIhRSIRhUIhSdLTTz89dIyPP/5Ypmnq3r17evTokb744otptw2HBINBHRwcWKr9q7/6K33/+9+fbkMAAAAAACy573//+/rrv/5rS7Vra2va39+fckf9ZnpGWCcEk6SdnR2tr6+PPVY2m9UzzzxjV2uYISsrKwoEAgNr+CsBAAAAAADT5/V6tba2NrCm0Wio1Wo51FG/mZ4Rtrq6qkePHimbzWpra2vi8f74j/9Yv/nNb2zoDLOgMyPMjQQZAAAAAACMx83X8yuOXm1EnSWQ169ft2W8YrFoyzgAAAAAAACYPzMdhHWWQh4eHto6HgAAAAAAAJbPTAdh3/ve99Rut1Wr1WwZ780337RlHAAAAAAAAMyfmQ7CLl++rB/+8Ie6e/eu/vEf/3Hi8WZpaaRhGIrH45Zr0+m0wuGwPB6P/H6/otGo0um0DMOYcqf2G+XeAQAAAAAA7DLTQZh0fNpjJBJRMpmcaJwPP/zQtpllg5imKY/HM/QRDoe7e6ANsr29rXA4rFKp1A29TNNUrVZTqVRSOBzW9va2pd78fr+l3s56pNNpx+8dAAAAAADAThfdbsCKnZ0dXbt2Tc8884xyuZxWV1ctfd3h4aFM01S9Xte9e/em3OWxUqlkuTaXyw38fDweV7Valc/nUywWUygUkmEYqtVqfTPBcrmcQqGQEonEuWNVKhWZpmm5t7N6GcbOewcAAAAAALDbXARhv/zlLyVJ9Xrd0syks7TbbXk8HjvbOtPW1paluk6wdZ5cLqdqtap8Pq9sNnvq89vb231hUjKZVLvdPne8SZeFDgrZOuy6dwAAAAAAgGmY+SDsW9/6VnemkcfjGRj2nMeJAEw6nhFlmqay2ezQGVTXrl0793OGYWh7e1s7OzuKxWJn1mSzWdXr9b5ZWLVaTZFI5MzxqtWqQqGQcrmcYrGYpVl10WhUhmGc20Mvu+4dAABgkbTbUmvIP19XPJJD/1wFAGDpedrjJEsOuX37dncG2LghWC+Px6MvvvjCjtbOFA6HJR3PXJtEMpnU9evXz5wJ1ss0Tfn9/u77580ey+VyqlQq2tvbk8/ns9SDYRjd+ykWi0qlUgPr7br3UQSDQR0cHGhtbU37+/uOXRcAADcRrMyXL1rSxdcH13z+hnRh5nfuBQDAPm6+np/pGWGd5XydECwWiykej8vn8428T9itW7f00UcfTa3XSqUiwzBsOZmyM7NqGJ/P1903rPP+WUqlku7fv285BJOO76fjxo0bQ2vtuncAADBYq20xWCEIAwAAOGWmg7BarSaPxyOfz6fd3V1duXJl7LE2Njb09NNP29hdv62tLfl8vqGhkRU7OzuWaw8PD7tvn7fksFwun7lkcpC7d+9KOt7Pa1iAZue9AwAAYDzMFgQAYLiZDsJ8Pp8ePXqkzc3NiUKwzlhXr161qbN+tVpNtVpNkuT3+xUKhbqz16xsMj8u0zS7J0HGYrFzwy4re3ydHLdzP8lkcmCtW/cOAACwKN56KL0cnXwcZgsCADDcTO9GEI0e/4vArhMGb9++bcs4J/We3igd769VKpWUTCbl8XiUTCa7YZGd7t27J+n4+1Mul20bt3cD/mGzvNy6dwAAcL63HrrdATqsPBevVHjOAABwykwHYel0Wu12u7sH1qSmMSOscyLjIJVKRdFotLvxvx1M01Q6nVYkEtHOzs5I+38N01kWGYlEBo7r1r0DbuvMxAQANxCszI9Hj6XX3rFW++rbUvPxVNsBAACa8SAskUjo6tWr3WBmUj//+c9tGadXKBRSsVhUPp9XKpUaOHutVCp1Z7lNwjAMRaNR+Xw+3b9/37YZc1L/ssibN28OrHXj3gG31Go15XI5+f1+bWxsuN0OgCVFsDJfflqTPv2dtdpPPpPuMIkeAICpm+kgTDre6H1vb09/93d/N/FYW1tbNnR0WiqVUjabVbFYVL1e19HRkfL5/JmzqWq1muLx+NjXqlQqCofDMgxDpmnK7/dre3t7gu77dZZbSrK0x5eT9w6cpxNSxeNxhcNh+f1+eTwe+f1+RaNRxeNx5XK5c2cwJpPJU0GtaZqqVCpKJpPdcba3t5kNBsBVBCvz5d0PRqt/b8R6AAAwOk+73R5ytoz7SqWSNjc39fHHH489xqNHj7S6uqovvvjCxs6Gq1Qq2tjYOPXiOZ/PK5vNWhrDNE2VSiUVi8V
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.35, 0.35, 35)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(5100, 5700)\n",
    "plt.xlabel(\"$t_x$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_tx_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACnRklEQVR4nOz9T2wj+X3n/7+omQQOkqiLmmABQQLiLmYC7ynponoPufwQi1xfYuTPkN2LGLYvFukxgj1oM+S0c4hzsGXSsS4LOEPKuTg/A6sm7U2QHLLD6tn7tliT3Ixds3oMSCCwwIjVcoIYiUf1PWhJkxL/FMniH4nPB0BAEj/81JuSekZ88fN5fyK+7/sCAAAAAAAA7ri1RRcAAAAAAAAAzANBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQRgAAAAAAABWAkEYAAAAAAAAVgJBGAAAAAAAAFYCQdgCuK6rZDI59TyxWEyRSETVanWsxzmOo2w223l8JBJRLBZTPp+X53kT11MsFpVMJhWNRjtzptNp2bY98ZwAAAAAAABhIQgLked5nWBp2C0Wi8k0zamulc/n5bru2PWl02nF43GVy+Wex7uuq2KxqGg0qnK5PNa8tm0rGo0qn89LkiqVihqNhgqFghzHUTKZVDKZnCpkAwAAAAAAmNariy7gLhknQGqHRpOwbVvFYnGsx3iep3g8Hig8y2azqtfrKpVKgWppr27LZDI9jzFNU6lUSvF4XLZtKx6Pq16vyzCMsWrv5xd/8Rf1k5/8RK+88or+3b/7d1PPBwAAAAAA5uP//t//q48++kgf+9jH9M///M9zvXbE931/rle8w6LRaKBVT4lEQrVabaJreJ6n+/fv91ynUqkolUoNfVwymZRt27IsS0+ePJFlWZKutkk+f/68b7A2at7uWkzTVKPR6DvOdV3FYjFJ0z33bq+88oouLy+nngcAAAAAACzG2tqaPvroo7lekxVhISmXy/I8T7lcbmT/r52dnYmvs7e3p42NDUkKvNWwXC7Ltm3lcjkVCoWe+9qrtrLZrNLptBzH6bnWsCAsnU53ahi2wq19jWq1Ktu2VS6XlclkAtU+SDsIW1tb0+bm5lRzBdFsNud6vSCo6fbVI1FTUNQUzLLVtGz1SNQUFDUFs2w1LVs9EjUFRU23rx6JmoKipmCWraZ519O+3iuvvDLza93gIxSmafqmac70GqVSyZfk1+t13zAMX5Ivya9UKiNrSyQSI+dvNBqdOdu3Wq0WaGyr1Ro6d6VS6Yw1DGNkLaNsbW35kvytra2p51rG6wVBTaMtWz2+T01BUVMwy1bTstXj+9QUFDUFs2w1LVs9vk9NQVHTaMtWj+9TU1DUFMyy1bRKr7Fplh+CarUq13Wn6vs1iuu6ymazyuVynW2NQTiOI9d1ValURo41TfPGirHuFWLdusclEomRfb+6V5Z5njf2SZcAAAAAAADTWujWyNdff32Rl+8RiUT0v//3/57osQcHBzIMQ48ePQq5qp9Jp9OyLOtGUDXK8fGxMplM4Ab1iUSi5/MPP/yw77jugwGCBnOmaXaa9R8fH4/sawYAAAAAABCmhQZhjUZDkUhE/gL79bevH4lEJnq84zidVVPRaFSmaSqRSCiZTIYW9OTzeTmOM7AZ/TCPHz+WaZqBx18PtdpN7rtdXyX28OHDwHO3gzBWhAEAAAAAgHlb+a2R04Zw17dDuq6rcrmsdDqtSCRyowH9uBzHUbFYVKlUGivQarMsK/BqMOlmA/5+17Rte+SYfq6Pm+b7AgAAAAAAMK6Fnxrp+75SqdREIU8YPM/r2eY3Dtd1b4RC11WrVVWrVWUyGZVKpbGvsbu7q1QqNfUpi0G1V2y1Xd8qKUnPnz/v+Txo0Pbaa6/1fH5ycjJWvzMAAAAAAIBpLDwIK5fL+sIXvrDQGlKplD71qU+N/TjTNFUqleR5nhqNhmzbvhEktZXLZZ2cnKherweeP51OS5KOjo7Grm1SJycnnY8HhW/Xn+OkK8Im2eoJAAAAAAAwqYUHYf1WHM1b0B5X/VwPi9orzA4ODm5sM3QcR8lkUrVabeS87ZVktVptrK2N0+petTboFMxBYd+4rn9/AAAAAAAAZmnhQdjGxsaiSwiVYRjK5XLK5XKqVqva29vrCXxs21axWFQulxs4h+d5SqfTymQycw0KXdft9O0qFAoDV3pNGmBdD/TOz88nmqdbs9nU9vb21PPs7+9rf39/6P0XFxdaX1+f+lphoabbaRm/R9QUzDLWtGyW8XtETcEsY03LaNm+T8tWj0RNt9myfZ+WrR6Jmm6zZfw+LVtNQes5PDzU4eHh1NdrNptTzzGpiL/AIxtfeeUVtVqthf/gX758qY2NDX300Uehz+15nnZ3d3sawxuGoVarNfAx8Xi8s91ykGg02gmkKpVKKCdUZrNZlctlmaY59NrXT9gM+itk27aSyWTn80QiEWh1XD/b29s6Ozub6LH9/Omf/qm+8pWvhDYflkP792Rra0unp6eLLgcIHb/juMv4/cZdxu837jJ+v++2r3zlK/qzP/uz0OZbxO/JQleELTCDu2FWtRiGoXq9rng83gnDPM+Tbdt9V3sVi0U5jjNWL7EwOI6jcrkswzBGhlOGYYSyrTGMLZ9ra2va3Nycep5Fh7EAAAAAACy79fV1bW1tTT1Ps9nU5eVlCBWNb6FBWKFQWIoA4t69eyoUCjO9xtHRkeLxeOfzWq12IwhzHEf5fF6FQmHupynu7e1Jkp49ezay+f3GxkYoQVgY22I3Nzd5lwEAAAAAgDkY1VYoqLB3eY1jbSFX/X/eeuutuV7vgw8+GHjfrGuxLKsn+OrXcD6dTsuyrKH9w2Yhm83KcRxVKpVAAdykK7muh2fzPAQAAAAAAABg4c3y5ymbzep//I//sbDrJ5NJ2bbd975isSjXdZVIJJROp0fO1R0qHRwc6Pj4uPP548ePA/cMK5fLKpfLKpVKgR+zs7PT0/PM87xAodb15vixWCzQ9QAAAAAAAMKwUkHYycnJQq/fveXw+rbADz/8UJIGBmXDOI7TE0yZphko1LJtW9lsVqVSSZlMJvD1urd4Sler24KsJLvegH+eJ2ICAAAAAADcyiDsgw8+GKtHleu6KpVKofS1mkZ3ELbobYGO4yiZTKpQKIwVgklXK8K6BQ3Cur//hmGM7EUGAAAAAAAQplsRhH3wwQcqFAqybbtvb60gfN9XJBIJubLxdK9ISyaTPfcVCoWxGvbHYrHO96JSqQTe1ihdBVe7u7vK5XIT9SOzLKvn5Mjnz58Hun73878epgEAAAAAAMza0gdhT548UbFYlHQVZt1m3VsDF7Ut0HVdxeNxZTKZwMGb67qqVqs9odmjR49ULpclqWdb5jDd4/L5/BhVA5PZ39/XxcXFUpxOC8wCv+O4y/j9xl3G7zfuMn6/sewi/hKnS9/73vc6jePbq7mmKTcSieijjz4KpbZJtFdx5XK5sVZ/DZtLCr4izPM8xeNxJRIJlUqlwNeKx+M6Ojrq2f7oOE5Pr7BRPxfbtjur4EzTvNEvbFzto1a3trZ0eno61VwAAAAAAGB+FvmafqlXhB0cHEi6CrB835dpmrIsq9Nb6rXXXhs5x4cffijP8/T06VO9fPlypvUOU61W5bquDMPQkydP5n79dghmmqby+XygLaau63ZWbl3vAWZZlhKJRKe5f7VaHRrGVSqVzsesBgMAAAAAAIuw1EGY4zidlWC1Wk27u7sTz5XL5fT666+HVZps21Y6nZbneUo
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dx/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACpLklEQVR4nOz9X4xad57n/7/A7shRshgqo5FQIc0YlFX6atpg70XfzRi2b2IlswF7lSiJtNsFmXx79oLdhrhz7a6mpoebmW/kwpmVfo4irQ2dSZRc9KZwX470XVfhnrtouzlOS1VCWm1XHTOJYiUO/C5qoKH+HA7FgcOf50Mqqf6863Pep/7ZvPj88bTb7bYAAAAAAACAOed1uwEAAAAAAABgEgjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEAjCAAAAAAAAsBAIwgAAAAAAALAQCMIAAAAAAACwEE673QBwUk899ZQePXqkU6dO6Y//+I/dbgcAAAAAANjwf/7P/9G3336rM2fO6Msvv5zotT3tdrs90SsCDjl16pRarZbbbQAAAAAAgBPwer369ttvJ3pNZoRhZnWCMK/Xq2AwONZrNRqNiV3LLnqyh55mrx+Jnuyip9nrR6Inu+hp9vqR6MkuerJn2nqatn4kerKLnqavn871Tp06NfZrHUQQhpn1x3/8x9rZ2VEwGNT29vZYrxUKhSZ2LbvoyR56mr1+JHqyi55mrx+Jnuyip9nrR6Inu+jJnmnradr6kejJLnqavn4613NjmyNXg7Bnn33Wzcv38Xg8+t//+3+73QYAAAAAAADGxNUgrF6vy+PxyM1tyjrX93g8rvUAAAAAAACA8fO63YDbOCsAAAAAAABgMbi+R1i73VYymVQ4HHbl+qZpqlQquXJtAAAAAAAATI7rQVipVNIPf/hDV3tIJpP6wQ9+4GoPOLlGo6FQKGRZk81mlc1mJ9QRAAAAAACLqVgsqlgsWtY0Go0JdXOY60FYPB53uwVdvHjR7RYwglarpZ2dHcuaZrM5oW4AAAAAAFhczWZz4GN0N7kehC0tLbndAmac1+tVMBi0rPH5fBPqBgAAAACAxeXz+bS8vGxZ02g01Gq1JtRRP1eDME5qhBOCwaC2t7fHeo1sNqtmszlVgRo9za5p+zpNWz8SPc2yafs6TVs/Ej3Nsmn7Ok1bPxI9zbJp/DpNW0/T1o80nT1No2n8Ok1bT072Y2drolAo5NqsMU/bxWMTvV6vTNN0/Rv/8OFDBQIB19JInEznF2d5eXnsQRhmFz8nmGf8fGOe8fONecbPN+YZP9+ww82fE+9Er3ZAoVBwPQSTpLNnz6pQKLjdBgAAAAAAAMbI1SDsxz/+8USv9/nnnx/7sUn3AgAAAAAAgMlyNQibtEwm43YLAAAAAAAAcMlCBWGbm5tutwAAAAAAAACXuHpq5El9/vnnMk3Tdr1hGFpfXx/qcwAAAAAAADBfZiII+/zzz1UoFFStVmUYxonGaLfb8ng8DncGAAAAAACAWTH1Qdi1a9e0trYmaT/MAgAAAAAAAE5iqoOwX/ziFyoUCpIkj8cjj8dDGAZgKNlsVs1mUz6fz+1WAMfx8415xs835hk/35hn/Hxj2nnaU5wsXbhwQbVarRuAhcNhRaNRhcNhSdIzzzwzcIzf//73Mk1Td+7c0cOHD/Xtt9+Ou21MSCgU0s7OjpaXl7W9ve12OwAAAAAAwAY3H89P9YywTggmSRsbG7p06dKJx8rlcnr22Wedag0AAAAAAAAzxut2A1b8fr+k/RBrlBBMksLhsM6dO+dAVwAAAAAAAJhFUx2EdZZAXrx40ZHx1tfXHRkHAAAAAAAAs2eqg7DOLLDd3V1HxwMAAAAAAMDimeog7Cc/+Yna7bZqtZoj47377ruOjAMAAAAAAIDZM9WnRkrS2tqaCoWCPv/8c/2bf/NvRhrr4sWLunfvnkOdwW2dUya8Xq+CwaBlbTabVTabnVBnAAAAAAAspmKxqGKxaFnTaDTUarU4NfIouVxOGxsbSqVS+uUvf3nicR48eODYzDJMl1arpZ2dHcuaZrM5oW4AAAAAAFhczWZz4GN0N019ECZJGxsbunDhgp599lnl83ktLS3Z+rzd3V2Zpql6va47d+6MuUu4xc6MMJ/PN6FuAAAAAABYXD6fT8vLy5Y1nRlhbpj6pZGS9Ktf/Uq5XE61Wk0ej+dEY7TbbXk8Hn377bcOdwe3dJZGujGVEgAAAAAAnIybj+enfkbYX/3VX6lUKkmSPB6PTpLbnTQ8AwAAwPRot6XWgP8Kej0S//UDAADHmeog7ObNm1pfX5d08hBM0ok/DwAAANOj1ZZOv21d8/i6dIogDAAAHGOqg7CDIVg8HlcikZDf7x96n7AbN27o888/H2O3AAAAOIhZXAAAYJpMdRDW2RPM7/drc3NT586dO/FYKysreuaZZxzsDgAAAIMwiwsAAEwTr9sNWPH7/ZKka9eujRSCdcY6f/68A10BAABgWr1/3+0OAADANJvqICwWi0mSwuGwI+PdvHnTkXEAAAAweXZCrtcrhGEAAOB4Ux2EZTIZtdttGYbhyHjMCAPgBNM03W4BAOaKneDq4SPpzY/sjffGh1Lz0UgtAQCAOTXVQVgymdT58+d1+/ZtR8b74IMPHBkHwOKp1WrK5/MKBAJaWVlxux0AmBlOzeJ6ryZ9+Y29a37xtXSrZq8WAAAslqkOwiSpXC5ra2tL//iP/zjyWKurqw50BGDadEKqRCKhSCSiQCAgj8ejQCCgWCymRCKhfD6varV65OenUqnuUuwO0zRVqVSUSqW646ytrTEbDACG4OQsro8/G+7anwxZDwAAFsNUnxop7e8PduPGDf3whz/UX/7lX554nIcPH6pW46lBTKdSqaRMJjP050WjUW1tbY18fc8Jzqz3+/3a29s78mOpVEqVSmXoMTc2NhSPx23Vmqap1dVVra2tWdZ0fu+r1Wq3NhqNKh6PKxKJaGNjQ5VKRdFotO9zV1ZWZBgGfzcAYAQnmcX1o+8f/fG9r4a79rD1AABgMUx1ENZZyvhHf/RHCgQCevbZZ5VMJocexzRN3blzx+n2AMek02lduXJFhmFodXXVMkSKx+PK5/O6cOFC92TVUXX24qtWq8rn88fOegqHw8rn84rH45aHWJTLZZmmqc3NTRUKhWNnYklSLpfT1atXFQ6Hbd9PPp8/MgBLJpPdsTrjGYYhwzBULpd1586dbjh2MOA6eM/lcrn7eiQScWyvQgBYJCeZxXVcEBZ4crixhq0HAACLwdNut9tuN3GcCxcu6P79P2wY0W63TzRzpfdzv/32W6fag8tCoZB2dna0vLys7e1tt9txVCaTUalUOvT+ZDLZF9CMg2EYikQiR35sb2/vROFbLBY7cmbV+vq60un0UL0lEolDoVQ6ndb6+rqtMdbW1pTP5w+9PxwOq16vH/k5B78fk/g+AMA8+Hf/r3RviH+i/11I+v/+n6M/9vf/JP31x/bH+rvLx4dqAADAXW4+np/qPcKuXLmidrutTlZ30hAMmDXHLZM8yfLJYYXD4UPLBKX98OekM9CO63uYEKxarR6ameX3+7W1tWU7BJP2Z6DV6/VD97K7u3vs5zg18w4AFo2Ts7hejUpPfcfeOE8/Ib12+J8yAACA6Q7COg+ePR5PNxA76QswS45bdmi1HHHc1x/l2kd97jDhUq1WUyKROPT5d+/ePTK0s9PPgwcP+npgE3wAcN7l54arf96i/uwZ6Z0X7I1z40XJd2a4awMAgMUw1XuEnT17VtFoVPf
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"$t_y$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_ty_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACdw0lEQVR4nOzdTWxbZ5rm/etQjuEgbopUGg0IEtAVctJwrXpM2rPIbkokahMjmQrpADFiAy/KZGKkZyFMyLiy6JWjUFWlTTcyJl39Aq/TBsYmqzpBsqix6PSuF2PpuHpnTDVpFyCBwAAlUuoEMRKb511oyNIneUjx8PDj/wMI6OPWc25JdGJefp77GJZlWQIAAAAAAABGnMftBgAAAAAAAIB+IAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYOOZ2A0A3XnjhBT158kQTExP6i7/4C7fbAQAAAAAANv2f//N/9OzZM504cULffPNNX69tWJZl9fWKQA9MTEyoXq+73QYAAAAAAOiSx+PRs2fP+npNdoRhKDWCMI/Ho+npacevV6lU+no9O+hp+PqR6MkuerJn0HoatH4kerKLnuwZtJ4GrR+Jnuyip+HrR6Inu+jJnkHrqd/9NK43MTHh+LX2IgjDUPqLv/gLra+va3p6Wmtra45fb3Z2tq/Xs4Oehq8fiZ7soid7Bq2nQetHoie76MmeQetp0PqR6Mkuehq+fiR6soue7Bm0nvrdT+N6bow6Ylg+AAAAAAAAxoKrO8JefvllNy+/i2EY+t//+3+73QYAAAAAAAAc4moQViqVZBiG3JzX37i+YRiu9QAAAAAAAADnjf3RSG6aCQAAAAAAMB5cH5ZvWZZisZgCgYAr16/Vasrlcq5cGwAAAAAAAP3jehCWy+X005/+1NUeYrGYfvzjH7vaAwAAAAAAAJzl+tHISCTidgs6e/ZsX69XLpcVjUaPvE4wGJRhGCoUCh1/7eLioqLRqPx+vwzDUDAYVDweV7FY7LofJ9YEAAAAAADoFdd3hE1NTbndQs/UajX5/X5btYlE4kjXSqfTKpfLHX9dsVhUPB5XrVZTJBJRPp9XIBCQaZpKp9OKRqPNj/t8PtfWHDTz8/Pa2tqS1+t1u5UmehpOg/gzoid7BrGnQTOIPyN6smcQexpEg/ZzGrR+JHoaZoP2cxq0fiR6GmaD+HMatJ4GrR8nGZaL0+InJiZUrVZd/0Fvbm5qampKz549O9I6i4uLSqfTtmpLpVLXc9GKxeKuHWX5fF6xWKyjr0skEspms/tqwuGwTNNUIBDQ6upq2+DKiTXtmJ2d1fr6umZmZrS2tnbk9TCaeJ5g1PEcxyjj+Y1RxvMbo4znN+xw83ni6tHIQbpjYy96WVhYsFUXiUS6DsFqtZri8fiRvi4QCBwYWEnboZq0fXyz3XWcWBMAAAAAAMAprh6NzGQyru8Gk6TJyUllMpkjrZHL5VSr1ZRKpdrO/zpz5kzX17l8+XLzOGmtVrP9dY2ji5Ja7loLBAKKxWIqFAoqFovK5XKHHuN0Yk0AAAAAAACnuHo0st8eP36sH/zgB46sHQwGJW0feXRKLpdTMpnU6uqq5ubmmiFUu6OR5XK52Z8kVavVlscTC4VCc+eWz+dTtVrty5qdYLst7OB5glHHcxyjjOc3RhnPb4wynt+wY2yPRvZbMpl0ZN1CoaByuWx7Plg3yuWyksmkUqmUQqFQR1+7c7dbJBJpO6NrZ6hWq9UOvCulE2sCAAAAAAA4aayCsJWVFUfWXVhYkM/n0/nz5x1ZX9o+hhgKhbo6wpnL5Zpv2w3Rds4wu337dl/WBAAAAAAAcJKrM8K69fjx447mY5XLZWWz2Y6+xi7TNGWapiTJ7/crEAgoEokoGo3aupOjHel0WqZpdnXsstFbw9mzZ219XSgUUrlclqR9u7ecWBMAAAAAAMBpQxGEPX78WJlMRsVisRmkdMqyLBmG0ePO9g+JL5fLyuVyzR1TsVhMV69e7fg4Y4NpmlpcXFQ2m+3qTpPFYnHX+3bX2Ftnmmbze3BiTQAAAAAAAKcNfBB29epVLS4uStoOswZJuVzeFwrtVSgUVCgUlEgklM1mO77G3NycYrFY13dZvH///q73283yanjxxRd3vb+ystIMrZxYE3DC/Py8tra2BuLutIATeI5jlPH8xijj+Y1RxvMbg26gg7Bf//rXzZlYhmHIMIyBCsMCgUDzyGWpVGq5Yy2Xy2llZUWrq6u212/cZfHGjRtd97i3n253b+08lunEmoAT5ufn3W4BcBTPcYwynt8YZTy/Mcp4fmPQDXQQtrCwIEnNACwQCCgUCjUDlb07jA7yxz/+UbVaTXfu3NHm5mbPe9y7U6tWqymXy2lhYWHfTDLTNBWNRrW8vNx23cZOsuXlZds7rg7S7VHSvXZ+L06sCQAAAAAA4LSBDsJM02zO9VpeXtbc3FzXa6VSKb388su9au1QPp9PqVRKqVRKhUJBly9f3hX4FItFLS4uKpVKHbpGrVZTPB5XIpFQJBI5Uj/dhk17w7eNjQ1H1+xWpVLR7OzskdeZn5/nXy4AAAAAAGhhaWlJS0tLR16nUqn0oJvuDHQQ5vP5tLm5qVQqdaQQTNo+lvfSSy/1qDN7YrGYIpGI5ubmdt1pcWFhoWUQNjc31zx2OSic2L3VizXr9brW19ePvM7W1taR1wAAAAAAYJRtbW315DW4mwY6CAsEAnrw4IHOnj3bk/XcCJZ8Pp9WV1cVDoebYVitVlOxWDxwt9fi4qJM0+xolli76/cicNq5m8uJNbvl8Xg0PT195HUY5AgAAAAAQGter1czMzNHXqdSqaher/ego84NdBA2NzenBw8e9OQIXWM9t9y4cUPhcLj5/vLy8r4gzDRNpdNpZTKZnt1NcWpqqieh1dTUlKNrdmt6elpra2tHXgcAAAAAALTWq7FCs7Ozru0s87hyVZt+9rOfybKsXccKj+JXv/pVT9bpRigU2hV8HTRwPh6PKxQKtTw22alud13tDbr27gjr9ZoAAAAAAABOG+gdYZOTk/r444+VyWSUyWT0Z3/2Z0daL5vN6qc//WmPuutcNBpVsVg88HOLi4sql8uKRCKKx+Nt19oZKi0sLOj27dvN9998803FYjFJ0pkzZ3YFibVazVYAtXcXXjAYbL7txJoAAAAAAABOG+ggTNq+2+Py8rLi8bh++9vfdr3Oo0ePerazrFuBQKD59t5jgX/84x8l6dCgrBXTNHd9b4FAoBmE7TyOKW3vRLNz7LJUKu16f+duNifWBAAAAAAAcNrAB2HS9jytM2fO6OWXX1Y6nbY9W2pjY0O1Wk2lUkl37txxuMv2dgZh/ToWeObMmV3v2w2tdu448/l8u3p3Yk0AAAAAAACnDUUQ9tVXX0na3lGUTCa7WsOyLBmG0cu2OraystJ8OxqN7vpc4/inXcFgsDlnLJ/PN3eA7RUKhXbd5fH+/fuH1h7W697gy4k1AQAA2rEsqW61rvEYkst/5QMAAANs4IOwd999V7lcTpJkGIYsq83ffg7gdgDWsPNoYD+PBZ4/f775M7R7PHRnXTqd7suaAAAArdQt6diHrWueXpMmBuOvfgAAYAAN9F0jb9y4oWw22wy/ugnBjvJ1vVYoFCSpp3eFtGPnLjo7M8h21gQCgQNDOyfWBAAAAAAAcNJA7wjLZrOS/rQTLBKJKBqNyufzdTwn7Pr163r8+LGD3bZWKBRULpfl8/l09erVvl47FAopEok0w6h
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACdWElEQVR4nOz9TWxbd57n/34oO4Fzk6YppTeCBFSF7Axcq4JJexZ1V1MiuxY3RtIV0gaSSQL8OyKToO4sNB0yrqwdhaoqbeb/T5t0zSycDjA2WTUJkgvUWHR6N4uxRGd2wVTzOA1I4CoirUoQI2Xz3IWGLFEP5CF1Hkie9wsgoocvf+dLU47FD38PAdM0TQEAAAAAAAATbsrrBgAAAAAAAAA3EIQBAAAAAADAFwjCAAAAAAAA4AsEYQAAAAAAAPAFgjAAAAAAAAD4AkEYAAAAAAAAfIEgDAAAAAAAAL5AEAYAAAAAAABfIAgDAAAAAACALxCEAQAAAAAAwBcIwgAAAAAAAOALBGEAAAAAAADwBYIwAAAAAAAA+AJBGAAAAAAAAHyBIAwAAAAAAAC+QBAGAAAAAAAAXyAIAwAAAAAAgC8QhAEAAAAAAMAXCMIAAAAAAADgCwRhAAAAAAAA8AWCMAAAAAAAAPgCQRgAAAAAAAB8gSAMAAAAAAAAvkAQBgAAAAAAAF8gCAMAAAAAAIAvEIQBAAAAAADAFwjCAAAAAAAA4AsEYQAAAAAAAPCFk143AAzjySef1HfffSfTNDU11TvP/au/+is99dRTLnUGAAAAAIB/ffPNN/rTn/7Us6bVaikQCOiJJ57Qt99+61JnuwKmaZquXhGwwYkTJ9RqtbxuAwAAAAAADGlqakqPHj1y9ZrMCMNY2huEDTMjrF6vq9VqaWpqSrOzs471uZ9X1/Xrtf34mP16bT8+Zi+v7cfH7OW1/fiYvby2Hx+zl9f242P28tp+fMx+vbYfH7OX1/bjY+51baszwqTd1/auM4ExNDc3Z0oy5+bmPLn/sLy6rl+v7cfH7Ndr+/Exe3ltPz5mL6/tx8fs5bX9+Ji9vLYfH7OX1/bjY/brtf34mL28th8f83Gv7WXfbJYPAAAAAAAAX/B0aeSzzz7r5eW7BAIB/e///b+9bgMAAAAAAAAO8TQIq9VqCgQCMj3cr799/UAg4FkPAAAAAAAAcJ7vl0Z6GcIBAAAAAADAPZ6fGmmappLJpMLhsCfXbzabKhaLnlwbAAAAAAAA7vE8CCsWi3r99dc97SGZTOpnP/uZpz0AAAAAAADAWZ4vjYzH4163oPPnz3vdAgAAAAAAABzm+YywmZkZr1uADy0tLWlnZ0fBYNDrVlzj5WP26tp+fMxe47n2z7W94sc/bz8+Zi/59c/bjz/fXvLjn7dfr+0Vv/55+/HvlpfG9XEHTA93iz9x4oQajYbnf2j379/XzMyMHj165GkfsG5+fl5bW1uam5vT5uam1+1YNq59Y7zwcwan8TMGN/BzBqfxMwY38HMGp43rz5iXfXu6NHKUTmwcpV4AAAAAAABgP0+DsHw+7/lsMEk6ffq08vm8120AAAAAAADAQZ4GYW+//bar1/vqq6+O/J7bvQAAAAAAAMBdnp8a6aZMJuN1CwAAAAAAAPCIr4Kw9fV1r1sAAAAAAACAR0563cAwvvrqKzWbTcv1hmGoUCgMdB8AAAAAAABMlrEIwr766ivl83lVKhUZhjHUGKZpKhAI2NwZAAAAAAAAxsXIB2GXL1/WysqKpN0wCxhnS0tL2tnZGYnTUjG5+DmD0/gZgxv4OYPT+BmDG/g5g9P4GRtcwBzhdOl3v/udUqmUJHVmcx2n3UAgoEePHtnSG7w1Pz+vra0tzc3NaXNz0+t2AAAAAACARV6+ph/pGWHLy8uSdgMs0zQVDocVjUYVDoclSU8//XTfMb7++ms1m03dvHlT9+/fd7RfAAAAAAAAjK6RDsKq1WpnJtja2poWFhaGHiubzerZZ5+1qzUAAAAAAACMmSmvG+glFApJ2g2xjhOCSVI4HNYzzzxjQ1cAAAAAAAAYRyMdhLWXQJ4/f96W8QqFgi3jAAAAAAAAYPyMdBDWngW2vb1t63gAAAAAAADwn5EOwn75y1/KNE1Vq1Vbxvvtb39ryzgAAAAAAAAYPyMdhJ0+fVrvv/++bty4oT/96U/HHo+lkQAAAAAAAP410kGYtLtRfjQaVSqVOtY49+7ds21mGQAAAAAAAMbPSa8bsGJtbU3nzp3Ts88+q1wup5mZGUv3297eVrPZVK1W082bNx3uEl6o1+uan5/vWbO0tKSlpSWXOgIAAAAAwL9WV1e1urras6Zer7vUzUFjEYR9/vnnkqRaraZMJjPUGKZpKhAI2NkWRkCr1dLW1lbPmp2dHZe6AQAAAABMAtOUWmbvmqmARMxw0M7OTt/X6V4a+SDszTffVLFYlCQFAgGZZp+fxEMQgE2uqakpzc7O9qwJBoMudQMAAAAAmAQtUzr5bu+ah1ekE8QNBwSDQc3NzfWsqdfrarVaLnXULWAOkyy55Nq1a50ZYMOGYHsFAgE9evTIjtbgsfn5eW1tbWlubk6bm5tetwMAAAAAmCCPWhaDsJHfeX00efmafqRnhLVPeWyHYPF4XIlEQqFQaOB9wq5evaqvvvrKwW4BAAAAAAAwykY6CKtWqwoEAgqFQlpfX9czzzwz9FiLi4t6+umnbewOAAAAAAAA42SkJ/GFQiFJ0uXLl48VgrXHOnv2rA1dAQAAAAAAYByNdBAWi8UkSeFw2Jbxrl27Zss4APyt2Wx63QIAAAAAYAgjHYRlMhmZpinDMGwZjxlhAIZVrVaVy+U0PT2txcVFr9sBAAAAAAxhpPcISyaTOnv2rG7cuKF/+Id/OPZ4v//97/Xzn//chs4AjJJqtaobN26oWq3KMIzOIRmhUEjhcFgzMzOKRqNKJBKKx+MH7p9KpWQYhjY2NjpfazabqlQqunHjhiqVCrPAAAAAAGACjHQQJkmlUkl/8zd/o//23/6b/u7v/u5YYy0vLxOEYSQVi0VlMpmB7xeNRrvCm2EFAoGB7xMKhdRoNA79XiqVUrlcHnjMtbW1Q4OqwzSbTS0vL2tlZaVnTbValSRVKpVObTQaVTweVyQS0dramsrlsqLRaNd9FxcXZRhG5/4AAAAAgPE30ksjpd39wa5evarXX3/9WOPcv3+fF7QYWel0Wo1GQxsbG0omkz1r4/G41tbWOvV2ME1TtVpNhUKhc0jFYcLhsAqFgmq12pEhmLQbYDcaDUvBVjab1cbGhhqNhuUQrL1EcX8IlkwmVSqVOuO1H9fa2prS6XTnsVWrVa2srCiTyXQCu/0zvtrjmKZp2z6FAAAAAABvBUzTNL1u4ii///3vOx9ns1kFAoG+IcFhms2mbt68qWazqUePHtnZIjwyPz+vra0tzc3NaXNz0+t2bJfJZFQsFg98vR30OMkwDEUikUO/12g0egZlR4nFYocG0YVCQel0eqDeEonEgX0D0+m0CoWCpTFWVlaUy+UOfD0cDqtWqx16n/3PhxvPAwAAAADvPGpJJ9/tXfPwinRi5KcXjSYvX9OP9NLI9957T3fv3u18bppmz2VQvZimOdTyL8ALRwVhwyyfHFQ4HFY0Gj0QXCWTyaFCMGm378N6HyQEq1QqSiQSXV8LhUK6ffv2gWWNvWSzWSWTScVisa5ZYNvb20feZ9jHDQAAAAAYLSOdXV68eFGmaao9aY0gC35x1FI8t5boHXad41z7sPsOEi5Vq1VbQrC9/dy7d6+rBzbDBwAAAIDJN9JBWHsGSSAQ6ARiw96AcXJUSDQzM+NuIyPAMAwtLCwc+HqpVBoqBGsLhUK27bE2Ssrl8oGlo/iLSqXidQsAAAAAPDTSSyNPnz6taDSqu3fvKp/PKxqNDhUEGIah9957T1988YX9TQIT6LC/Z08//bSt41n9u5xKpQ7M1kqn05Y31u8lHA4rm80OveTaTdVqVYVCQZVKRYZhKBQKaWZmRqFQSOFwWJcuXZK0e9plPp/vuey0UqkolUppZmZGhULBlj/LUWUYhiqVikqlUicEG6U3R/z
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_x\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACf1ElEQVR4nOz9T4wieX7n/7/I7rXGsk0F2dZKKZB2OnBbsye7IWsPe/MkrC+2dj0NVZIt24edAvdotQfkga7x2ZMD4+Gy0m5Djg+/8c/SZsGMbdkHf5uo9n0rkx7frF0T1SMlQlppMqMYWx7ZbuJ7qC8YMhMI/iQBGc+HlBJkfvjEmz9Z3fHK9+cTIdd1XQEAAAAAAAD33J7fBQAAAAAAAACbQBAGAAAAAACAQCAIAwAAAAAAQCAQhAEAAAAAACAQCMIAAAAAAAAQCARhAAAAAAAACASCMAAAAAAAAAQCQRgAAAAAAAACgSAMAAAAAAAAgUAQBgAAAAAAgEAgCAMAAAAAAEAgEIQBAAAAAAAgEAjCAAAAAAAAEAgEYQAAAAAAAAgEgjAAAAAAAAAEAkEYAAAAAAAAAoEgDAAAAAAAAIFAEAYAAAAAAIBAIAgDAAAAAABAIBCEAQAAAAAAIBAIwgAAAAAAABAIBGEAAAAAAAAIBIIwAAAAAAAABAJBGAAAAAAAAAKBIAwAAAAAAACBQBAGAAAAAACAQCAIAwAAAAAAQCAQhAEAAAAAACAQ3vS7AGAZP/VTP6Uf//jHeuONN/Sv//W/9rscAAAAAADg0f/9v/9Xn332mT73uc/p7//+7zd67JDruu5GjwiswRtvvKHBYOB3GQAAAAAAYEl7e3v67LPPNnpMOsKwk4ZB2N7eng4ODu70WL1eb2PH2rbjB/XYfh8/qMf2+/gcm/c8KMf2+/hBPbbfxw/qsf0+flCP7ffxg3psv48f1GP7ffxdPfbwsW+88cYdVTeDC+ygaDTqSnKj0ei9Ota2HT+ox/b7+EE9tt/H59i850E5tt/HD+qx/T5+UI/t9/GDemy/jx/UY/t9/KAe2+/j7+qx/aybzfIBAAAAAAAQCL4ujXznnXf8PPyEUCik//2//7ffZQAAAAAAAOCO+BqEdTodhUIhuT7u1z88figU8q0GAAAAAAAA3L3AL430M4QDAAAAAADA5vh+1UjXdZXJZGSapi/HdxxH9Xrdl2MDAAAAAABgc3wPwur1ur785S/7WkMmk9Ev//Iv+1oDAAAAAAAA7pbvSyNTqZTfJejhw4d+lwAAAAAAAIA75ntH2P7+vt8lAFutUCio3+8rHA5z7AAd3y9+P++gfub8ft39FNTX3e/33O/j+8Xv5x3Uz5zfr7ufgvy6B/m5+ymor7vf77nfx/fLrj7vkOvjbvFvvPGGrq6ufH/RXr16pf39fX322We+1gHvYrGYut2uotGoLi4u7s2xAInPHDaLzxs2jc8cNo3PHDaJzxs2bVc/c37W7evSyG26YuM21QIAAAAAAID18zUIK5fLvneDSdKDBw9ULpf9LgMAAAAAAAB3yNcg7Ktf/epGj/fpp59O/dmmawEAAAAAAMBm+X7VyE3K5/N+lwAAAAAAAACfBCoIOzs787sEAAAAAAAA+ORNvwtYxqeffirHcTyPt21btVptoccAAAAAAADgftmJIOzTTz9VuVyWZVmybXupOVzXVSgUWnNl8Fuv11MsFps5plAoqFAobKgiAAAAAACCq1qtqlqtzhzT6/U2VM1NWx+EPX36VJVKRdLrMAsYNxgM1O12Z47p9/srHaNQKKjf72/FFU4RDHzmsEl83rBpfOawaXzmsEl83rBp2/iZ6/f7c8/T/RRytzhd+u53v6tsNitJo26uVcoNhUL67LPP1lIb/BWLxdTtdrW3t6eDg4OZY+kIAwAAAABgM7x2hA0GA0WjUV1cXGyoste2Ogg7PDxUu91WKBSS67oyTVOJREKmaUqS3nrrrblz/PCHP5TjOHr27JlevXpFEHZPDIMwP35pAAAAAADA8vw8p9/qpZHDEEySWq2Wjo6Olp6rWCzqnXfeWVdpAAAAAAAA2DF7fhcwi2EYkl6HWKuEYJJkmqbefvvtNVQFAAAAAACAXbTVQdhwCeTDhw/XMl+tVlvLPAAAAAAAANg9Wx2EDbvALi8v1zofAAAAAAAAgmerg7Cvfe1rcl1X7XZ7LfN9+9vfXss8AAAAAAAA2D1bHYQ9ePBA3/jGN3R6eqof/ehHK8/H0kgAAAAAAIDg2uogTHq9UX4ikVA2m11pnpcvX66tswwAAAAAAAC7502/C/Ci1Wrp8PBQ77zzjkqlkvb39z097vLyUo7jqNPp6NmzZ3dcJQAAAAAAALbZTgRhH3/8sSSp0+kon88vNYfrugqFQussCwAAAAAAYGNcVxq4s8fshSTij+m2Pgh7//33Va/XJUmhUEiuO+cdvwUBGAAAAAAA2HUDV3rz92aP+effl94gBplqq/cIOzk5Ua1WG4Vfy4RgqzwOAAAAAAAA98dWd4QNr/I47ARLpVJKp9MyDGPhfcI+/PBDffrpp3dYLQAAAAAAgL/++BPpt5J+V7G9tjoIa7fbCoVCMgxDZ2dnevvtt5ee68mTJ3rrrbfWWB0AAAAAAMDm/PEn88f8dlN6Y0/6jXfvvp5dtNVLIw3DkCQ9ffp0pRBsONe77/IpAAAAAAAAu+fVj6Wv/Jm3sb/zp1L/x3dazs7a6iAsmXzdy2ea5lrmOzk5Wcs8AILNcRy/SwAAAAAQMH/Ulv7+n7yN/bt/lL7Tvtt6dtVWB2H5fF6u68q27bXMR0cYgGW1222VSiVFIhE9efLE73IAAAAABMyf/81i4/9iwfFBsdV7hGUyGb377rs6PT3V7/7u76483/e+9z196UtfWkNlALZJu93W6emp2u22bNseXSTDMAyZpqn9/X0lEgml02mlUqkbj89ms7JtW+fn56PvOY4jy7J0enoqy7LoAgMAAADgq6t/uNvxQbHVQZgkNRoN/dzP/Zz+5E/+RL/2a7+20lzHx8cEYfdMr9dTLBabOaZQKKhQKGyoouXU63Xl8/mFH5dIJCbCm2WFQqGFH2MYhq6urm79WTabVbPZXHjOVqt1a1B1G8dxdHx8rEqlMnNMu/26H9iyrNHYRCKhVCqleDyuVqulZrOpRCIx8dgnT57Itu3R4wEAAADAT5GfvNvx61KtVlWtVmeO6fV6G6rmpq0PwkzT1Icffqgvf/nLKwVhr1694oT2HhoMBup2uzPH9Pv9DVWzvFwup0ePHsm2bR0fH88MkVKplEqlkg4PD0cXlFjVcAmyZVkqlUpTu59M01SpVFIqlZq5d1+j0ZDjODo7O1O5XJZlWVPHFotFPX78WKZpen4+pVLp1gAsk8mM5hrOZ9u2bNtWo9HQs2fPRuHY9X8Prj/nRqMxuh2Px9e2RBsAAAAAlvGrX5A++j/ex//KF+6ulln6/f7c83Q/bXUQ9r3vfU+S9LM/+7OKRCJ65513lMlkFp7HcRw9e/Zs3eVhC+zt7eng4GDmmHA4vKFqVmMYhhKJhBqNhvL5vOr1+o0xmUxmIqBZJ9M0lcvlRp1Stzk/P/ccVhmGoVQqpVQqpWQyeWsQXavVlMvlPNdo27bS6fSNUCqXy6lWq936mGEolkqlVKvVVKlUVCqVPB9Teh0+3vZ+AAAAAMCm/GZC+uAvvW2Y/9M/If1WYv64uxAOhxWNRmeO6fV6GgwGG6poUsh1XdeXI3tweHioTz75ZHTfdd2llnCNP/azzz5bV3nwUSwWU7fbVTQa1cXFhd/lrF273R5dNXXcIksHV3FbcLVKCDdt6eci//xYlqV0Oj3xPcMw9Pz58xvLGuexbVvJZHKiC2zWUs/rHWh3GUgCAAAAwDTfOZd+28MuNP//R9JvbPH1Av08p9/qq0Y+evRIruuOTpaXDcGAXTNt2eGs5Yh3ffxVjn3bYxdZ1tlut9cWgg3refny5UQNbIYPAAAAYNt5Cbe+k9nuEMxvWx2EDTtIQqHQKBBb9gvYJdNCov39/c0WsgVs29bR0dGN7zcajaVCsCHDMNZyoYFt02w22c9
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"x_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACtMElEQVR4nOz9XWxj953m+z6LVTHKHTeLkhsDCCTQXVR74L453UWq9kX6qlvkzsXYcLpNVgHx2AZmUpTbJzsXnIloxRj0lSNTnRAYzD5GkeUe4JTbwC6TcduwZyOx6MzdBs6UxErfGZOQ5QAiCAzG0iq2DRf8wnUuNGRISiIX39biy/cDECGlH//rp5dy93r0fzEsy7IEAAAAAAAAzDmP2w0AAAAAAAAATiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQiAIAwAAAAAAwEIgCAMAAAAAAMBCIAgDAAAAAADAQjjvdgPAsL75zW/qwYMHOnfunP7Fv/gXbrcDAAAAAABs+B//43/o66+/1oULF/TZZ585em3DsizL0SsCY+LxeMSvLwAAAAAAs8kwDDUaDUevyYwwzKz2EMzv97vYyWBqtZoajYY8Ho9WVlbcbsc2+nYWfTuLvp1F386ib2fRt7Po21n07Sz6dhZ9O6tarUqSK5NbCMIwszweT+sf/MHBgdvt2BYIBFStVrWyskLfDqBvZ9G3s+jbWfTtLPp2Fn07i76dRd/Oom9nzWrf586da93PO83VIOyxxx5z8/IdDMPQf//v/93tNgAAAAAAADAhrgZh5XJZhmG4us9T8/qGYbjWAwAAAAAAACbP+TloU4bN1gEAAAAAABaD63uEWZalWCymYDDoyvVN01Qul3Pl2gAAAAAAAHCO60FYLpfT9773PVd7iMVi+va3v+1qDwAAAAAAAJgs15dGRiIRt1vQlStX3G4BAAAAAAAAE+Z6ELa8vOx2CwAAAAAAAFgAri6N5KRGjOL3f//3df/+ff3+7/++260shGQyqXq9Lq/X63YrA6FvZ81q37NqVr/f9O2sWe17Vs3q95u+nTWrfc+qWf1+07ezZrXvWeXm/bxhuXhsosfjkWmarv+i3b9/X0tLS2o0Gq72gcEEAgFVq1X5/X4dHBy43Y5ts9o3MM/4dwlMH/5dAtOHf5fA9JnVf5du9u3qjLB0Ou16CCZJFy9eVDqddrsNDKlWqykQCPSsSSaTSiaTDnUEAAAAAMBiymQyymQyPWtqtZpD3ZzkahD2wx/+0NHrffzxx/qjP/qjqegF49NoNFStVnvW1Ot1h7oBAAAAAGBx1ev1vvfobnI1CHPaxsaGfvGLX7jdBsbM4/FoZWWlZ800zDwEAAAAAGDeeb1e+f3+njW1Ws217akWKgjb29tzuwVMwMrKykythQYAAAAAYF7Z2ZqouUeYG2YyCPv4449lmqbt+kqlomw2O9B7AAAAAAAAMF9mIgj7+OOPlU6nVSwWValUhhrDsiwZhjHmzgAAAAAAADArpj4I29ra0s7OjqTjMAsAAAAAAAAYxlQHYT/72c+UTqclSYZhyDAMwjDMvGQyqXq9zgb+wBTh3yUwffh3CUwf/l0C04d/l4MzrClOltbW1lQqlVoBWDAYVCgUUjAYlCQ9+uijfcf45JNPZJqm3nrrLd2/f19ff/31pNuGQ5qb6/n9fjbLBwAAAABgRrh5Pz/VM8KaIZgk7e7uan19feixNjc39dhjj42rNQAAAAAAAMwYj9sN9OLz+SQdh1ijhGCSFAwGdenSpTF0BQAAAAAAgFk01UFYcwnklStXxjJeNpsdyzgAAAAAAACYPVMdhDVngR0eHo51PAAAAAAAACyeqQ7CfvSjH8myLJVKpbGM9/rrr49lHAAAAAAAAMyeqQ7CLl68qFdffVW3b9/WP//zP488HksjAQAAAAAAFtdUB2HS8Ub5oVBI8Xh8pHHu3bs3tpllAAAAAAAAmD3n3W7Ajt3dXa2tremxxx5TKpXS8vKyrfcdHh7KNE2Vy2W99dZbE+4SAAAAAAAA02wmgrBf/vKXkqRyuayNjY2hxrAsS4ZhjLMtAAAAAAAAzJCpD8L+5m/+RrlcTpJkGIYsyxp4DAKw+Var1RQIBHrWJJNJJZNJhzrqZFlSo8+vrceQ+DUFAAAAAMy6TCajTCbTs6ZWqznUzUlTHYTdvHmztcH9sCGYpKHfh9nQaDRUrVZ71tTrdYe6OalhSedf7l3z1SvSOYIwAAAAAMCMq9frfe/R3TTVQVh3CBaJRBSNRuXz+QbeJ+zGjRv6+OOPJ9gt3OLxeLSystKzxuv1OtQNAAAAAACLy+v1yu/396yp1WpqNBoOddTJsKZ4upTH45FhGPL5fNrb29OlS5eGHss0TT366KP6+uuvx9gh3BQIBFStVuX3+3VwcOB2O2f6umFzRtjUn+EKAAAAAMDo3Lyfn+pbb5/PJ0na2toaKQRrjnX58uUxdAUAAAAAAOA8yzqebNHrMb3TnabDVC+NDIfD+uUvf6lgMDiW8W7evDmWcQAAAAAAAJzGHtSjm+oZYRsbG7IsS5VKZSzjMSMMwDiYpul2CwAAAACAIUx1EBaLxXT58mXdvn17LOO9/fbbYxkHGLc377rdAfoplUpKpVJaWlrS9evX3W4HAAAAADCEqV4aKUn5fF5//Md/rH/8x3/UX/3VX4001vb2tv76r/96TJ0B9tgJuZ4vHG+W/wyTFodSKpV0+/ZtlUolVSqV1mmxPp9PwWBQy8vLCoVCikajikQiJ94fj8dVqVS0v7/f+phpmioWi7p9+7aKxSKzwAAAAABgDkx9EBYMBnXjxg1973vfGykIu3//vkql0hg7A/q7/0B68V0bhb/K6V+/uqF/PeD4oVCoI7wZlmEMvoDc5/Pp6Ojo1M/F43EVCoWBx9zd3T01qDqNaZra3t7Wzs5Oz5rmv/tisdiqDYVCikQiWl1d1e7urgqFgkKhUMd7r1+/rkqlwn83AAAAAGCOTHUQ1lzK+Ad/8AdaWlrSY489plgsNvA4pmnqrbfeGnd7QF9vlKTPvrRR+GcJ6fGr2vx/VVT5v7d7hkiRSESpVEpra2utk1VH1dyLr1gsKpVKnTn7KRgMKpVKKRKJ9DzEIp/PyzRN7e3tKZ1Oq1gsnlm7ubmpa9euKRgM2v56UqnUqQFYLBZrjdUcr1KpqFKpKJ/P66233mqFY90BV/fXnM/nW89XV1fHtlchAAAAAMA9hmVN78Gaa2trunv3d+vKLMsaauZK+3u//vrrcbUHlwUCAVWrVfn9fh0cHLjdzqm+/Z+lD349QP1j0s//zfFBEblc7sTnY7FYR0AzCZVKRaurq6d+7ujoaKjwLRwOnzqzKpvNKpFIDNRbNBo9EUolEglls1lbY+zs7CiVSp34eDAYVLlcPvU93T8PJ34OAAAAANDt64bNUyOnekd4d+/np/pbc/XqVVmWpWZWN2wIBrjl6PPh6jc2Nk79/FkfH6dgMHhimaB0HP4MOwPtrL4HCcGKxeKJmVk+n0/7+/u2QzDpeAZauVw+8bUcHh6e+Z5xzbwDAAAAgEnjMLbepjoIa948G4bRCsSGfQBuWHp4uPqzlh32Wo44TqddZ5Rrn/beQcKlUqmkaDR64v0ffvjhqaGdnX7u3bvX0QOb4QMAAACYdnYPYyMMO9tU7xF28eJFhUIh3b17V+l0WqFQSMvLywOPU6lU9OMf/1i/+tWvxt8k0MOTjw+2NPKJx4//96yQaJjf/1lXqVS0vr5+4uP5fH6oEKypOZvsrGWgs6y5+b9TwemsKRa
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 51)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_x$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_deltatx_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACLo0lEQVR4nOz9T2yj6X0n+v6oqgQOnFFR6sEAggRMmpoe+OAugqKqZpHdWOTxJr7JxGIVEF87G5eUzg3OBYSx5LLXbVlKrM3BzbRYnsV1YGCqRfskSC4w02L3rO+pErtnZ5xjsTuABAIHcImtxIiRdJN3UREjqSSKkvhH1Pv5AARI6eHv/VGkqvv96nmeN9VsNpsBAAAAADfcyKAbAAAAAIB+EIQBAAAAkAiCMAAAAAASQRAGAAAAQCIIwgAAAABIBEEYAAAAAIkgCAMAAAAgEQRhAAAAACSCIAwAAACARBCEAQAAAJAIgjAAAAAAEkEQBgAAAEAiCMIAAAAASARBGAAAAACJIAgDAAAAIBEEYQAAAAAkgiAMAAAAgEQQhAEAAACQCIIwAAAAABJBEAYAAABAIgjCAAAAAEgEQRgAAAAAiSAIAwAAACARBGEAAAAAJIIgDAAAAIBEEIQBAAAAkAiCMAAAAAASQRAGAAAAQCLcHnQDcBmf//zn45e//GXcunUr/tW/+leDbgcAAADo0P/1f/1f8dlnn8XnPve5+MUvftHXY6eazWazr0eELrh161Y0Go1BtwEAAABc0sjISHz22Wd9PaYZYQylwyBsZGQkJiYmul6/Vqv1tL5jDd9xHGu4jnUTX5NjDc9xHMuxBn0cxxquY93E1+RYw3McxxquY92k13RY/9atW12vfa4mDKHJyclmRDQnJyeHsr5jDd9xHGu4jnUTX5NjDc9xHMuxBn0cxxquY93E1+RYw3McxxquY92k19TP13KSzfIBAAAASISBLo184403Bnn4Y1KpVPwf/8f/Meg2AAAAAOiRgQZhOzs7kUqlojnA/foPj59KpQbWAwAAAAC9l/ilkYMM4QAAAADon4FfNbLZbMbc3FxkMpmBHL9er0exWBzIsQEAAADon4EHYcViMb7xjW8MtIe5ubn40pe+NNAeuJxarRZTU1NtxywuLsbi4mKfOgIAAIDkWl9fj/X19bZjarVan7p51cCDsFwuN+gW4v79+4NugUtqNBqxt7fXdszBwUGfugEAAIBkOzg4OPc8fZAGHoSNj48PugWG2MjISExMTLQdMzo6euG6i4uLcXBwcKnnOlZ/3dSf3009Vr/c1J/fTT1Wv9zUn99NPVY/9et13dT3yudiOI5zk4/VTz4Xw3OsfvK5OG50dDQmJyfbjqnVatFoNC7b3pWkmgPcLf7WrVuxv78/8F+CTz75JMbHx+Ozzz4baB90bmpqKvb29mJycjJ2d3cH3Q6n8B5xGp8LTuNzwWl8LjiNzwWn8bngND4X19sg35+BXjXyOl2x8Tr1AgAAAED3DTQIW11dHfhssIiIO3fuxOrq6qDbAAAAAKCHBhqEffOb3+zr8T7++OMzv9fvXgAAAADor4EGYf22sLAw6BYAAAAAGJBEBWHPnz8fdAsAAAAADMjtQTdwGR9//HHU6/WOx1er1djY2LjQcwAAAAC4WYYiCPv4449jdXU1yuVyVKvVS9VoNpuRSqW63BkAAAAAw+LaB2GPHz+OtbW1iHgZZgHDYXFxMQ4ODq7FlWG5PnwuOI3PBafxueA0PhecxueC0/hccJZU8xqnSz/+8Y+jUChERLRmc12l3VQqFZ999llXemOwpqamYm9vLyYnJ2N3d3fQ7QAAAAAdGuQ5/bWeEbayshIRLwOsZrMZmUwmstlsZDKZiIh47bXXzq3x85//POr1erzzzjvxySef9LRfAAAAAK6vax2EVSqV1kywra2tmJ2dvXStpaWleOONN7rVGgAAAABDZmTQDbSTTqcj4mWIdZUQLCIik8nE66+/3oWuAAAAABhG1zoIO1wCef/+/a7U29jY6EodAAAAAIbPtQ7CDmeBvXjxoqv1AAAAAEieax2Effvb345msxmVSqUr9X7wgx90pQ4AAAAAw+daB2F37tyJ733ve/H06dP427/92yvXszQSAAAAILmudRAW8XKj/Gw2G4VC4Up1Pvroo67NLAMAAABg+NwedAOd2Nrainv37sUbb7wRy8vLMT4+3tHzXrx4EfV6PXZ2duKdd97pcZcAAAAAvdNsRjSa7ceMpCJSqf70M4yGIgh7//33IyJiZ2cnFhYWLlWj2WxGyicBAAAAGFKNZsTt77Qf8+lbEbfEH2e69kHYm2++GcViMSIiUqlUNJvnRJ+nEIABAAAAcK33CHvy5ElsbGy0wq/LhGBXeR4AAAAAN8e1nhF2eJXHw5lguVwu8vl8pNPpC+8T9vbbb8fHH3/cw24BAAAAuM6udRBWqVQilUpFOp2O58+fx+uvv37pWo8ePYrXXnuti90BAAAAMEyu9dLIdDodERGPHz++Ugh2WOvu3btd6AoAAACAYXStZ4TNzMzE+++/H5lMpiv1njx50pU6XB+1Wi2mpqbajllcXIzFxcU+dUQS1Ov1VlAPAADAP1tfX4/19fW2Y2q1Wp+6edW1DsIWFhbivffei2q12pV6ZoTdPI1GI/b29tqOOTg46FM33GSVSiWePn0axWIxcrlcbG5uDrolAACAa+fg4ODc8/RButZB2NzcXNy9ezeePn0a//E//scr1/vJT34Sv/d7v9eFzrguRkZGYmJiou2Y0dHRPnXDoByGVJVKJarVausiGel0OjKZTIyPj0c2m418Ph+5XO6V5xcKhahWq7G9vd36Wr1ej3K5HE+fPo1yuRz1er2PrwgAAGA4jY6OxuTkZNsxtVotGo1Gnzo6LtVsNpsDOXKHqtVq/Jt/82/ixz/+cfyH//AfrlTr/v378ezZsy51xiBNTU3F3t5eTE5Oxu7u7qDbubJisRgLCwsXfl42mz0W3lxWKpW68HPS6XTs7++f+r1CoRClUunCNbe2tk4Nqk5Tr9djZWUl1tbWLnycbDYbuVwupqenY2trK0ql0is/y8NwrFKpvPL8ubk5M8IAAIC++6wRcfs77cd8+lbErWu9I/xgz+mv+Y8mIpPJxNtvvx3f+MY3rlTnk08+OfWEFq6D+fn52N/fj+3t7Zibm2s7NpfLxdbWVmt8NzSbzdjZ2YmNjY22e19lMpnY2NiInZ2dM0OwiIjNzc3Y39/vKNhaWlqK7e3t2N/f7zgEW15ejrGxsVdCsMOA6rDe4eva2tqK+fn51murVCqxtrYWCwsLrcDu5IyvwzrNZrNr+xQCAAAwWNd6RthPfvKT1v2lpaVIpVLnhgSnqdfr8c4770S9Xo/PPvusmy0yIDdtRthJCwsLUSwWX/l6P2YiVavVmJ6ePvV7+/v7l9okfmZm5tQgemNjI+bn5y/UWz6ff2XfwPn5+djY2OioxtraWiwvL7/y9UwmEzs7O6c+5+T7YUYYAAAwCGaEXd213iPsu9/9bnzwwQetx81m81LLoA6fe5nlXzAIZwVhl1k+eVGZTCay2ewrwdXc3Nylr5S4sLBwau8XCcHK5XLk8/ljX0un0/Hee+9FNpvtuM7S0lLMzc3FzMzMsVlgL168OPM5rhAJAABwM1zrjPDBgwfRbDbjcNKaIIukOGspXr+W6J12nKsc+7TnXiRcqlQqXQnBjvbz0UcfHevBZvgAAAA337UOwg5nkKRSqVYgdtkbDJOzQqLx8fH+NnINVKvVmJ2dfeXrm5ublwrBDqXT6a7tsXadlEqlV5aO8s/K5fKgWwAAAAboWi+NvHPnTmSz2fjggw9idXU1stnspYKAarUa3/3ud+PDDz/sfpNwA532e/baa691tV6nv8uFQuGV2Vrz8/Mdb6zfTiaTiaWlpUsvue6nSqUSGxsbUS6Xo1qtRjqdjvHx8Uin05HJZOLhw4cREfHo0aNYXV1tu+y0XC5HoVCI8fHx2NjY6MrP8rqqVqtRLpdjc3OzFYJdpz+OJOm9AACA6+BaB2EREQ8fPozp6en45je/eekad+/eja9
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.5, 0.5, 100)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_yEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_y$\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intercept= 5308.975458144673\n",
      "coef= {'dSlope_xEndT_abs': 84.02382370039953, 'x_EndT_abs': 0.03552632399369171, 'ideal_state_770_tx^2': 68529.10928948055, 'ideal_state_770_tx dSlope_xEndT': 60538.43896560162, 'dSlope_xEndT^2': 13284.852569734045, 'x_EndT_abs^2': -0.0008174639168529617, 'ideal_state_770_ty^2': -2163.1533980289814, 'ideal_state_770_tx^2 x_EndT_abs': -15.687680488172175, 'ideal_state_770_tx dSlope_xEndT x_EndT_abs': -13.288664637206063, 'dSlope_xEndT^2 x_EndT_abs': -3.077894760025427, 'dSlope_xEndT_abs ideal_state_770_ty^2': 3221.0446732918413, 'x_EndT_abs^3': 1.8433768191744582e-07}\n",
      "r2 score= 0.8143147645331705\n",
      "RMSE = 36.0572267645116\n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso, Ridge, ElasticNet, LassoCV\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.metrics import mean_squared_error\n",
    "import numpy as np\n",
    "\n",
    "features = [\n",
    "    \"ideal_state_770_tx\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_xEndT_abs\",\n",
    "    \"x_EndT_abs\",\n",
    "    \"ideal_state_770_ty\",\n",
    "]\n",
    "\n",
    "target_feat = \"z_mag_xEndT\"\n",
    "order = 3\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(sel_array[feat]) for feat in features])\n",
    "target = ak.to_numpy(sel_array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data,\n",
    "                                                    target,\n",
    "                                                    test_size=0.2,\n",
    "                                                    random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=order, include_bias=True)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_tx^2\",\n",
    "#     \"dSlope_xEndT^2\",\n",
    "#     \"dSlope_xEndT_abs\",\n",
    "#     \"x_EndT_abs\",\n",
    "# ]\n",
    "\n",
    "keep = [\n",
    "    # \"ideal_state_770_tx\",\n",
    "    # \"dSlope_xEndT\",\n",
    "    # \"ideal_state_770_ty\",  # no\n",
    "    # \"ideal_state_770_tx x_EndT_abs\",  # maybe no\n",
    "    # \"ideal_state_770_tx ideal_state_770_ty\", # dont keep\n",
    "    # \"dSlope_xEndT_abs x_EndT_abs\",  # dont keep\n",
    "    # \"ideal_state_770_tx^2 dSlope_xEndT_abs\", # dont keep\n",
    "    # \"ideal_state_770_tx dSlope_xEndT dSlope_xEndT_abs\",  # dont keep\n",
    "    # \"dSlope_xEndT^2 dSlope_xEndT_abs\", # no\n",
    "    \"dSlope_xEndT_abs\",  # keep\n",
    "    \"x_EndT_abs\",  # keep\n",
    "    \"ideal_state_770_tx^2\",  # do keep\n",
    "    \"ideal_state_770_tx dSlope_xEndT\",  # do keep\n",
    "    \"dSlope_xEndT^2\",  # keep\n",
    "    \"x_EndT_abs^2\",  # do keep\n",
    "    \"ideal_state_770_ty^2\",  # keep\n",
    "    \"ideal_state_770_tx^2 x_EndT_abs\",  # do keep\n",
    "    \"ideal_state_770_tx dSlope_xEndT x_EndT_abs\",  # do keep\n",
    "    \"dSlope_xEndT^2 x_EndT_abs\",  # do keep\n",
    "    \"dSlope_xEndT_abs ideal_state_770_ty^2\",  # maybe keep\n",
    "    \"x_EndT_abs^3\",  # do keep\n",
    "]\n",
    "\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\n",
    "\n",
    "X_train_model = np.delete(X_train_model, remove, axis=1)\n",
    "X_test_model = np.delete(X_test_model, remove, axis=1)\n",
    "poly_features = np.delete(poly_features, remove)\n",
    "# print(poly_features)\n",
    "\n",
    "lin_reg = LinearRegression(fit_intercept=True)\n",
    "# lin_reg = Lasso(alpha=0.00001)\n",
    "# lin_reg = LassoCV()\n",
    "# lin_reg = ElasticNet()\n",
    "# lin_reg = Ridge(alpha=0)\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features, lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['x_EndT_abs^3']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "koeffs = dict(zip(poly_features, lin_reg.coef_))\n",
    "\n",
    "remove = []\n",
    "for itr in koeffs.items():\n",
    "    if abs(itr[1]) < 1e-4:\n",
    "        remove.append(itr[0])\n",
    "remove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['dSlope_xEndT_abs',\n",
       " 'x_EndT_abs',\n",
       " 'ideal_state_770_tx^2',\n",
       " 'ideal_state_770_tx dSlope_xEndT',\n",
       " 'dSlope_xEndT^2',\n",
       " 'x_EndT_abs^2',\n",
       " 'ideal_state_770_ty^2',\n",
       " 'ideal_state_770_tx^2 x_EndT_abs',\n",
       " 'ideal_state_770_tx dSlope_xEndT x_EndT_abs',\n",
       " 'dSlope_xEndT^2 x_EndT_abs',\n",
       " 'dSlope_xEndT_abs ideal_state_770_ty^2',\n",
       " 'x_EndT_abs^3']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ls_koeffs = []\n",
    "for itr in koeffs.items():\n",
    "    ls_koeffs.append(itr[0])\n",
    "ls_koeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9T2wbeZ7neX8o17OYxbMlheS+2UClgunGPKe2glIDtbeyyMw5zGC3yqSceAY7e8gUaXct9qBKk5b70DWHtkzZ6cugukzKeZnFACORdu5i91Bl0q5jAWWSct4GncmQEkiftlMMsXaxg91OxR7UjCL1h+J/UuL7BRDNPxG/3zdIytX85O+Pz3VdVwAAAAAAAMAFNzHsAgAAAAAAAIBBIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCAMAAAAAAMBYIAgDAAAAAADAWCAIAwAAAAAAwFggCANwITmOM+wSAAAAAAAjhiAMwIVRKpWUSCQ0PT2t5eXlYZcDAAAAABgxPxh2AQDGQ6lU0ubmpkqlkmzb1t7enhzHkWEYMk1TMzMzsixLoVBIwWDw2PmRSES2batYLHrPOY6jfD6vzc1N5fN5RoEBAAAAAJryua7rDrsI4DxIp9OKxWJtn2dZVkN40ymfz9f2OYZhqFKpnPhaJBJRNpttu81cLndiUHUSx3G0tram9fX1tvuxLEvBYFB+v1+5XE7ZbPbYe1kLx0ql0rHzw+GwMplM2/0CAAAAAC4upkYCLYpGo6pUKioWiwqHw02PDQaDyuVy3vG94LquyuWyUqmUDMM49TjTNJVKpVQul08NwSQpk8moUqm0FGzF43EVi0VVKpWWQ7DaFMWjIVgtoKq1V7uuXC6naDTqXVupVNL6+rpisZgX2B0d8VVrx3VdmabZUl0AAAAAgPHFiDCgQ7FYTOl0+tjzgxiJZNu2/H7/ia9VKpWmQdlpAoHAiSOrUqmUotFoW7WFQiHZtt3wfDQaVSqVaqmN9fV1JRKJY8+bpqlyuXziOUc/D0aEAQAAAACOYkQY0KHTpkl2Mn2yXaZpyrKsY8+Hw+GOQjDp9LrbCcHy+bz8fn9DCGYYhorFYsshmHQ4Aq1cLh+7lr29vVPP6fS6AQAAAADjgyAM6NBpU/EGNUXvpH666fukc9sJl0qlkkKh0LHzX716dWJo10o9Ozs7DTWwGD4AAAAAoBsEYUCHTguJZmZmBlvICLBtW4uLi8eez2QyHYVgNbXRZBdRNps9Nn0Uf5LP54ddAgAAAIALiCAMOKdOCtwuX77c0/ZaDfUikcix0VrRaLTlhfWbMU1T8Xi863YGpVQqKRaLye/3y+fzaXp6Wn6/X4FAwNupM5vNanl5+cywJ5/Pe+df9GDItm2l02mFQiH5fL5jowuHbZw+CwAAAOAi+8GwCwBwvuXz+RMX2U8mkz3rY3V11dt90nGckVwPzHEcLS8veztcmqbp7S5q27ZKpZJKpZL3eu2cZmoBo+M4isVip24UcJ7Ztq1IJCLbtkd66us4fBYAAADAOGBEGICunLTIfjeL9p/EMIy2Fu0fhtpoL+lwp81yuaxMJqNMJqNisahyuXxshNxZYco4TLM1TVPFYlGVSmWkR/6Nw2cBAAAAjAOCMAAdy+fzJ65ztbq62vO+IpGIpOY7RzaTz+cViUQapiyGQiGl0+mua1tfX/emy8Xj8RNDO9M0lcvlGkbKnbVGWCqVkmmaMk1TmUym6zpH3ahNh6w3bp8FAAAAcFExNRJAx04LBLpZIP80wWBQyWSy7ZE5juNocXHx2PRNx3GUz+eVz+eVTCa7Wth/bW3Nu3/SCLl68Xhc3333ndbX188M9YLB4FhNwRvWqKtYLKZUKtX0mHH7LAAAAICLihFhADp20qLhvVgg/zTxeLytKZe2bWt2dvbENcyOHhcIBDpaBL2Tta2SyaQsy2LXyBGQTqd7MioQAAAAwPlAEAagI7Ztnxjk9GM0WCccx/EWOA+Hw0qlUioWi8rlcqeuRRUKhdoOp46GYK2GaaurqyO9OPw4sG37zBF8AAAAAC4WpkYC6MhpgZHf7x9wJSerBVK5XO7YKLVgMKhYLHZi8BWJRFQsFlvuxzTNhseJREJLS0tnjlyr7Sg5qrtgXnSO44z0mmQAAAAA+oMRYQA6cloQNkq7650UgtXUFq8/qlQqtTVF0jCMhiDLcRwFAoGWRpa5rksINgS1deOYmgoAAACMH4Iw4AJJJBLy+Xwd3QKBQFt9nTatb1SCHcMwzlyvzDTNhl0ca056rpmjx9u2Lb/f35O1p0qlkmKx2Jkj7RzHUTqdViAQaOi3Nv1venpaPp9Pfr9f6+vrx863bVuJRMLbVdPv9ysWi534Odf6OfodOioUCnnt1d96OSU0nU439DM9Pa1AIHDiNUqH7+dJ68bV13fae93qZ1HjOI7W19cVCoW8979WXyKRaCmIs21b6+vr8vv9DTXXfwZ+v9+bBgwAAADgDC6Ajkk6dqtUKgPpOxqNnth/L2+mabbdfy6XG8j1HxWPxxvqCAaDLZ9rGMax62hXMBg88f2wLMstFotttVUsFt1oNNpQl2EYx46rVCpuMpl0Lctq6DOVSrmue/w9qb9Fo1GvnWbHmaZ56nf66HfgNMlksqW/kWKx2PJnkMlkvPclk8k0tGGa5qm1VyoVt1gsHqupWCw23Orba+WzOCqVSnnHJ5NJt1gsuuVy2c1kMg2fVzgcPnZupVJx4/G4dx31NVYqlWOfd31d5XL5zNoAAACAccaIMOACiUajKhaLHd1SqVRbfZ028mtURqW0MzLtpJFjZ+00eVQul/PW/TraTiAQaHnEjm3b2tzcbHnUkWVZxxZ8LxaLCgQCKpVKymQyKpfLqlQqDZsEpNNplUolhUIh5fN55XI5VSoVVSoV5XI57/2zbVtra2sn9h2JRFqq8aT3pRv5fN7re3V1taF9y7KUyWQkHda+vLzccK5hGLIs69imDrXn6l9r97OoSSQSisViMgxD5XJZ8XhclmXJNE2Fw2EVi0Wv5mw2K7/f3/Dd2NvbUygUOva9LBQKmp2d9ab1lstlpVIp77NyHIfF/wEAAICzDDuJA84zjdiIsGQy2XF7R0fj6IwRYUdH1NRutdFIg3Z0VNNJI21OUxtd1IvrOO196aTd+rbOGoV0dJTQaSPzjo5cqx8ZVq/+PTmt71wu19IIrkql0tMRYeFw+Mzv/FnttFp7TaufRf1IsLM+6/oRX6eNYKyv0TCMEz/Xo9/fQf0bBAAAAJxHjAgD0JFRHxHWjqM7P0qdX0c8HlelUjl1FFRtt8pW2j+prtPUb1IQjUZPXR/t6E6Jp40ErK/fcZyR/VxP+x7WP9+LRfFb/SwSiYTXfzQabelY6XCU20mbNNRfx8bGxomf69HvWqFQaKlWAAAAYBwRhAHoyGnBQLlcHnAl3TspTPnuu++6ai+TyahYLB6bgicdhh6Li4s9DZdanQpa/7m1M310lHZY3NjYUDKZVCaTOTNskgYXzqbTaa+vVoKzaDTa8Bm0u0lDvfr+RumzAgAAAEYNQRiAjpw24uikUS2j7qTQ4vLly123a1nWqeuvlUqlY+tXDUKnu3ru7e31tpAuGIaheDx+6pps7a7v1iu5XM673+oIsvq/o/P4twMAAACcNwRhADp20mgn27ZHdhpdO9qZlniWaDSqcrl8LITKZrMjHX50GpoNmm3bSqfTyufzJy6EP8g62nX0e3YR/nYAAACAUUYQBqBjt27dOvH5UQ53TnM09OllEFZrr1gsHuun3d068Se1XS9TqVTTddE
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(5150, 5700, 50)\n",
    "ax = sns.regplot(\n",
    "    x=y_test,\n",
    "    y=y_test - y_pred_test,\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax.set_ylim(-150, 175)\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    y_test,\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "ax.set_ylabel(\n",
    "    r\"$\\left\\langle z_{\\mathrm{Mag}}-z_{\\mathrm{Mag}}^{\\mathrm{pred}}\\right\\rangle$ [mm]\"\n",
    ")\n",
    "ax2.set_ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_regression_plot.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9b2wb6Z0ven4pdXI7J4lYoufFLuxFrGJ7tudeYGMVpQCZ8yqtYvdZnPROOialYDLdAW5bpJ3kzgJKTEodXJw5WMQSZbeBi9kkIuV+cToT4Eik3clNBtttltzvNntisuR+s+g7bZbUOPbdF2dMFpkEaSRt1r5QWM3/LJLF/98PQLREPfU8vyLpOPr59/weh2EYBoiIiIiIiIiIiMbc1KADICIiIiIiIiIi6gcmwoiIiIiIiIiIaCIwEUZERERERERERBOBiTAiIiIiIiIiIpoITIQREREREREREdFEYCKMiIiIiIiIiIgmAhNhREREREREREQ0EZgIIyIiIiIiIiKiicBEGBERERERERERTQQmwoiIiIiIiIiIaCIwEUZERERERERERBOBiTAiIiIiIiIiIpoITIQREREREREREdFEYCKMiIiIiIiIiIgmAhNhREREREREREQ0EZgIIyIiIiIiIiKiicBEGBERERERERERTQQmwoiIiIiIiIiIaCIwEUZERERERERERBOBiTAiIiIiIiIiIpoITIQREREREREREdFEYCKMiIiIiIiIiIgmAhNhREREREREREQ0EZgIIyIiIiIiIiKiicBEGBERERERERERTQQmwoiIiIiIiIiIaCIwEUZERERERERERBOBiTAiGku6rg86BCIiIiIiIhoyTIQR0dhQVRXhcBizs7NYXV0ddDhEREREREQ0ZJ4adABENBlUVcXe3h5UVYWmachms9B1HYIgQBRFuFwuSJIEr9cLWZZrrvf7/dA0Del02nxO13UoioK9vT0oisIqMCIiIiIiImrKYRiGMeggiEZBLBZDMBhs+zpJkiqSN51yOBxtXyMIAnK5XN2f+f1+JBKJtudMJpN1E1X16LqOzc1NbG9vt72OJEmQZRlutxvJZBKJRKLmtSwlx1RVrbne5/MhHo+3vS4RERERERGNL26NJLIoEAggl8shnU7D5/M1HSvLMpLJpDneDoZhIJPJIBqNQhCEhuNEUUQ0GkUmk2mYBAOAeDyOXC5nKbEVCoWQTqeRy+UsJ8FKWxSrk2ClBFVpvtJ9JZNJBAIB895UVcX29jaCwaCZsKuu+CrNYxgGRFG0FBcRERERERFNLlaEEXUoGAwiFovVPN+PSiRN0+B2u+v+LJfLNU2UNeLxeOpWVkWjUQQCgbZi83q90DSt4vlAIIBoNGppju3tbYTD4ZrnRVFEJpOpe031+8GKMCIiIiIiIqrGijCiDjXaJtnJ9sl2iaIISZJqnvf5fB0lwYDGcbeTBFMUBW63uyIJJggC0um05SQYcFKBlslkau4lm802vKbT+yYiIiIiIqLJwUQYUYcabcXr1xa9eut0s3a9a9tJLqmqCq/XW3P9wcFB3aSdlXiOjo4qYmAzfCIiIiIiIuoGE2FEHWqUJHK5XP0NZAhomoalpaWa5+PxeEdJsJJSNdk4SiQSNdtH6ROKogw6BCIiIiIiGkNMhBGNqHoJt1OnTtk6n9Wknt/vr6nWCgQClhvrNyOKIkKhUNfz9IuqqggGg3C73XA4HJidnYXb7YbH4zFP6kwkElhdXW2Z7FEUxbx+3BNDmqYhFovB6/XC4XDUVBcO2iS9F0RERERE4+ypQQdARKNNUZS6TfYjkYhta2xsbJinT+q6PpT9wHRdx+rqqnnCpSiK5umimqZBVVWoqmr+vHRNM6UEo67rCAaDDQ8KGGWapsHv90PTtKHe+joJ7wURERER0SRgRRgRdaVek/1umvbXIwhCW037B6FU7QWcnLSZyWQQj8cRj8eRTqeRyWRqKuRaJVMmYZutKIpIp9PI5XJDXfk3Ce8FEREREdEkYCKMiDqmKErdPlcbGxu2r+X3+wE0PzmyGUVR4Pf7K7Yser1exGKxrmPb3t42t8uFQqG6STtRFJFMJisq5Vr1CItGoxBFEaIoIh6Pdx3nsBu27ZDlJu29ICIiIiIaV9waSUQda5QQ6KZBfiOyLCMSibRdmaPrOpaWlmq2b+q6DkVRoCgKIpFIV439Nzc3za/rVciVC4VCePz4Mba3t1sm9WRZnqgteIOqugoGg4hGo03HTNp7QUREREQ0rlgRRkQdq9c03I4G+Y2EQqG2tlxqmoa5ubm6Pcyqx3k8no6aoHfS2yoSiUCSJJ4aOQRisZgtVYFERERERDQamAgjoo5omlY3kdOLarBO6LpuNjj3+XyIRqNIp9NIJpMNe1F5vd62k1PVSTCrybSNjY2hbg4/CTRNa1nBR0RERERE44VbI4moI40SRm63u8+R1FdKSCWTyZoqNVmWEQwG6ya+/H4/0um05XVEUaz4PhwOY3l5uWXlWulEyWE9BXPc6bo+1D3JiIiIiIioN1gRRkQdaZQIG6bT9eolwUpKzeurqara1hZJQRAqElm6rsPj8ViqLDMMg0mwASj1jePWVCIiIiKiycNEGNEYCYfDcDgcHT08Hk9bazXa1jcsiR1BEFr2KxNFseIUx5J6zzVTPV7TNLjdblt6T6mqimAw2LLSTtd1xGIxeDyeinVL2/9mZ2fhcDjgdruxvb1dc72maQiHw+apmm63G8FgsO77XFqn+jNUzev1mvOVP+zcEhqLxSrWmZ2dhcfjqXuPwMnrWa9vXHl8jV5rq+9Fia7r2N7ehtfrNV//UnzhcNhSIk7TNGxvb8PtdlfEXP4euN1ucxswERERERG1YBBRxwDUPHK5XF/WDgQCdde38yGKYtvrJ5PJvtx/tVAoVBGHLMuWrxUEoeY+2iXLct3XQ5IkI51OtzVXOp02AoFARVyCINSMy+VyRiQSMSRJqlgzGo0ahlH7mpQ/AoGAOU+zcaIoNvxMV38GGolEIpb+jKTTacvvQTweN1+XeDxeMYcoig1jz+VyRjqdrokpnU5XPMrns/JeVItGo+b4SCRipNNpI5PJGPF4vOL98vl8NdfmcjkjFAqZ91EeYy6Xq3m/y+PKZDItYyMiIiIimmSsCCMaI4FAAOl0uqNHNBpta61GlV/DUpXSTmVavcqxVidNVksmk2bfr+p5PB6P5YodTdOwt7dnuepIkqSahu/pdBoejweqqiIejyOTySCXy1UcEhCLxaCqKrxeLxRFQTKZRC6XQy6XQzKZNF8/TdOwublZd22/328pxnqvSzcURTHX3tjYqJhfkiTE43EAJ7Gvrq5WXCsIAiRJqjnUofRc+c/afS9KwuEwgsEgBEFAJpNBKBSCJEkQRRE+nw/pdNqMOZFIwO12V3w2stksvF5vzecylUphbm7O3NabyWQQjUbN90rXdTb/JyIiIiJqZdCZOKJRhiGrCItEIh3PV12NgxYVYdUVNaVHqRqp36qrmupV2jRSqi6y4z4avS6dzFs+V6sqpOoqoUaVedWVa+WVYeXKX5NGayeTSUsVXLlcztaKMJ/P1/Iz32oeq7GXWH0vyivBWr3X5RVfjSoYy2MUBKHu+1r9+e3X/wYREREREY0iVoQRUUeGvSKsHdUnPwKd30coFEIul2tYBVU6rdLK/PXiaqT8kIJAINCwP1r1SYmNKgHL49d1fWjf10afw/Ln7WiKb/W9CIfD5vqBQMDSWOCkyq3eIQ3l97G7u1v3fa3+rKVSKUuxEhERERFNIibCiKgjjRIDmUymz5F0r14y5fHjx13NF4/HkU6na7bgASdJj6WlJVuTS1a3gpa/b+1sHx2mExZ3d3cRiUQQj8dbJpuA/iVnY7GYuZaVxFkgEKh4D9o9pKFc+XrD9F4REREREQ0bJsKIqCONKo7qVbUMu3pJi1OnTnU9ryRJDfuvqapa07+qHzo91TObzdobSBcEQUAoFGrYk63d/m52SSaT5tdWK8jK/xyN4p8dIiIiIqJRw0QYEXWsXrWTpmlDu42uHe1sS2wlEAggk8nUJKESicRQJz8
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(5150, 5700, 50)\n",
    "ax = sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]) -\n",
    "    ak.to_numpy(sel_array[\"match_zmag\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax.set_ylim(-150, 175)\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "ax.set_ylabel(\n",
    "    r\"$\\left\\langle z_{\\mathrm{Mag}}-z_{\\mathrm{Mag}}^{\\mathrm{old}}\\right\\rangle$ [mm]\"\n",
    ")\n",
    "ax2.set_ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_old_regression_plot.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    label=\"True\",\n",
    ")\n",
    "plt.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    label=\"True\",\n",
    ")\n",
    "plt.xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "# plt.show()\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_old_new.pdf\",\n",
    "    format=\"PDF\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 330,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAADmiklEQVR4nOz9T2yb6X3vf38oTYMpmkqU3J0HyOjmuMhZ1bopB0h3GZEzmwzajEl50CAJcDIi7Tz9nYUTi9Zk0y5imfLEm/N7EpNyDnCmJ8BPIu1pMV1MTcrZ/YJnTN5ydkUzvKUBxruOdEtp0aCNdT8LhRySkij+//t+AUT05+J1fylSmvDj6/peHtd1XQEAAAAAAABDbqzXBQAAAAAAAADdQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAYAAAAAAICRQBAGAAAAAACAkUAQBgAAAAAAgJFAEAZgoDmO0+sSAAAAAAADgiAMwMCxLEuxWExTU1NaXFzsdTkAAAAAgAHxQq8LADBcLMvS+vq6LMuSbdva3d2V4zjyer0yDEPT09MyTVPBYFCBQODY/cPhsGzbVj6fL33NcRxls1mtr68rm82yCgwAAAAA0BSP67pur4sA+kkymVQ0Gm34fqZpVoQ3zfJ4PA3fx+v1am9v78TvhcNhpdPphufMZDInBlUncRxHKysrWl1dbfg6pmkqEAjI5/Mpk8konU4f+1kWwzHLso7dPxQKKZVKNXxdAAAAAMDoYWskUCUSiWhvb0/5fF6hUKjm2EAgoEwmUxrfDq7rqlAoKJFIyOv1njrOMAwlEgkVCoVTQzBJSqVS2tvbqyvYWlpaUj6f197eXt0hWHGLYnUIVgyoivMVH1cmk1EkEik9NsuytLq6qmg0Wgrsqld8FedxXVeGYdRVFwAAAAAA1VgRBpwhGo0qmUwe+3o3ViLZti2fz3fi9/b29moGZafx+/0nrqxKJBKKRCIN1RYMBmXbdsXXI5GIEolEXXOsrq4qFosd+7phGCoUCifep/r5YEUYAAAAAKBerAgDznDaNslmtk82yjAMmaZ57OuhUKipEEw6ve5GQrBsNiufz1cRgnm9XuXz+bpDMOloBVqhUDj2WHZ3d0+9T7OPGwAAAAAAgjDgDKdtxevWFr2TrtPKtU+6byPhkmVZCgaDx+6/ubl5YmhXTz3b29sVNdAMHwAAAADQCQRhwBlOC4mmp6e7W0gfsG1b8/Pzx76eSqWaCsGKiqvJhk06nT62dRSfy2azvS4BAAAAwIghCAP63EmB27lz59o6X72hXjgcPrZaKxKJ1N1YvxbDMLS0tNTyPN1gWZai0ah8Pp88Ho+mpqbk8/nk9/tLp3Sm02ktLi6eGfZks9nS/Yc9GLJtW8lkUsFgUB6P59jKwl4bpecCAAAAGFUv9LoAAIMhm82e2GQ/Ho+37RrLy8ul0ycdx+m7fmCO42hxcbF0uqVhGKWTRW3blmVZsiyr9P3ifWophouO4ygajZ56SMAgs21b4XBYtm339bbXUXguAAAAgFHHijAAdTmpyX4rTftP4vV6G2ra323F1V7S0SmbhUJBqVRKqVRK+XxehULh2Oq4s8KUUdhiaxiG8vm89vb2+nrV3yg8FwAAAMCoIwgDcKZsNntir6vl5eW2XyscDkuqfXJkLdlsVuFwuGLbYjAYVDKZbKmu1dXV0na5paWlEwM7wzCUyWQqVsmd1SMskUjIMAwZhqFUKtVSjYOg37ZDlhu15wIAAAAYRWyNBHCm00KBVhrknyYQCCgejze8OsdxHM3Pzx/bvuk4jrLZrLLZrOLxeNON/VdWVkofn7Q6rtzS0pI+++wzra6unhnoBQKBkdqC16tVV9FoVIlEouaYUXsuAAAAgFHEijAAZzqpcXg7GuSfZmlpqaEtl7Zta2Zm5sQeZtXj/H5/w43Qm+ltFY/HZZomp0b2gWQy2fKKQAAAAADDgSAMQE22bZ8Y5nRiNVgzHMcpNTkPhUJKJBLK5/PKZDKn9qMKBoMNBVTVIVi9Qdry8nJfN4cfBbZtn7mCDwAAAMDoYGskgJpOC4x8Pl+XKzlZMZTKZDLHVqkFAgFFo9ETg69wOKx8Pl/XNQzDqPg8FotpYWHhzFVrxRMl+/EEzFHgOE5f9yQDAAAA0H2sCANQ02lBWD+dsHdSCFZUbGBfzbKsuld2eb3eiiDLcRz5/f66VpW5rksI1gPFnnFsTQUAAABQjiAMGECxWEwej6epm9/vb+hap23t65dwx+v1ntmvzDCMipMci0762mmqx9q2LZ/P15beU5ZlKRqNnrnKznEcJZNJ+f3+iusWt/9NTU3J4/HI5/NpdXX12P1t21YsFiudqOnz+RSNRk98jovXqX79VAsGg6X5ym/t3BKaTCYrrjM1NSW/33/iY5SOfp4n9Ywrr++0n3W9z0WR4zhaXV1VMBgs/fyL9cVisbqCONu2tbq6Kp/PV1Fz+XPg8/lKW4ABAAAAtMAFcCZJx257e3tduXYkEjnx+u28GYbR8PUzmUxXHn+1paWlijoCgUDd9/V6vcceRyMCgcCJPwvTNN18Pt/QXPl83o1EIhU1eb3eY+P29vbceDzumqZZcc1EIuG67vGfR/ktEomU5qk1zjCMU1/P1c//aeLxeF2/H/l8vu6ffyqVKv1cUqlUxRyGYZxa+97enpvP54/VlM/nK27l89XzXFRLJBKl8fF43M3n826hUHBTqVTF8xUKhY7dd29vz11aWio9jvIa9/b2jj3f5XUVCoUzawMAAABwMlaEAQMoEokon883dUskEg1d67SVX/2yMqWRlWknrRw766TJcplMptT3q3oOv99f94od27a1vr5e96oj0zSPNXzP5/Py+/2yLEupVEqFQkF7e3sVBwQkk0lZlqVgMKhsNqtMJqO9vT3t7e0pk8mUfna2bWtlZeXEa4fD4bpqPOnn0opsNlu69vLycsX8pmkqlUpJOqp9cXGx4r5er1emaR470KH4tfLvNfpcFMViMUWjUXm9XhUKBS0tLck0TRmGoVAopHw+X6o5nU7L5/NVvDZ2d3cVDAaPvSZzuZxmZmZKW3oLhYISiUTpuXIch+b/AAAAQCt6ncQBg0B9tiIsHo83PV/1ihydsSKselVN8VZckdRt1SubTlptc5riCqNWH8dpP5Nm5iyf66xVSNWrhE5blVe9cq18ZVi58p/HadfOZDJ1reDa29tr64qwUCh05uv9rHnqrb2o3ueifCXYWc91+Yqv01Yvltfo9XpPfF6rX7vd+vsDAAAADBtWhAGoqd9XhDWi+vRHqbnHsbS0pL29vVNXQRVPqqxn7pNqOk35AQWRSOTU3mjVJyWetgqwvH7Hcfr2OT3tNVj+9XY0xa/3uYjFYqXrRyKRusZKR6vcTjqgofxxrK2tnfi8Vr/WcrlcXbUCAAAAqEQQBqCm08KBQqHQ5Upad1Kg8tlnnzU9VyqVUj6fP7YFTzoKPebn59saLtW7DbT8OWtk62g/nbC4trameDyuVCp1ZtgkdS+YTSaTpWvVE5xFIpGK56CRAxqqlV+vn54rAAAAYJAQhAGo6bRVRyetbOl3JwUX586da2lO0zRP7b1mWdax/lXd0OyJnru7u+0tpAVer1dLS0un9mRrpLdbO2UymdLH9a4gK/8dGsTfGwAAAGCYEIQBONNJK55s2+7brXSNaGRrYi2RSESFQuFYCJVOp/s6/Gg2NOs227aVTCaVzWZPbITfzToaVf0aG4bfGwAAAGBQEYQBONOVK1dO/Ho/BzynqQ5+2hWEFefK5/PHrtHoSZ34XPHUy0Q
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xlim = 1\n",
    "bins = np.linspace(-xlim, xlim, 50)\n",
    "ax = sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    bins=bins,\n",
    "    color=\"#FB8F67\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    range=[-xlim, xlim],\n",
    ")\n",
    "ax.set_xlabel(r\"$\\Delta t_x$\")\n",
    "ax.set_ylabel(r\"$\\left\\langle z_{\\mathrm{Mag}}\\right\\rangle$ [mm]\")\n",
    "ax2.set_ylabel(f\"Number of Tracks (normalised)\", labelpad=15.0)\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "# plt.show()\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_zmag_deltatx_dist.pdf\",\n",
    "    format=\"PDF\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(\n",
    "    sel_array[\"z_mag_corr\"],\n",
    "    bins=100,\n",
    "    # range=[5100, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    density=True,\n",
    ")\n",
    "plt.xlabel(r\"z$_{Mag}$ correction [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.4, 0.4, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(5100, 5700)\n",
    "plt.xlabel(\"dx/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_x\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"x_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"y_EndVelo_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"y_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_x$\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_yEndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_y$\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = [\n",
    "    # # \"ideal_state_770_tx\",\n",
    "    # \"ideal_state_770_ty\",\n",
    "    # \"ideal_state_9410_ty\",\n",
    "    # # \"dSlope_yEndT\",\n",
    "    # # \"dSlope_yEndT_abs\",\n",
    "    # # \"ideal_state_770_x\",\n",
    "    # # \"ideal_state_9410_x\",\n",
    "    # # \"x_EndVelo_abs\",\n",
    "    # \"y_EndVelo_abs\",\n",
    "    # # \"y_EndT_abs\",\n",
    "    # \"dSlope_xEndT\",\n",
    "    # \"dSlope_xEndT_abs\",\n",
    "    # \"x_EndT_abs\",\n",
    "    \"ideal_state_770_tx\",\n",
    "    \"ideal_state_770_ty\",\n",
    "    \"ideal_state_9410_ty\",\n",
    "    \"ideal_state_9410_tx\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_xEndT_abs\",\n",
    "    \"x_EndT_abs\",\n",
    "    \"ideal_state_770_x\",\n",
    "    \"x_EndVelo_abs\",\n",
    "]\n",
    "\n",
    "target_feat = \"z_mag_corr\"\n",
    "order = 4\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(sel_array[feat]) for feat in features])\n",
    "target = ak.to_numpy(sel_array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    data, target, test_size=0.1, random_state=42\n",
    ")\n",
    "\n",
    "poly = PolynomialFeatures(degree=order, include_bias=True)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_tx^2\",\n",
    "#     \"ideal_state_770_x\",\n",
    "#     # \"x_EndVelo_abs\",\n",
    "#     \"dSlope_xEndT^2\",\n",
    "#     \"dSlope_xEndT_abs\",\n",
    "#     \"x_EndT_abs\",\n",
    "# ]\n",
    "# remove = [i for i, f in enumerate(poly_features) if f not in keep]\n",
    "# remove_term = [\n",
    "#     \"ideal_state_770_x^2 ideal_state_9410_x\",\n",
    "# ]\n",
    "# remove = [i for i, f in enumerate(poly_features) if f in remove_term]\n",
    "\n",
    "# X_train_model = np.delete(X_train_model, remove, axis=1)\n",
    "# X_test_model = np.delete(X_test_model, remove, axis=1)\n",
    "# poly_features = np.delete(poly_features, remove)\n",
    "# print(poly_features)\n",
    "\n",
    "# lin_reg = LinearRegression()\n",
    "lin_reg = Lasso(alpha=0.1, max_iter=2000)\n",
    "# lin_reg = LassoCV()\n",
    "# lin_reg = ElasticNet(alpha=1)\n",
    "# lin_reg = Ridge(alpha=1)\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features, lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-200, 200, 50)\n",
    "ax = sns.regplot(\n",
    "    x=y_test,\n",
    "    y=abs(y_test - y_pred_test),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(y_test, bins=30, range=[-200, 200], color=\"#2A9D8F\", alpha=0.8, align=\"left\")\n",
    "ax.set_xlabel(r\"z$_{Mag}$ correction [mm]\")\n",
    "ax.set_ylabel(\"Mean Deviation [mm]\")\n",
    "ax2.set_ylabel(\"Number of Tracks\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6",
   "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.10.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}