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

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "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)\n",
    "                  ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACxsElEQVR4nOz9T2wbd57n/79K3YcFFi1R8mAAQQI6KsZ7+OUwFiX3vS0y3deOSTnndETGcwrUbdHKZWcvtskk+uU0bVJOzmtRdq69JuW+j8Vy+pDLdFhSAAkEdmGpzMYCe4nqexDIFvWHosQiqyg+HwAxovSpz+ddlt0DvvL+fMpwXdcVAAAAAAAAcMUN+V0AAAAAAAAA0AsEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAgEYQAAAAAAABgIBGEAAAAAAAAYCARhAAAAAAAAGAg/97sAwGv/9b/+V/2///f/9LOf/Uz//M//7Hc5AAAAAADgiP/9v/+3fvrpJ/2X//Jf9H//7//t6dqG67puT1cEuuxnP/uZDg4O/C4DAAAAAAC0MDQ0pJ9++qmna9IRhiunHoQNDQ1pfHzc73KaVKvVwNbmpUG4z0G4R2kw7nMQ7lEajPschHuUuM+rZBDuURqM+xyEe5S4z6tkEO5RGoz77Nd7rNf9s5/9rOdr910Qtr29LcuytLe3J8dx9ObNG0nStWvXFAqFZJqmZmdnNTw87HOl8Ms///M/a3d3V+Pj49rZ2fG7nCaTk5OBrc1Lg3Cfg3CP0mDc5yDcozQY9zkI9yhxn1fJINyjNBj3OQj3KHGfV8kg3KM0GPfZr/dYr9uP44z6Igh7/vy5isWi1tbW5DhOW9eYpqloNKpEIqFbt251t0AAAAAAAAAEXqCDsCdPniiTyci2bUnS0ePMDMM49Zr6mEqlItu2lc/nZZqm7t+/r9///vfdLxoAAAAAAACBNOR3Aad5+fKlrl+/rlQqdWoIVn9/2uu0MbZtK5lM6r/9t/+mv/71rz27DwAAAAAAAARH4DrClpeXlc1mT4Ra0WhUkUhE165dk2maCoVCkqSxsTFJ0t7eniTJcRzZtq03b97Itm2VSqXGdsoffvhBkUhE6XRaDx486Nk9AQAAAAAAwH+BCcJqtZrm5uZULpclHQZfsVhM0WhU09PTHc399u1bPX36VPl8XpZlKZPJqFQqqVQqcag+AAAAAADAgAjE1sjvvvtOU1NTKpfLSiaTqlQqevHihe7du9dxCCZJIyMjSiaT2tzc1A8//KBbt25pc3NTs7Oz+vvf/+7BHQAAAAAAACDofA/CXr9+rUgkoqmpKVUqFT1+/FhTU1NdW880TRWLRb148UL/5//8H83MzKhWq3VtPQAAAAAAAASDr1sjt7a2NDc3p3g8rrW1tZ6uHY1GtbW1pZmZGc3Pz+vPf/5zT9cHrrLFxUXVarUrvfV4EO5RGpz7HASD8LschHuUBuc+B8Gg/C4H4T4H4R4HySD8PgfhHqXBuU9cjOEeP5W+h959911Fo1E9fvzYrxLkOI5mZ2d19+5d/eEPf/CtDnhncnJSu7u7mpiY0M7Ojt/lNAlybcCg4t8lEEz82wSCh3+XQPD0679LP+v2rSPs/v37Mk3T1xBMkkKhkDY3N2Wapubm5nTjxg1f6wEAAAAAAEB3+HZG2LVr11QoFPxavkkoFNLGxoZs2/a7FAAAAAAAAHSJbx1h9+7d82vpU01PT3vyhEoAAAAAAAAEk+9PjeyVu3fv+l0CAAAAAAAAfDQwQVg+n9f29rbfZQAAAAAAAMAnAxGEbWxsyHVdPXv2zO9SAAAAAAAA4BNfzgj75JNPtLW11fV19vb25DhO4xD8x48f6w9/+EPX1wUAAAAAAEDw+BKEhUIhFYtFGYbR1XVc1216b9u2vvvuO924caOr6yIYqtWqJicnW45ZXFzU4uJijyo6XK9Wq2l4eLhnawJojX+XQDDxbxMIHv5dAsETxH+XKysrWllZaTmmWq32qJqTDPd4WtQDr1+/1szMTNtBmOu6HYdm9TnS6bQePHjQ0VwItsnJSe3u7rY19r//9/+uf/u3f+tuQQAAAAAADIh/+7d/0//4H/+jrbETExPa2dnpckXNfOkIm56elmma2traOtG1dRrDMNoa1+p66TAMKxQKBGEDYmhoSOPj4y3HBCk1BwAAAACg3w0PD2tiYqLlmGq1qoODgx5V1MyXjjBJSqVSevLkiYrFom7dunXmuM8//1zpdFqRSESZTEazs7MaGRlpex3LsjQ/P69YLKY//elPXpSOgKt3hPmRLAMAAAAAgNb8/NzuS0eYJM3Pz8uyrJYh2NbWlu7fv69UKnXpECsSiejFixd69913ZRiG/v3f//2yJQMAAAAAAKCPDfm18NzcnEZHR1uOyWQyMk2z404u0zS1sLCgXC6nv/zlLx3NBQAAAAAAgP7kWxAmSS9evGj5883NTaVSKU/Wmp2dleu6yuVynswHAAAAAACA/uJrEHYey7IUiUQ8mWtsbEySVCwWPZkPAAAAAAAA/SXQQVgoFNLW1pYnc9UDMMdxPJkPAAAAAAAA/SXQQdjs7KweP37c8Txv375VPp+XdBiuAQAAAAAAYPAEOghLJpMql8v68ssvLz1HrVZTNBqVJBmGIdM0vSoPAAAAAAAAfSTQQVg8HtfU1JSWlpb04Ycfant7+0LXP3/+XFNTU7Isq/G9O3fueFwlAAAAAAAA+oHhuq7rdxGt2Latd999V4ZhSJIikYhSqZRM05RpmnrnnXckSdvb23IcR7Ztq1gsam1tTY7jyHVdGYYh13U1OjqqN2/e+Hg36IXJyUnt7u5qYmJCOzs7fpcDAAAAAACO8PNz+897utolmKapzc1Nzc7OSjp8kmQqlTr3unq+Vw/QDMNQoVDoXqEAAAAAAAAItEBvjayLRCL64YcfND09Ldd123oZhtHoBHNdV2tra7p165bftwIAAAAAAACf9EUQJh12hpXLZa2trSkSiZz4eT34qqsHYPF4XPv7+7p9+3YvywUAAAAAAEDABH5r5HHxeFzxeFxbW1sqlUoql8uybVt7e3uSpLGxMZmmqVgspmg0qpGREZ8rBgAAAAAAQBAE/rB84KI4LB9+chxHoVDI7zIAeOC9b746d8z3H33a9ToAAACuGg7LB4A+ZlmWnj59qnw+r2g0yoM5gIBrJ+ACAADA1XQlg7Dl5WV98skn+uUvf+l3KQB8Vg+pLMtqbKOud22ZpqmxsTFFIpHGdurjEomEbNtWuVxufM9xHJVKJT19+lSlUkmO4/TwjgAAAAAAl9U3h+XXPX/+XHfv3tXNmzf1/PnzU8fMzMxoenpaH374oX788cceVwh4K5/PNx4GcZHXzMyMJ+tfZu3R0dEz50skEpeas1QqtV2z4zhKp9ONP4dsNqtSqSTbthuhleM4sixLpVJJ2WxWsVisMT6dTiufzyuRSGh9ff3E/AsLC3r48KHW19cJwQAAAACgj/RNEPbFF1/o2rVrSiQSyufzsiyrcUD+cfF4XLZt629/+5tM09S3337b42oRBNVqVZOTky1fKysrfpd5rmQyqf39fZXLZcXj8ZZjo9GoisViY7wXXNdVpVJRLpdrefaVaZrK5XKqVCra398/c1yhUND+/r6KxeKpHVhHLS0tqVwua39//9yxdel0WqOjo8pms03fj8fjKhQKjfnq91UsFpVMJhv3ZlmWstmsUqlUIwQ7HnbV53FdV6ZptlUXAAAAAAyClZWVcz+LV6tV3+oL/NbIt2/fanZ2VrZt6+i5/oZ
      "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": 3,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACcKklEQVR4nOz9T4xb953ne39YUgwF8VBkuTEAUQS6RbYHzmpapHQXvesWOdnYSGZMSoMYtjcp0u2n7yw4HTIVr+0KK2lu+j6BSLkXj3INjETG7cBZdKeozO4Cd1RFpXfGJDxygCoQGIxVR+wYMRKbfBbVZMj6Qx6Sh+eQPO8XQKD+fOt3vqfIKul86vf7HV+32+0KAAAAAAAAWHFrbjcAAAAAAAAAOIEgDAAAAAAAAJ5AEAYAAAAAAABPIAgDAAAAAACAJxCEAQAAAAAAwBMIwgAAAAAAAOAJBGEAAAAAAADwBIIwAAAAAAAAeAJBGAAAAAAAADyBIAwAAAAAAACeQBAGAAAAAAAATyAIAwAAAAAAgCcQhAEAAAAAAMATCMIAAAAAAADgCQRhAAAAAAAA8ASCMAAAAAAAAHgCQRgAAAAAAAA8gSAMAAAAAAAAnkAQBgAAAAAAAE8gCAMAAAAAAIAnEIQBAAAAAADAEwjCAAAAAAAA4AkEYQAAAAAAAPAEgjAAAAAAAAB4AkEYAAAAAAAAPIEgDAAAAAAAAJ5AEAYAAAAAAABPIAgDAAAAAACAJ1x0uwFgWl/5ylf02Wef6cKFC/q3//bfut0OAAAAAACw4H/9r/+lL774QpcuXdKnn37q6LF93W636+gRAZtcuHBBnU7H7TYAAAAAAMAU1tbW9MUXXzh6TGaEYWn1grC1tTWFQqH+x1ut1pkfnwenjsVxFv9Yq3YcJ4+1asdx8lgcZ/GPtWrHcfJYHGfxj7Vqx3HyWKt2HCePxXEW/1irdhwnj+Wl4/Q+d+HChbkd/1xdYEltbGx0JXU3NjYsfdzJHjjOYhzHyWOt2nGcPNaqHcfJY3GcxT/Wqh3HyWNxnMU/1qodx8ljrdpxnDwWx1n8Y63acZw8lpeO4+Tzd5KrM8Kef/55Nw8/xOfz6X/+z//pdhsAAAAAAACYE1eDsGazKZ/Pp66L25T1ju/z+VzrAQAAAAAAAPO35nYDbnMzhAMAAAAAAIBzXN8sv9vtKpVKKRKJuHJ80zRVqVRcOTYAAAAAAACc43oQVqlU9K1vfcvVHlKplL72ta+52gMAAAAAAADmy/WlkYlEwu0WdP36dbdbAAAAAAAAwJy5HoStr6+73QIAAAAAAAA8wNWlkdypEfOQy+XUbrfl9/tX5lirdhwnrdr3bhVf305Zxe/dqh3HSav2vVvF17dTeI4W/zhOWrXv3Sq+7pzCc7T4x3HSqn3vVu04k/J1Xbxt4tramkzTdP2b8vTpUwWDQXU6HVf7wGTC4bAODw+1sbGhg4MDt9vBBHjuMIjXA3p4LWAQrwf08FrAIF4P6OG1sNzcfP5cnRFWLBZdD8Ek6fLlyyoWi263gSm1Wi2Fw+GRNblcTrlczqGOAAAAAADwplKppFKpNLKm1Wo51M1prgZh3/72tx093scff6w/+ZM/WYheYJ9Op6PDw8ORNe1226FuAAAAAADwrna7PfYa3U2uBmFOy2az+qd/+ie324DN1tbWFAqFRtYswsxDAAAAAABWnd/v18bGxsiaVqvl2vZUngrC9vb23G4BcxAKhVgTDgAAAADAArCyNVFvjzA3LGUQ9vHHH8s0Tcv1hmGoXC5P9DUAAAAAAABYLUsRhH388ccqFouq1+syDGOqMbrdrnw+n82dAQAAAAAAYFksfBC2tbWlnZ0dScdhFgAAAAAAADCNhQ7CfvzjH6tYLEqSfD6ffD4fYRiwAnK5nNrtNjcxgCReD/gDXgsYxOsBPbwWMIjXA3p4LWBavu4CJ0vXrl1To9HoB2CRSESxWEyRSESS9Nxzz40d45NPPpFpmrp//76ePn2qL774Yt5twyG9zfU2NjbYLB8AAAAAgCXh5vX8Qs8I64VgkrS7u6sbN25MPVY+n9fzzz9vV2sAAAAAAABYMmtuNzBKIBCQdBxizRKCSVIkEtGVK1ds6AoAAAAAAADLaKGDsN4SyOvXr9syXrlctmUcAAAAAAAALJ+FDsJ6s8CePHli63gAAAAAAADwnoUOwr773e+q2+2q0WjYMt67775ryzgAAAAAAABYPgsdhF2+fFnf+973dO/ePf3Lv/zLzOOxNBIAAAAAAMC7FjoIk443yo/FYkqn0zON8/jxY9tmlgEAAAAAAGD5XHS7ASt2d3d17do1Pf/88yoUClpfX7f0dU+ePJFpmmo2m7p///6cuwQAAAAAAMAiW4og7Oc//7kkqdlsKpvNTjVGt9uVz+ezsy0AAAAAAAAskYUPwv7qr/5KlUpFkuTz+dTtdicegwAMAACsim5X6oz579CaT+K/P4uB5wsAgMWy0EHYnTt3+hvcTxuCSZr66wAAABZNpytdfGt0zedvSxcIVhYCzxcAAItloYOwkyFYIpFQMplUIBCYeJ+w27dv6+OPP55jt3BLq9VSOBweWZPL5ZTL5RzqCAAAAAAAbyqVSiqVSiNrWq2WQ92cttBBWKPRkM/nUyAQ0N7enq5cuTL1WJubm3ruueds7A6LotPp6PDwcGRNu912qBsAAAAAALyr3W6PvUZ300IHYYFAQE+fPtXW1tZMIVhvrKtXr9rUGRbJ2tqaQqHQyBq/3+9QNwAAAJN575H0WtztLgAAsIff79fGxsbImlarpU6n41BHwxY6CIvH4/r5z3+uSCRiy3h37tyxZRwsllAopIODA7fbAABgYRCsLI73Ho2veb0mXViTXuFvtgCAFWBla6JwOOzarLE1V45qUTabVbfblWEYtozHjDAAdjBN0+0WAHiY1WDFSh3m6+ln0ps/sVb7xgdS+7PZjtftSl90Rj+4hxQAwOsWOghLpVK6evWq7t27Z8t477//vi3jAPCeRqOhQqGgYDCozc1Nt9sB4FFOByuYzY8a0qe/t1b7m99JdxuzHa93h8pRjw5BGADA4xY6CJOkarWq/f19/cM//MPMY21vb9vQEYBF0wupksmkotGogsGgfD6fgsGg4vG4ksmkCoWC6vX6mV+fTqcVjw+vITJNU7VaTel0uj/Ozs4Os8EAuMrpYAWz+fCjyep/OmE9AACY3MIHYZFIRLdv39a3vvWtmcZ5+vSpGo3F+d+gYRhKJpOWa7PZrKLR6NDFfTabtW3ZqJMmOXevqFQq8vl8Ez9OhjfTmubYwWDw3PHS6fRUY54XVJ3FNE0VCoX+92FnZ0f1el2GYfTDKtM01Wg0VK/XtbOzo2Qy2a8vFAqqVCpKp9Oq1Wqnxt/c3NT29rZqtRrhF4CFQbCyXI5+O996AAAwuYXeLL+3lPGP/uiPFAwG9fzzzyuVSk08jmmaun//vt3tnXusUQHBoEwmM7ZmZ2dHhULh1DEajYYajYYqlYqKxaLy+fzYsYLB4NQX9JlMRuVyeWSN3efuJZlMRjdv3pRhGP3w5TyJREKFQkHXrl1TIBCw5fi9vfjq9boKhcK5r5NIJKJCoaBEIjHyJhbValWmaWpvb0/FYnFkwJXP53Xr1i1FIhHL51MoFLSzs3Pq46lUqj9WbzzDMGQYhqrVqu7fvz/08zPo5DlXq9X+29FodClDZwCrh2BluQS/PN96AAAwuYUOwt555x09evSHnV673e6ZF79WdLtd+Xw+u1o7V6VSsVx7MuA6KZlMql6vKxAI9IMHwzDUaDSGLsoLhYIikcjIkHDWWS1WZnDZee5eFAgEFIvFVK1Wlc1mz/x+plKpoYDGTpFIRJlMRolEQtFo9Mya/f19y2FV73WbSCQUj8fPnJFZLpcnCkV7swlPhlKjgtpeKJZIJFQul88Ml8dJJBITvb4BYF4IVpbLSy9IP/ul9foXX5hfLwAA4NhCL428efOmut2uuv96exsngqxZWd2HbNyMmt5+RsViUUdHR6pWqyoWi6pWq2o2myoWi0P16XR65PHGzeYax8pMPLvOHcd3TJ3k43a
      "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": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACmxElEQVR4nOz9T2wj6X3v+3+ongQTOFEXNUEAQQLsLmYCZ5VfF9V3kd2xSHiTQf4M2Q3EsL24FunxLzgLnWPS7aw9Mmlbm+AmQ8pZXAcGTjdpnwTJwmdYPVndxT0t1kx2xklYPQEkEDjASNVyggwSj/hb6EealPinSBb/qd4vQIAoPnzqK1E9I374PN8n0mq1WgIAAAAAAABuubVFFwAAAAAAAADMA0EYAAAAAAAAQoEgDAAAAAAAAKFAEAYAAAAAAIBQIAgDAAAAAABAKBCEAQAAAAAAIBQIwgAAAAAAABAKBGEAAAAAAAAIBYIwAAAAAAAAhAJBGAAAAAAAAEKBIAwAAAAAAAChQBAGAAAAAACAUCAIAwAAAAAAQCgQhAEAAAAAACAUCMIAAAAAAAAQCgRhAAAAAAAACAWCMAAAAAAAAIQCQRgAAAAAAABCgSAMAAAAAAAAoUAQBgAAAAAAgFAgCAMAAAAAAEAoEIQBAAAAAAAgFAjCAAAAAAAAEAoEYQAAAAAAAAgFgjAAAAAAAACEAkEYAAAAAAAAQoEgDAAAAAAAAKFAEAYAAAAAAIBQeGXRBQCT+NSnPqWPP/5Yd+7c0W/8xm8suhwAAAAAAODT//7f/1uffPKJXn31Vf3rv/7rXK8dabVarbleEQjAnTt3dHl5uegyAAAAAADAhNbW1vTJJ5/M9ZqsCMNKGicI+7Vf+zWtr69Pdb1ms6nLy0utra1pc3NzqrmCQk2rV49ETX5Rkz/LVtOy1SNRk1/U5M+y1bRs9UjU5Bc1rV49EjX5RU3+LFtNQdZzcXGhn/3sZ77G3rlzZ6prTYIgDCvpN37jN3R6eurrH+n+/r729/enut729rZOT0+1ubmpk5OTqeYKCjWtXj0SNflFTf4sW03LVo9ETX5Rkz/LVtOy1SNRk1/UtHr1SNTkFzX5s2w1BVnP4eGhDg8Ph45pB2+LaHVEEIaVtiz/0QAAAAAAAP4Wo7SDt0VYaBD2+uuvL/LyPSKRiP7X//pfiy4DAAAAAAAAM7LQIKzRaCgSiWiR/frb149EIgurAQAAAAAAALO3tugCFo1DMwEAAAAAAMJh4T3CWq2WUqmUTNNcyPU9z1O5XF7ItQEAAAAAADA/Cw/CyuWyvvKVryy0hlQqpc9//vMLrQEAAAAAAACztfCtkYlEYtEl6MGDB4suAQAAAAAAADO28BVhGxsbiy4BGGl/f18XFxdaX19fdCkd1LSalvFnRE3+LGNNy2YZf0bU5M8y1rSMlu3ntGz1SNS0ypbt57Rs9UjUtMqW8ee0bDUtWz2zFGktsFv8nTt3dH5+vvAf9MuXL7WxsaFPPvlkoXXAv+3tbZ2enmpra0snJyeLLgdLit8T3Hb8juM24/cbtxm/37jN+P2GH4v8PVno1shlOrFxmWoBAAAAAABA8BYahBUKhYWvBpOku3fvqlAoLLoMAAAAAAAAzNBCg7Cvf/3rc73ehx9+OPC+edcCAAAAAACA+Vr4qZHzlM1mF10CAAAAAAAAFiRUQdjx8fGiSwAAAAAAAMCCvLLoAibx4YcfyvM83+Nd11WpVBrrMbPkuq6y2axqtdpU88RiMbmuq0qlolQq5ftxjuOoVCrJtm25ritJMk1TqVRKjx8/lmEYE9VTLBZVq9V0fHwsz/NkmqYsy1I2m1UikZhoTgAAAAAAgKCsRBD24YcfqlAo9AQ342q1WopEIgFX1svzPEWjUV9jM5nMVNfK5/Nj/yw8z9Pe3p6q1eqN+1zXVbFYVLFYVKlUGqs+27aVTqfleZ4SiYQqlYpM05TjOMrn80omk52vTxqyAQAAAAAATGvpg7DHjx+rWCxKugqzllm5XPY9Np/PT3wd27Y7PxO/PM9TPB73FZ5ls1nV63WVSiVftSSTSUlX4V73Y9qrzOLxuGzbVjweV71eJwzD3Ozv7+vi4mIpTqcFZoHfcdxm/H7jNuP3G7cZv99YdpHWEqdLP/rRj5ROpyWps5prmnIjkYg++eSTQGrrJxqN+tp+mUgkJt4W6Xme7t2713MdP1sjk8mkbNuWZVl6/PixLMuSdLVN8vnz532DtVHzdtdimqYajUbfca7rKhaLSZrue++2vb2t09NTbW1t6eTkZOr5AAAAAADAfCzyNf1Srwg7ODiQdBVgtVqtTs8p0zQlSa+99trIOT766CN5nqenT5/q5cuXM6u1XC7L8zzlcrnOCqlBdnZ2Jr7O3t6eNjY2JMl3z7NyuSzbtpXL5VQoFHrua6/aymazSqfTchyn51rDgrD2dkhp+Aq39jWq1aps21a5XJ56aygAAAAAAMC4lnpF2NraWmcl2Lvvvqvd3d2J53JdV6+//vrMVoS1VzwNWhUVhHK53Nm2uLu72wmhRq3cisViMk1z5Eqs7pVbbbVarW+j++tjz8/Ph255rFarndV9hmHo/Px8aC2jsCIMAAAAAIDVtMjX9GtzvdqY2sFKLpebKgSTrlYl3bt3L4CqbqpWq3Jdd6q+X6O0T5rM5XKdbY1+OI7TOVlyFNM0b6wY614h1q17XCKRGNn3qzuo8zyvb8N+AAAAAACAWVrqIKy9BfLBgweBzOen+fskDg4OZBiGHj58OJP5pattiJZl3QiqRnny5IkymYzvBvXXV3999NFHfcd1HwzgN5hrP5/tugAAAAAAAOZpqXuE7e7u6v3339fZ2Vlg8wXNcZzOqqloNCrTNJVIJJRMJkc2sPcrn8/LcZyJtl0+evSoJ4Aa5XqodX2rpHRzlZjfoNKyrM6plawIAwAAAAAA87bUK8K++c1vqtVqDdyeN67vf//7gczT7fp2SNd1VS6XlU6nFYlEbjSgH5fjOCoWiyqVSmMFWm2WZfleDSbdbMDf75q2bY8c08/1cUE9rwAAAAAAAH4sdRB29+5dffvb39aTJ0/0s5/9bOr5gt4a6brujVDoumq1qng8rmw2O9E1dnd3lUql5nbKYnvFVlu/RvnPnz/vue03aLt+yufx8fF4xQEAAAAAAExhqYMwSZ3m8O0TByf14sWLwFcgmaapUqmkQqGgTCYzdGVUuVxWPB4fa/7293x0dDRVnePoDqcGhW/Xw7JJV4TN8oRNAAAAAACA65a6R1hbrVbTzs6OXn/9deXzeW1sbPh63NnZmTzPU6PR0NOnT2dS2/WwyPM8lctlHRwc3Nhm6DiOksmkarXayHmr1aqq1apqtdpYWxun1b1qbtApmNeDsEld//kAAAAAAADM0koEYe+9956kqxVEk24xbLVaikQiQZbVl2EYyuVyyuVyqlar2tvb6wl8bNtWsVhULpcbOIfneUqn08pkMn23Js6K67qdVXOFQmHgSq9JA6zrgV4QhyA0m01tb29PPc/+/r729/ennue6Vku6bA0fsxaR5vCrCQAAAADAVA4PD3V4eDj1PM1mM4BqJrP0Qdhbb72lcrksSYpEImq1RqQKfcwjAOsnlUopkUhod3e3Z1vmwcHB0CBsd3e3s+1yngqFgqSrLYzD6gtKECvCLi8vdXp6OvU8FxcXU8/Rz2VLeuVPh4/5+bekOwRhAAAAAIAld3FxEchr8EVa6iDs6OioEwZNGoJJmvhxQTAMQ/V6XfF4vBOGeZ4n27b7rvYqFotyHEf1en2udTqOo3K5LMMwRm7dNAwjkBAriC2fa2tr2tzcnHqe9fX1qecAAAAAAOA2W19f19bW1tTzNJtNXV5eBlDR+JY6CLsegiUSCSWTSRmGMXafsHfeeUcffvjhDKsd7ujoqKdZfq1WuxGEOY6jfD6vQqEgy7LmWt/e3p4k6dmzZyOb329sbAQShPl9DofZ3NzUycnJ1PMAAAAAAIDhgmortL29vbCVZUsdhDmOo0gkIsMwdHx8rHv37k08197enl577bUAqxuPZVlKJBKybVtS/4bz6XRalmX
      "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": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACphElEQVR4nOz9TYxid57n+38gs6y0ykNCuNVSCKTpBHnkXnUnpO+idt0JXRtbdo8hc2TLtnSnElz+150FMwWO8tqOIrqaTd9rJaTrSp1uL9JQ7irZixoHWbO6V/pPBmT1zuoqTrqlCCGNpjJO0raccjk5dxEDBfEABzhwIHi/JKR4+MbvfE8E+cAnfg8ey7IsAQAAAAAAAKec1+0GAAAAAAAAgHkgCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEogCAMAAAAAAMBKIAgDAAAAAADASiAIAwAAAAAAwEo463YDwKS+/e1v6+HDhzpz5oz++I//2O12AAAAAACADf/jf/wPPXr0SOfOndOXX34512t7LMuy5npFwCFnzpxRp9Nxuw0AAAAAADABr9erR48ezfWazAjD0uoGYV6vV+vr6zO9VqvVmtu17KIne+hp+fqR6Mkuelq+fiR6souelq8fiZ7soid7Fq2nRetHoie76Gnx+ule78yZMzO/1mEEYVhaf/zHf6y9vT1btdlsVtlsduJrhUIh7e3taX19Xbu7uxOP4yR6soeelq8fiZ7soqfl60eiJ7voafn6kejJLnqyZ9F6WrR+JHqyi57m20+xWFSxWLRV68Y2R64GYU899ZSblx/g8Xj0z//8z263gQl0Op2RgVi73Z5TNwAAAAAArK52u2170oobXA3Cms2mPB6P3NymrHt9j8fjWg+Yjp2pmz6fb07dAAAAAACwunw+n4LB4NCa7tJIN6z80kjOClh+izKVFAAAAACAVWdna6LuUkw3uB6EWZalZDKpcDjsyvVN01S5XHbl2gAAAAAAAJgf14Owcrms733ve672kEwm9d3vftfVHgAAAAAAADBbXrcbiMfjbregZ555xu0WAAAAAAAAMGOuB2Fra2tutwAAAAAAAIAV4OrSSE5qxLLIZrNqt9sLdfokPS2vRfs+LVo/Ej0ts0X7Pi1aPxI9LbNF+z4tWj8SPS2zRfw+LVpPi9aPtJg9LaJF/D4tWk+L1s8seSwXj030er0yTdP1b/SDBw8UCARcO7oTk+meMhEMBjk1EifieYLTjOc3TjOe3zjNeH7jNOP5DTvcfJ64ujSyUCi4HoJJ0vnz51UoFNxuAwAAAAAAADPkahD2wx/+cK7X+/zzz0/83Lx7AQAAAAAAwHy5vln+PGUyGbdbAAAAAAAAgEtWKgjb2dlxuwUAAAAAAAC4xNVTIyf1+eefyzRN2/WGYahUKo31NQAAAAAAADhdliII+/zzz1UoFFSr1WQYxkRjWJYlj8fjcGcAAAAAAABYFgsfhG1sbGhra0vSQZgFAAAAAAAATGKhg7Cf/exnKhQKkiSPxyOPx0MYBmAs2WxW7XZbPp/P7VYAx/H8xmnG8xunGc9vnGY8v7HoPNYCJ0uXLl1So9HoBWDhcFjRaFThcFiS9OSTT44c43e/+51M09SHH36oBw8e6NGjR7NuG3MSCoW0t7enYDCo3d1dt9sBAAAAAAA2uPl6fqGDMK/X29vX69NPP9Xly5cnHsswDD311FMEYadI9w+O1+vV+vr60NpsNqtsNjunzgAAAAAAWE3FYlHFYnFoTavVUqfTcSUIW+ilkX6/Xw8ePFAul5sqBJOkcDisCxcuONQZFkmn09He3t7Qmna7PaduAAAAAABYXe12e+RrdDctdBAWDod19+5dPfPMM46MVyqVHBkHi8XOjDDWpwMAAAAAMHs+n0/BYHBoTXdGmBsWOgi7fPmy7t69q/v37zs2Hk6f9fV19ggDAAAAAGAB2NmaqLvVkRu8rlzVph/96EeyLEuNRsOR8d577z1HxgEAAAAAAMDyWegg7Pz58/rxj3+sW7du6V//9V+nHo+lkQAAAAAAAKtroYMwScrlcopGo0qlUlONc+/ePcdmlgEAAAAAAGD5LPQeYV3b29u6dOmSnnrqKeXzea2trdn6uvv378s0TTWbTX344Ycz7hIAAAAAAACLbCmCsF/96leSpGazqUwmM9EYlmXJ4/E42RYAAAAAAACWyMIHYd///vdVLpclSR6PR5ZljT0GARgAAMDysyypM+K/gl6PxH/9AADASRY6CLtx40Zvg/tJQzBJE38dAAAApuNkeNWxpLNvDa/55m3pDEEYAAA4wUIHYYdDsHg8rkQiIb/fP/Y+YdevX9fnn38+w24BAABwGOEVAABYJAsdhDUaDXk8Hvn9fu3s7OjChQsTj3Xt2jU9+eSTDnYHAAAAAACAZeJ1u4Fh/H6/JGljY2OqEKw71sWLFx3oCgAAAIvqg7tudwAAABbZQgdhsVhMkhQOhx0Z78aNG46MAwAAgPmzE3K9ViUMAwAAJ1voICyTyciyLBmG4ch4zAgD4ATTNN1uAQBOFTvB1YOH0hu/sDfe6z+X2g+nagkAAJxSCx2EJZNJXbx4Ubdu3XJkvI8++siRcQCsnkajoXw+r0AgoGvXrrndDgAsDadmcb3fkL78vb1rfvG1dLNhrxYAAKyWhQ7CJKlSqaher+sf//Efpx5rc3PTgY4ALJpuSJVIJBSJRBQIBOTxeBQIBBSLxZRIJJTP51Wr1Y79+lQq1VuK3WWapqrVqlKpVG+cra0tZoMBwBicnMX18WfjXfuTMesBAMBqWOhTI6WD/cGuX7+u733ve/rrv/7ricd58OCBGg1+NYjFVC6Xlclkxv66aDSqer0+9fU9nvHPrPf7/drf3z/2c6lUStVqdewxt7e3FY/HbdWapqnNzU1tbW0Nren+ua/Var3aaDSqeDyuSCSi7e1tVatVRaPRga+9du2aDMPg7w0AmMIks7h+8J3jP7//1XjXHrceAACshoUOwrpLGf/oj/5IgUBATz31lJLJ5NjjmKapDz/80On2sCBarZZCodDQmmw2q2w2O6eOxpdOp3XlyhUZhqHNzc2hIVI8Hlc+n9elS5d6J6tOq7sXX61WUz6fP3HWUzgcVj6fVzweH3qIRaVSkWma2tnZUaFQOHEmliTlcjldvXpV4XDY9v3k8/ljA7BkMtkbqzueYRgyDEOVSkUffvhhLxw7HHAdvudKpdJ7OxKJOLZXIQCskklmcZ0UhAUeH2+scesBAIAzisWiisXi0JpWqzWnbo5a6CDsnXfe0d27f9gwwrKsobM/hrEsa6JZL1h8nU5He3t7Q2va7facupmc3+9XNBpVpVJRJpNRuVw+UpNMJgcCGieFw2Gl0+neTKnj1Ot122GV3+9XPB5XPB5XLBY7dmZVqVRSOp223aNhGEokEkdCqXQ6rVKpdOzXdEOxeDyuUqmkra0t5fN529eUDsLH434eAIDhnJzF9dzT0qe/sT/Ws0+Pd20AAOCMdrs98jW6mxZ6j7ArV67IsixZliVpsuVbOP28Xq+CweDQh8/nc7vNsZy0THKS5ZPjCofDR5YJSgch3KQz0E7qe5wQrFarHZmZ5ff7Va/XTwzBjpPL5dRsNo/cy/3790/8Gqdm3gHAqnFyFtcrUenb37I3zhOPSa8e/acMAADMgc/nG/ka3et1L45a6BlhmUxGb775pjweTy8MAw5bX1/X7u6u22046qRlh8OWIzp9/cMzuKa59nFfO0641Gg0lEgkjnz97du3jw3t7PRz7949Xbhwobckkk3
      "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": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACcfUlEQVR4nOz9TWxbd57n+38O5QQOkqZIpzGAIAHToToN16oh0p5FLUskahMj6QppAzFsA4MSmeRfdxbCFBlXFr1yVFRVtJl70xZdPQunDbRNVk2CZJGx6PTuAv+xRFfvgqnWcQqQQGCAMilVghiJzXMXajKiHshD8hwePrxfAAHJ+vJ3vhIp2/zw92BYlmUJAAAAAAAAGHE+rxsAAAAAAAAA+oEgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY4EgDAAAAAAAAGOBIAwAAAAAAABjgSAMAAAAAAAAY+GE1w0A3Xj++ef1+PFjTUxM6D/8h//gdTsAAAAAAMCm//N//o+ePn2qkydP6uuvv+7rtQ3Lsqy+XhFwwMTEhGq1mtdtAAAAAACALvl8Pj19+rSv12RGGIbS/iDM52u9wvcv/uIv9MILL/R0vXK5rFqtJp/Pp6mpqZ7Gcgo9DV8/Ej3ZRU/2DFpPg9aPRE920ZM9g9bToPUj0ZNd9DR8/Uj0ZBc92TNoPTnZz1dffaU///nPLWvqr+cnJiZ6ulZXLGAITU9PW5Js3f7+7//esetNT0/33rxD6Km9QevHsujJLnqyZ9B6GrR+LIue7KInewatp0Hrx7LoyS56am/Q+rEserKLnuwZtJ6c7Ofv//7vbb9e9+L7Z0YYhpqdtNrv9/epGwAAAAAAxpvf79f09HTLmvoMNC94GoS9/PLLXl6+iWEY+t//+3973QY6NDU1pa2tLa/bAAAAAAAAkhYXF7W4uNiyZmZmRtvb233qqJmnQdjm5qYMw5Dl4X799esbhuFZDwAAAAAAAHBf613Gx4CXIRwAAAAAAAD6x/M9wizLUjweVygU8uT61WpVuVzOk2sDAAAAAACgfzwPwnK5nH7605962kM8HtePf/xjT3sAAAAAAACAuzxfGhmNRr1uQWfPnvW6BQAAAAAAALjM8xlhp06d8roFoK3FxUXt7u7K7/d73UoDPQ2nQfwZ0ZM9g9jToBnEnxE92TOIPQ2iQfs5DVo/Ej0Ns0H7OQ1aPxI9DbNB/DkNWk+D1o+bDMvD3eInJiZUqVQ8/0Hv7Ozo1KlTevr0qad9wL76UavT09Pa2tryuh0MKJ4nGHU8xzHKeH5jlPH8xijj+Q07vHyeeLo0cpBObBykXgAAAAAAAOA8T4OwbDbr+WwwSZqcnFQ2m/W6DQAAAAAAALjI0yDs5z//eV+v9+WXXx77tX73AgAAAAAAgP7y/NTIfkqlUl63AAAAAAAAAI+MVRC2vr7udQsAAAAAAADwyAmvG+jGl19+qWq1arveNE2trq52dB8AAAAAAACMlqEIwr788ktls1kVi0WZptnVGJZlyTAMhzsDAAAAAADAsBj4IOzq1ataXl6WtBdmAYBdi4uL2t3dHYjTaQE38BzHKOP5jVHG8xujjOc3Bp1hDXC69Nvf/laJREKSGrO5emnXMAw9ffrUkd7grZmZGW1vb2t6elpbW1tetwMAAAAAAGzy8jX9QM8IW1pakrQXYFmWpVAopHA4rFAoJEl68cUX247xpz/9SdVqVXfu3NHOzo6r/dplmqZSqZTW1tZ6Gmd2dlamaSqfzysej3d03+XlZa2trWl9fV3VarXxs02lUopGo13148aYAAAAAAAAThnoIKxUKjVmgq2trWl+fr7rsdLptF5++WWnWjtStVpVMBi0VZtMJnu6ViaT6Wq/tGKxqEQioWq1qmg0qnw+r1AopFKppEwmo1gs1vjzQCDg2ZgAAAAAAABOG+ggLBAIaGdnR+l0uqcQTJJCoZBeeuklhzo7Wi6Xs12byWS6vk6xWGzsm9bp/WKxmKS9IG51dbXxtVAopHg8rkgkomKxqEgkoo2NjbbBlRtjAgAAAAAAuMHndQOt1JdAnj171pHx9oc0bqgv5WwnGo02vrdOVavVxr5p3d4vFAod+7PI5/OS9pZvtruOG2MCAAAAAAC4ZaBnhM3Pz+vBgwd69OiRY+O5JZfLqVqtKp1ON2ZIHefMmTNdX2dhYUGnTp2StBdE2VVfuii1no1Wn8VVKBRULBaVy+WOXcbpxpgAAAAAAABuGehTI3d2dhQMBvXmm2/qgw8+6Hm83/zmN/rpT3/qQGeHzc7OSpI2NzddGV/aC9tSqZQ2NjY0Pz/fCKHabZZvmmajP0mqVCotlycWCoXGzK1AIKBKpdKXMTvBqZEAAAAAAAwnL1/TD/TSyMnJSf3yl7/U7du39ec//7nn8dxaGlkoFGSaZk/7frVTP2kynU4rHA53dN9sNtv4OBqNtt2ja3+oVq1WVSgU+jImAAAAAACAmwY6CJPUCH563Vvq4cOHKpVKDnXVbGlpSYFAQOfPn3dlfGlvGWI4HG4KoOzav4m/3RBt/x5mt2/f7suYAAAAAAAAbhroPcLq1tbWdObMGb388svKZDKNPbLaefTokarVqjY3N3Xnzh1XeiuVSo2ALRgMKhQKKRqNKhaLtVyu2IlMJqNSqdTVssuD4Z/dgwfC4bBM05SkQ7O33BgTAAAAAADAbUMRhH3++eeS9vbfSqVSXY1hWZYMw3CyLUmHN4k3TVO5XK4xYyoej+vq1asdL2esK5VKWl5e1urqalcnTRaLxabP7Y5xsK5UKjW+BzfGBAAAaMeypFqb3W19huTCf/kAAMCIGPgg7K233mqESoZhqJu9/d0IwKS90OtgKHRQoVBQoVBQMpnsao+y+fl5xePxrk9ZvH//ftPn7fbyqnvxxRebPl9fX2+EVm6MCQAARpOT4VXNkk6827rmyTVpgiAMAAAcY6CDsBs3bjTCo25DMEld36+dUCik1dXVxvLLYrHYWPp3UC6X0/r6ujY2NmyPX98X7caNG133eLCfbmdv7V+W6caYAABgNBFeAQCAQTLQQdjBEKy+91YgEOh4n7Dr16/ryy+/dLzHgzO1qtWqcrmclpaWVK1Wm75WKpUUi8W0trbWdtz6TLK1tTXbM66Oclww16n934sbYwIAADjh1gPpcsTrLgAAwKAa6CCsVCrJMAwFAgGtr6/rpZde6nqshYWFQ0vz3BAIBJROp5VOp1UoFLSwsNAU+BSLRS0vLyudTh87RrVaVSKRUDKZVDQa7amfbsOmg+Hbo0ePXB2zW+VyWTMzMz2Ps7i4qMXFxZ7HAQAA7rn1oH3NlYI04ZMuzrnfDwAA42ZlZUUrKys9j1Mulx3opjsDHYQFAgHt7Ozo6tWrPYVg9bHm5vr7P6J4PK5oNKr5+fmmkxaXlpZaBmHz8/ONZZeDwo3ZW06MWavVtL293fM4u7u7PY8BAADcs/NYevtje7VvfiSd+4HkP+lqSwAAjJ3d3V1HXoN7aaCDsEgkos8//7yr0xKP0steW90KBALa2NhQJBJphGHValXFYvHI2V7Ly8sqlUod7SXW7vpOBE77Z3O5MWa3fD6fpqameh7H7/f3PAbGR7VadeT5CwDYY2c544cl6evv7I331bfSzZL0sx/23hsAAPie3+/X9PR0z+OUy2XVajUHOurcQAdhqVRK9+7dc2xPqn7PCNvvxo0bikS+/x/e2traoSCsVCopk8kom806dpriqVOnHAmt9u/J5saY3ZqamtLW1lbP4wDtlEol3b59W7lcTtFoVPl83uu
      "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": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACbDElEQVR4nOz9TWxbd57v+X9IO4EzSdGUczeCBHSF7DRSq74m7buo2UVk1wUmRtId0gaSSQL8u0QmQaYXQoeMK2uXQlWXNj2TNunqWTgd4NpkpRMkA1RbdHp3F9cSXb0LbhVPUoAErkqi2QliJDbPf6FLNqkH8pA8h4fkeb8AInr46ne+FCnF/Oj34DNN0xQAAAAAAAAw4/xuNwAAAAAAAACMA0EYAAAAAAAAPIEgDAAAAAAAAJ5AEAYAAAAAAABPIAgDAAAAAACAJxCEAQAAAAAAwBMIwgAAAAAAAOAJBGEAAAAAAADwBIIwAAAAAAAAeAJBGAAAAAAAADyBIAwAAAAAAACeQBAGAAAAAAAATyAIAwAAAAAAgCcQhAEAAAAAAMATCMIAAAAAAADgCQRhAAAAAAAA8ASCMAAAAAAAAHgCQRgAAAAAAAA8gSAMAAAAAAAAnkAQBgAAAAAAAE8gCAMAAAAAAIAnEIQBAAAAAADAEwjCAAAAAAAA4AkEYQAAAAAAAPAEgjAAAAAAAAB4AkEYAAAAAAAAPIEgDAAAAAAAAJ5w0u0GgGE8/vjj+vbbb2Wapvz+3nnuD37wAz3xxBNj6gwAAAAAAO/6+uuv9e///u89a5rNpnw+nx577DF98803Y+psn880TXOsVwRscOLECTWbTbfbAAAAAAAAQ/L7/Xr48OFYr8mMMEylziBsmBlhtVpNzWZTfr9f8/PzjvV5kFvX9eq1vXifvXptL95nN6/txfvs5rW9eJ/dvLYX77Ob1/bifXbz2l68z169thfvs5vX9uJ97nVtqzPCpP3X9mNnAlNoYWHBlGQuLCy48vXDcuu6Xr22F++zV6/txfvs5rW9eJ/dvLYX77Ob1/bifXbz2l68z25e24v32avX9uJ9dvPaXrzPo17bzb7ZLB8AAAAAAACeQBAGAAAAAAAAT3B1j7Cnn37azct38fl8+p//83+63QYAAAAAAAAc4moQVq1W5fP5ZLp4cGXr+j6fz7UeAAAAAAAA4DzPL410M4QDAAAAAADA+Lg6I0zaD6ISiYRCoZAr16/X6yoUCq5cGwAAAAAAAOPjehBWKBT005/+1NUeEomEfvKTn7jaAwAAAAAAAJzlehAWi8XcbkHnz593uwWM2crKihqNhgKBgNutjI2b99mta3vxPruNx9o713aLF7/fXrzPbvLq99uLz283efH77dVru8Wr328v/my5aVrvt890cZMsv9+ver3u+jft3r17OnPmjB4+fOhqH7BucXFROzs7WlhY0Pb2ttvtWDatfWO68DyD03iOYRx4nsFpPMcwDjzP4LRpfY652berm+VzUiMAAAAAAADGxdUgbJJObJykXgAAAAAAAGA/V4OwXC7n+rJISTp9+rRyuZzbbQAAAAAAAMBBrgZhb7/99liv99VXXx37uXH3AgAAAAAAgPFyNQgbt3Q67XYLAAAAAAAAcImngrDNzU23WwAAAAAAAIBLTrrdwDC++uor1et1y/WGYSifzw/0NQAAAAAAAJgtUxGEffXVV8rlciqXyzIMY6gxTNOUz+ezuTNgMCsrK2o0GhNxSARmF88zOI3nGMaB5xmcxnMM48DzDE7jOTY4n2mapttN9HL58mWtra1J2g+zRuHz+fTw4UM72oLLFhcXtbOzo4WFBW1vb7vdDgAAAAAAsMjN1/QTvUfYr3/9a+VyuXYAxowuAAAAAAAADGuil0aurq5K2g/ATNNUKBRSJBJRKBSSJD355JN9x/jjH/+oer2umzdv6t69e472CwAAAAAAgMk10UFYpVJpzwLb2NjQ0tLS0GNlMhk9/fTTdrUGAAAAAACAKTPRSyODwaCk/RBrlBBMkkKhkJ566ikbugIAAAAAAMA0muggrLUE8vz587aMl8/nbRkHAAAAAAAA02eig7DWLLDd3V1bxwMAAAAAAID3THQQ9rOf/UymaapSqdgy3q9+9StbxgEAAAAAAMD0megg7PTp03rvvfd048YN/fu///vI47E0EgAAAAAAwLsmOgiT9jfKj0QiSiaTI43z5Zdf2jazDAAAAAAAANPnpNsNWLGxsaFz587p6aefVjab1ZkzZyx93e7urur1uqrVqm7evOlwlwAAAAAAYBaYptQ0e9f4fZLPN55+YJ+pCMI+//xzSVK1WlU6nR5qDNM05eMZOnNqtZoWFxd71qysrGhlZWVMHQEAAAAApl3TlE6+27vmwRXpBDHDIevr61pfX+9ZU6vVxtTNYRMfhL3xxhsqFAqSJJ/PJ9PsE8kegQBsdjWbTe3s7PSsaTQaY+oGAAAAAABvazQafV+nu2mig7Br1661N7gfNgSTNPTXYfL5/X7Nz8/3rAkEAmPqBgAAAAAAbwsEAlpYWOhZU6vV1Gw2x9RRN585wSnRuXPnVKlU2iFYLBZTPB5XMBgceJ+wq1ev6quvvtLDhw8d7hrjsLi4qJ2dHS0sLGh7e9vtdgAAAAAAM+Rh0+LSyIk/gnAyufmafqJnhLVCsGAwqM3NTT311FNDj7W8vKwnn3zSxu4AAAAAAAAwTSY6uwwGg5Kky5cvjxSCtcY6e/asDV0BAAAAAABgGk10EBaNRiVJoVDIlvGuXbtmyzgAvK1er7vdAgAAAABgCBMdhKXTaZmmKcMwbBmPGWEAhlWpVJTNZjU3N6fl5WW32wEAAAAADGGi9whLJBI6e/asbty4ob/9278debyPPvpIf/VXf2VDZwAmSaVS0Y0bN1SpVGQYRvuQjGAwqFAopDNnzigSiSgejysWix36+mQyKcMwtLW11f5YvV5XuVzWjRs3VC6XmQUGAAAAADNgooMwSSoWi/rTP/1T/fM//7P+8i//cqSxVldXCcIwkQqFgtLp9MBfF4lEusKbYfl8voG/JhgMam9v78jPJZNJlUqlgcfc2Ng4Mqg6Sr1e1+rqqtbW1nrWVCoVSVK5XG7XRiIRxWIxhcNhbWxsqFQqKRKJdH3t8vKyDMNofz0AAAAAYPpN9NJIaX9/sKtXr+qnP/3pSOPcu3ePF7SYWKlUSnt7e9ra2lIikehZG4vFtLGx0a63g2maqlaryufz7UMqjhIKhZTP51WtVo8NwaT9AHtvb89SsJXJZLS1taW9vT3LIVhrieLBECyRSKhYLLbHa92vjY0NpVKp9n2rVCpaW1tTOp1uB3YHZ3y1xjFN07Z9CgEAAAAA7vKZpmm63cRxPvroo/bbmUxGPp+vb0hwlHq9rps3b6per+vhw4d2tgiXLC4uamdnRwsLC9re3na7Hdul02kVCoVDH28FPU4yDEPhcPjIz+3t7fUMyo4TjUaPDKLz+bxSqdRAvcXj8UP7BqZSKeXzeUtjrK2tKZvNHvp4KBRStVo98msOPh7jeBwAAAAAuOdhUzr5bu+aB1ekExM/vWgyufmafqKXRv785z/X3bt32++bptlzGVQvpmkOtfwLcMNxQdgwyycHFQqFFIlEDgVXiURiqBBM2u/7qN4HCcHK5bLi8XjXx4LBoG7fvn1oWWMvmUxGiURC0Wi0axbY7u7usV8z7P0GAAAAAEyWic4uL168KNM01Zq0RpAFrzhuKd64lugddZ1Rrn3U1w4SLlUqFVtCsM5+vvzyy64e2AwfAAAAAGbfRAdhrRkkPp+vHYgNewOmyXEh0ZkzZ8bbyAQwDENLS0uHPl4sFocKwVqCwaBte6xNklKpdGjpKP5DuVx2uwUAAAAALpropZGnT59WJBLR3bt3lcvlFIlEhgoCDMPQz3/+c/32t7+1v0lgBh31c/bkk0/aOp7Vn+VkMnlotlYqlbK8sX4voVBImUxm6CXX41SpVJTP51Uul2UYhoLBoM6cOaNgMKhQKKRLly5J2j/tMpfL9Vx2Wi6XlUwmdebMGeXzeVu+l5PKMAyVy2UVi8V2CDZJfxzx0mM
      "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": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOUCAYAAABAB5/EAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACo3klEQVR4nOz9T2xjaZrf+f6oTBs98DTzUFkwIFBAVxx2Ar6rdhwqvLg7l0jXZoyZ7iQVwBhjA+MW6WxceMF2kRnlhduLThVZVdwYmE4yqu/Ccw1YQVb1vZhZNIpH2XuHdLK8G/gWT2QBIggYSPEEexrTsDt57iIu2aTEfyIpHknn+wEEkNLL9zz6E9nNXz3vcyK+7/sCAAAAAAAAHrmdoAsAAAAAAAAAtoEgDAAAAAAAAKFAEAYAAAAAAIBQIAgDAAAAAABAKBCEAQAAAAAAIBQIwgAAAAAAABAKBGEAAAAAAAAIBYIwAAAAAAAAhAJBGAAAAAAAAEKBIAwAAAAAAAChQBAGAAAAAACAUCAIAwAAAAAAQCgQhAEAAAAAACAUCMIAAAAAAAAQCgRhAAAAAAAACAWCMAAAAAAAAIQCQRgAAAAAAABCgSAMAAAAAAAAoUAQBgAAAAAAgFAgCAMAAAAAAEAoEIQBAAAAAAAgFAjCAAAAAAAAEAoEYQAAAAAAAAgFgjAAAAAAAACEAkEYAAAAAAAAQoEgDAAAAAAAAKFAEAYAAAAAAIBQIAgDAAAAAABAKLwfdAHAKv7W3/pb+su//Eu99957+tt/+28HXQ4AAAAAAFjSf/7P/1nffPONfu3Xfk1/8Rd/sdVrR3zf97d6Rch1XeXzebVaraXWlstl2bYt13VlGIZM09TBwYFKpZJM07z19T3P08nJiRzHkeu6cl1XpmnKsiyl02nlcrlVvi1VKhW1Wi2dn5/L87zRnvl8XqlUaqU9Z3nvvfc0GAw2uicAAAAAANienZ0dffPNN1u9JkHYBnmep1gsttTaXC6nWq02d02lUlGpVJq7plwuq1gsLl3jcE/DMPTixQtZlqXd3V2dn5+rVqvJcRyZpqlarbZ0eGXbtrLZrDzPUyqVGgV0juOoVCrJdV2lUik1Gg0ZhrF0rfP8zb/5N/Vf/+t/1c7Ojvb29jay5yzdbleDwWAr17pv1w/rtYO+flivHfT1uTa/87BcO+jrh/XaQV8/rNcO+vphvXbQ1w/rtYO+flivHfT1H+q1h6/9G3/jb+i//Jf/ckcVzuBjY8rlsi9pqY92uz13r1Qq5UvyDcPwM5mMXywW/Uwm45umeWOvRqOxVH2ZTMaX5FuWNXNNLpcb7XtxcbFwz1arNVqfy+WmrrEsy5fkm6bp93q9pWpdJB6P+5L8eDy+kf3uy7Xu2/XDeu2grx/Wawd9fa7N7zws1w76+mG9dtDXD+u1g75+WK8d9PXDeu2grx/Wawd9/Yd67SDrZlj+Bp2cnCy1LpVKzT3SWCqVZNu2yuWyer2eGo2GyuWyGo2G2u22yuXyxPpsNrvwmpVKRc1mU5J0dnY2c12tVhvVdnh4OHdPz/NG1x52kU3TaDQkvTvmuUytAAAAAAAAd4Fh+RtSr9fleZ6KxaLS6fTctQcHBzO/5rruaNbWrKOJxWJR7XZb9Xp99DnHcWRZ1sx9h0csLctaeDwxk8moUqnI8zw1m01lMpmp64bHIcf3n8Y0TWUyGTWbTdm2rXq9vvIcMgAAAAAAgFUFGoR99NFHQV5+QiQS0X/6T/9p5deXy2WZpnmjW+u2SqWSyuXywvlc5XJ5IgizbXtmEGbb9ujxMsP1nz17Nnr8+vXrqUGY67oT+x4dHc3d8/nz56OOtFKpRBAGAAAAAAC2LtAgrN1uKxKJyA9wXv/w+pFIZOU9ms2mXNddOPx+GcOuskWGd490XXf0fJbhGuld59htfPjhh1M/Px74pVKppbrMhhZ1mgEAAAAAANyF0M8I20QId3JyIsMwFnZFLaPVai299urqavR43nHL3d3d0WPXdSeCsWlev349ejyrg2y8G23ekcxx43udnp4u9RoAAAAAAIBNCXxGmO/7ymQySx3Zuwue502EOrflOM6oyyoWi8k0TaVSKaXT6TvtePI8bzSfK5VKzQ2jrv9s8/n83MBteITRMIyp38P1rrLxo5TzWJY1CuGG1wAAAAAAANiWwIOwer2u3/3d3w20hkwmo+9+97srvfb6kHjXdVWv10fhWiaT0YsXL5bumlrWq1evJL0LuYZ3ZZzFsixZljUKsGzbVjabnfq6er0+CqtmzTsbnw02rGEZ19ctGvAPAAAAAACwSYEHYYuGwm/Dsh1N110fGD9Ns9lUs9lULpfbyAwx6V03WD6fl2VZajQaC+dzSdLLly+VTCYn6roehtm2rXw+L0mq1WozB9qPH52U5s8nG3d93tj5+TlB2BIKhYL6/b6i0SjXDtH1gxL09x3Wv7mgf+5BCuvPPejfedDXD0rQ33dY/+aC/rkHKcw/9zB/70EK68896N950NcPykP9viN+gJPqd3Z25Hle4D+0t2/fand3V998882tX1uv1+V5ntrttmzbnjt/y7IsXVxcrFOqXNdVOp3W1dWV3rx5s3QIJf11+DVu2FFm27ZKpZJM01StVpsbUCaTyYnjkcv+CV2/frFYXPkum/v7++p0OorH47q8vFxpj/t4LUDibw7bxd8bto2/OWwbf3PYJv7esG0P9W8uyLoD7Qhb506N98X1rqnhzLGTk5PRDK8hx3GUTqdvNRB/3PUgKRaLqVwuL3WXSendMc1arTbq+pLeBWvDTrFlg6lFw/aXdf3ns4put6v9/f219ykUCioUCmvvAwAAAADAY1WtVlWtVtfep9vtbqCa1QQahAXYjHbDpmoxDEPFYlHFYlHNZlPHx8cTgY9t26pUKkuHV8NgrVarTQ2gSqWSXr9+vXBO2NAwuBsPw8Zr8zxvYZfZqgHW9X3H73q5qsFgoE6ns/Y+/X5/7T0AAAAAAHjM+v3+Rt6DBynQIKxcLgd+LFKSPvjgg5WP6M2TyWSUSqV0eHg4cZTw5ORk6SDMtm21222lUqmZM8mazeatwrVcLqeLi4sbd8t0HEdPnjzR2dnZVmZ3baIjbGdnR3t7e2vvcx/+DgEAAAAAuM+i0aji8fja+3S7XQ0Ggw1UdHuBBmHf+973tnq9r776St/+9re3WothGLq4uJiYq+V5nmzbXupGAZlMRplMZuJz9XpdpVJpIkgqlUrK5XJLzQzL5/Oq1+vK5XI35pp5nqdkMqlWqzWzPsMwNhJi3Wa+2Sx7e3sP6hw0AAAAAAAP1abGCg1nhAVhJ5CrBmTaccBtefny5cTzVeeESX/d0XU9SLre4TVNOp1WvV5XsVhUrVZTu92eenfIdDo9846Yu7u7K9V9V/sAAAAAAAAsI1RB2Pn5eWDXtixrosNq3YHzpmnq7Oxs4nOvX7+e+5pkMjnqRBs/Clqr1abOGMtms1M7v1bt5Lq+1yY6wgAAAAAAAJYV6NHIVX311Ve3Oprnuq5qtdpGjvOtY16X1Sosy1Imk1Gz2ZQ0P1zLZrOjo5m1Wu3G1zOZjC4uLnR4eDj6OXmep1KpdGP9wcHBxMyzZQbsSzeH4ycSiYWvAQAAAAAA2JQHEYR99dVXKpfLN+ZZ3Ybv+4pEIhuu7HZM0xw93tSxwOfPn4+CsFlBn+u6ozWWZU3UMc6yLF1cXEwEVPV6XeVyeSLoSiaTN/ZfZrh+u92eeL7MjLT7oFAoqN/vM1AfW8PfHLaJvzdsG39z2Db+5rBN/L1h2/ibu717H4S9ePFClUpF0rsw6yEbD6A2dSxwPICated4R9fBwcHc/UzTVK1Wm5indn5+PhFaXd9j2SBsPKgzDGNmIHffbGIQIHAb/M1hm/h7w7bxN4dt428O28TfG7aNv7nbu9czwn7605+qXC6PArCgO7rWNT6jLJ1Ob3z/WSHXeAC1zHHE68Pzr3fhWZY1Ebotmk02NP79LwrkAAAAAAAANu1ed4SdnJxIeheA+b4v0zQnjvZ9+OGHC/f4+uuv5XmeXr16pbdv395pvYuMHw3c1LHA8ZBqVrg
      "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": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACtxUlEQVR4nOz9X4xb953f/78OpRjyxktxxkWBAQlsxVkXDn7Ab1fkqBfZq12RzUXtX7JrUiri2gbaiHTcfPeCbYZWjC+6vXAmnKQEvui3qYbKFvjJ6wuZtNeG3R9SD529+wKtZqjsndGElAMMQaBoxCNuDAuJxfO7mCVDcmbIwz/nHHL4fABEODNvfs57/jrnpc8fw7IsSwAAAAAAAMAp5/O6AQAAAAAAAMANBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKBGEAAAAAAABYCgRhAAAAAAAAWAoEYQAAAAAAAFgKZ71uAJjUF7/4RT18+FBnzpzRP/yH/9DrdgAAAAAAgA3/63/9Lz169Ejnzp3Tp59+6uq1DcuyLFevCMyIz+cTP74AAAAAACwmwzDUbrddvSYzwrCwekOwYDDoYSfjaTQaarfb8vl8Wltb87od2+jbXfTtLvp2F327i77dRd/uom930be76Ntd9O2uer0uSZ5MbiEIw8Ly+XzdX/iDgwOv27EtFAqpXq9rbW2Nvl1A3+6ib3fRt7vo21307S76dhd9u4u+3UXf7lrUvs+cOdO9n3ebp0HYU0895eXl+xiGof/5P/+n120AAAAAAADAIZ4GYdVqVYZheLrPU+f6hmF41gMAAAAAAACc5/4ctDnDZusAAAAAAADLwfM9wizLUiKRUDgc9uT6pmmqUCh4cm0AAAAAAAC4x/MgrFAo6Bvf+IanPSQSCX3lK1/xtAcAAAAAAAA4y/OlkbFYzOsWdOnSJa9bAAAAAAAAgMM8nxG2urrqdQtYcO12W6FQaGhNJpNRJpNxqSMAAAAAAJZTPp9XPp8fWtNut13q5ihPgzBOasQ0fvd3f1cPHjyQJNXr9aG1rVbLjZZOtUwmo1arJb/f73UrY6Fvdy1q34tqUb/e9O2uRe17US3q15u+3bWofS+qRf1607e7FrXvedRqtUbeo3f87u/+rsPdHGVYHh6b6PP5ZJqm5z9oDx480MrKiqeJJMYXCoVUr9fl8/m0trY2tHaeZoR1+g4Ggzo4OPC6HQDi9xKYR/xeAvOH30tg/szj76WdGWGNRkPtdtuTvj2dEZbL5TwPwSTp/PnzyuVyXreBCa2trc3NLzwAAAAAAMvMzkSUToDnBU83y//2t7/t6vU++eSTEz/mdi8AAAAAAABwl+enRropnU573QIAAAAAAAA8slRB2N7entctAAAAAAAAwCOe7hE2qU8++USmadqur9Vq2tnZGes1AAAAAAAAOF0WIgj75JNPlMvlVC6XVavVJhrDsiwZhjHjzgAAAAAAALAo5j4Iu379ura3tyUdhlkAAAAAAADAJOY6CHv77beVy+UkSYZhyDAMwjAsvEwmo1arJb/f73UrAP4ev5fA/OH3Epg//F4C84ffy/EZ1hwnSxsbG6pUKt0ALBwOKxKJKBwOS5KefPLJkWP88pe/lGmaeuutt/TgwQM9evTI6bbhklAopHq9rmAwqIODA6/bAQAAAAAANnh5Pz/XM8I6IZgk7e7u6vLlyxOPtbm5qaeeempWrQEAAAAAAGDB+LxuYJhAICDpMMSaJgSTpHA4rAsXLsygKwAAAAAAACyiuQ7COksgL126NJPxdnZ2ZjIOAAAAAAAAFs9cB2GdWWD379+f6XgAAAAAAABYPnMdhH3nO9+RZVmqVCozGe9HP/rRTMYBAAAAAADA4pnrIOz8+fP63ve+p9u3b+vv/u7vph6PpZEAAAAAAADLa66DMOlwo/xIJKJkMjnVOPfu3ZvZzDIAAAAAAAAsnrNeN2DH7u6uNjY29NRTTymbzWp1ddXW6+7fvy/TNFWtVvXWW2853CW80mg0FAqFhtZkMhllMhmXOgIAAAAAYDnl83nl8/mhNY1Gw6VujlqIIOwnP/mJJKlarSqdTk80hmVZMgxjlm1hTrTbbdXr9aE1rVbLpW4AAAAAAFherVZr5D26l+Y+CPvmN7+pQqEgSTIMQ5ZljT0GAdjp5vP5tLa2NrTG7/e71M1RliW1R/zY+gyJH1MAAAAAwDCLcH/p9/sVDAaH1jQaDbXbbZc66mdYkyRLLrl582Z3BtikIVgvwzD06NGjWbSGORAKhVSv1xUMBnVwcOB1Oyd61JbOvja85vPXpTNzv2MfAAAAAMBLp+X+0sv7+bmeEdY55bETgsViMcXjcQUCgbH3Cbtx44Y++eQTB7sFAAAAAADAPJvrIKxSqcgwDAUCAe3t7enChQsTj3Xt2jU9+eSTM+wOAAAAAAAAi2SuJ8sFAgFJ0vXr16cKwTpjXbx4cQZdAQAAAAAAYBHNdRAWjUYlSeFweCbj3bx5cybjAAAAAAAAYPHMdRCWTqdlWZZqtdpMxmNGGIBZME3T6xYAAAAAABOY6yAskUjo4sWLun379kzGe+edd2YyDjBrb971ugOMUqlUlM1mtbKyomvXrnndDgAAAABgAnO9Wb4kFYtF/f7v/77++q//Wn/6p3861VhbW1v6sz/7sxl1BthjJ+R6qXR4vO3zTFqcSKVS0e3bt1WpVFSr1bqnxQYCAYXDYa2urioSiSgejysWix15fTKZVK1W0/7+fvd9pmmqXC7r9u3bKpfLzAIDAAAAsBDevCu9GPW6i/k190FYOBzWjRs39I1vfGOqIOzBgweqVCoz7AwY7cFD6ZX3bBT+tKB/8b20/sWY40cikb7wZlKGYYz9mkAgoGazeezHksmkSqXS2GPu7u4eG1QdxzRNbW1taXt7e2hN5/e+XC53ayORiGKxmNbX17W7u6tSqaRIJNL32mvXrqlWq/F3AwAAAMDcYKLF9OY6COssZfwH/+AfaGVlRU899ZQSicTY45imqbfeemvW7QEjvVGRPv2NjcI/TElPX9Hm/7um2v9va2iIFIvFlM1mtbGx0T1ZdVqdvfjK5bKy2eyJs5/C4bCy2axisdjQQyyKxaJM09Te3p5yuZzK5fKJtZubm7p69arC4bDtzyebzR4bgCUSie5YnfFqtZpqtZqKxaLeeuutbjg2GHANfs7FYrH7fH19fWZ7FQIAAADAJGxPtJD08rvSs1+S/OccbWkhGZZlWV43cZKNjQ3dvfvbuNOyrIlmrvS+9tGjR7NqDx4LhUKq1+sKBoM6ODjwup1jfeW/SB/+bIz6p6Qf/8vDgyIKhcKRjycSib6Axgm1Wk3r6+vHfqzZbE4UvkWj0WNnVu3s7CiVSo3VWzwePxJKpVIp7ezs2Bpje3tb2Wz2yPvD4bCq1eqxrxn8frjxfQAAAACAXv/3/yP9H+/br/+Pz0rf+rJz/UzDy/v5ud4s/8qVK7IsS52sbtIQDPBK87PJ6tPp9LEfP+n9sxQOh48sE5QOw59JZ6Cd1Pc4IVi5XD4yMysQCGh/f992CCYdzkCrVqtHPpf79++f+JpZzbwDAAAAgEm9//F49R+MWb8s5joI69w8G4bRDcQmfQBeWHl8svqTlh0OW444S8ddZ5prH/faccKlSqWieDx+5PUfffTRsaGdnX7u3bvX1wOb4QMAAACYZ5NOtEC/ud4j7Pz584pEIrp7965yuZwikYhWV1fHHqdWq+m73/2ufvrTn86+SWCIZ58eb2nkM08f/u9JIdEkP/+Lrlar6fLly0feXyw
      "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": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMAAAAOWCAYAAAAJOu7yAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACCGklEQVR4nOz9X0yk6X0v+v6KblsT2aELZikSAimZIrPlpXMRddG9L3K3TbF9E5/8GWpGirftGzd4rJx1JPYyuO2LdWUzRWJu1lEyVHtdHOdY2t1g70TJkSJTM7k7F3ugxvvOWivUjCVQSUtyU01ixXJmqHPRmzLQUBRQf18+H6mkKvjV8/6AKuj328/zvKl6vV4PAAAAAEiooV43AAAAAACdJAADAAAAINEEYAAAAAAkmgAMAAAAgEQTgAEAAACQaAIwAAAAABJNAAYAAABAognAAAAAAEg0ARgAAAAAiSYAAwAAACDRBGAAAAAAJJoADAAAAIBEE4ABAAAAkGgCMAAAAAASTQAGAAAAQKIJwAAAAABINAEYAAAAAIkmAAMAAAAg0QRgAAAAACSaAAwAAACARBOAAQAAAJBoAjAAAAAAEk0ABgAAAECiCcAAAAAASDQBGAAAAACJJgADAAAAINEEYAAAAAAkmgAMAAAAgES73esG4Co+9alPxS9/+cu4detW/NZv/Vav2wEAAIAb67//9/8eH3/8cbz00kvxi1/8otftnClVr9frvW4CLuvWrVtxeHjY6zYAAACA/8vQ0FB8/PHHvW7jTGaAMZCOArChoaEYGxvr6LGq1WrXjtVug9x7xGD3r/feGeT+B7n3iMHuX++9M8j9D3LvEYPdv957Z5D7H+TeIwa7f733Trf6PzrOrVu3OnaM6xKAMZB+67d+K/b29mJsbCx2d3c7eqyJiYmuHavdBrn3iMHuX++9M8j9D3LvEYPdv957Z5D7H+TeIwa7f733ziD3P8i9Rwx2/3rvnW71f3Scft6iyCb4AAAAACRaT2eAvfrqq708/AmpVCr+63/9r71uAwAAAIA262kAtrOzE6lUKnq5D//R8VOpVM96AAAAAKBzbvwSSBfBBAAAAEi2nm+CX6/XY3Z2NjKZTE+OX6vVolgs9uTYAAAAAHRezwOwYrEYX/nKV3raw+zsbHzuc5/raQ8AAAAAdEbPl0DmcrletxD379/vdQsAAAAAdEjPZ4CNjo72ugVIrIWFhTg4OIjh4eFet3Ilg9y/3rmKQf/eD3L/eucqBv17P8j9652rGPTv/SD3r3f6Qarew13gb926Ffv7+z1/IT179ixGR0fj448/7mkftG5iYiL29vZifHw8dnd3E3Ms4Pq8Z2GweM/CYPGehcHSrffsIPxu6OkSyH66AmM/9QIAAABA+/Q0ACsUCj2f/RURcefOnSgUCr1uAwAAAIAO6GkA9vWvf72rx/vwww/P/Vy3ewEAAACgO3p+Fchump+f73ULAAAAAHTZjQrAtra2et0CAAAAAF12u9cNXMWHH34YtVqt5fpKpRJra2uXeg4AAAAAyTAQAdiHH34YhUIhSqVSVCqVK41Rr9cjlUq1uTMAAAAA+l3fB2APHz6MlZWViHgeYkG3LSwsxMHBQV9csRS4mPcsDBbvWRgs3rMwWLxnfy1V7+NU6Yc//GHk8/mIiMbsreu0m0ql4uOPP25Lb/TWxMRE7O3txfj4eOzu7va6HQAAALixBuEcva9ngC0vL0fE8+CqXq9HJpOJbDYbmUwmIiJefvnlC8f4+c9/HrVaLZ48eRLPnj3raL8AAAAA9J++DsDK5XJj5tfm5mZMT09feazFxcV49dVX29UaAAAAAANiqNcNNJNOpyPieXh1nfArIiKTycQrr7zShq4AAAAAGCR9HYAdLXW8f/9+W8ZbW1tryzgAAAAADI6+DsCOZn09ffq0reMBAAAAcHP0dQD2zW9+M+r1epTL5baM973vfa8t49A/qtVqTExMNL2trq72uk0AAAAYSKurqxeed1er1V63eaG+3gT/zp078dZbb0WhUIhCoRC/+Zu/ea3x1tbW4itf+UqbuqMfHB4ext7eXtOag4ODLnUDAAAAyXJwcHDhefcg6OsALOL5Bvibm5uRz+fjH/7hH648zgcffNC2mWT0j6GhoRgbG2taMzw83KVuAAAAIFmGh4djfHy8aU21Wo3Dw8MudXQ1qXq9Xu91E624d+9ePHv2LJaWlmJ0dLSl5zx9+jRqtVrs7OzEkydPolarxccff9zhTumGiYmJ2Nvbi/Hx8djd3e11O+eq1yMOL3iHDaUiUqnu9AMAAADtNgjn6H0/Aywi4t13342IiJ2dnZifn7/SGPV6PVJSBrrssB5x+1vNaz76dsQtL00AAADomL4PwN58880oFosREZFKpeIqE9YEXwAAAAA3V19fBfLRo0extrbWCL2uulpzQFZ5AgAAANABfT0DbG1tLSJ+PfMrl8vFzMxMpNPpS+8D9vbbb8eHH37YwW4BAAAA6Ed9HYCVy+VIpVKRTqdja2srXnnllSuP9eDBg3j55Zfb2B0AAAAAg6Cvl0Cm0+mIiHj48OG1wq+jse7evduGrgAAAAAYJH0dgE1NTUVERCaTact4jx49ass4wM1Vq9V63QIAAACX1NcB2Pz8fNTr9ahUKm0Zzwww4CrK5XIsLS3FyMhIPHjwoNftAAAAcEl9vQfY7Oxs3L17Nx4/fhz/8T/+x2uP96Mf/Sj+5E/+pA2dAf2iXC7H48ePo1wuR6VSaVz4Ip1ORyaTidHR0chmszEzMxO5XO6F5+fz+ahUKrG9vd34WK1Wi1KpFI8fP45SqWTWFwAAwIDr6wAsImJ9fT1+93d/N/73//1/jz/+4z++1ljLy8sCMPrOo2Ix3nxz/tLPy2azJ0Kbq0qlUpd+Tjqdjv39/TM/l8/nY2Nj49Jjbm5unhlQnaVWq8Xy8nKsrKw0rSmXyxERUSqVGrXZbDZyuVxMTk7G5uZmbGxsRDabPfHcBw8eRKVSaTwfAACAwdbXSyAjnu//9fbbb8dXvvKVa43z7NkzJ7P0pQdzc7G/vx/b29sxOzvbtDaXy8Xm5majvh3q9Xrs7OzE2tpa48ITZ8lkMrG2thY7Ozvnhl8Rz0Pr/f39lgKtxcXF2N7ejv39/ZbDr6OliKfDr9nZ2VhfX2+Md/R1bW5uxtzcXONrK5fLsbKyEvPz842g7vQMr6Nx6vV62/YgBAAAoHdS9Xq93usmzvOjH/2ocX9xcTFSqdSFAcFZarVaPHnyJGq1Wnz88cftbJEemZiYiL29vRgfH4/d3d1et3Oujw8jbn+rec1H3464dSyKnp+fj2Kx+ELdUcDTSZVKJSYnJ8/83P7+ftOA7DxTU1Nnhs9ra2sxNzd3qd5mZmZe2BNwbm4u1tbWWhpjZWUllpaWXvh4JpOJnZ2dM59z+ufRjZ8DAADAIBmEc/S+XgL5ne98J95///3G43q93nTJUzP1ev1KS72g284LwObnL79M8rIymUxks9kXAqvZ2dkrhV8Rz/s+q/fLhF+lUilmZmZOfCydTsc777zzwvLFZhYXF2N2djampqZOzPp6+vTpuc+56tcNAABA/+jrJZCvv/561Ov1OJqkJsDiJjhvyV23luKddZzrHPus514mVCqXy20Jv47388EHH5zowSb3AAAAydbXAdjRrJFUKtUIwq56g0FxXjg0Ojra3Ub6QKVSienp6Rc+vr6+fqXw60g6nW7bHmr9ZmNj44VlovxaqVTqdQsAAEAP9PUSyDt37kQ2m433338/CoVCZLPZK4UAlUolvvOd78RPfvKT9jcJCXPWe+zll19u63itvo/z+fwLs7Pm5uZa3jC/mUwmE4uLi1deVt1t5XI51tbWolQqRaVSiXQ6HaOjo5FOpyOTycQbb7wREc+vYFkoFJouMS2VSpHP52N0dDTW1tba8v3sV5VKJUqlUqyvrzfCr376T5Gb9LMAAIBe6usALCLijTfeiMnJyfj6179+5THu3r0br7322o2cQQODqlQqnbl5fqFQaNsxHj582AjAarVaX+73VavV4sGDB40rVmYymcbFQCqVSpT
      "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": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intercept= 5308.689740859701\n",
      "coef= {'dSlope_xEndT_abs': 94.8390840129271, 'x_EndT_abs': 0.03952176022974936, 'ideal_state_770_tx^2': 70418.55504271125, 'ideal_state_770_tx dSlope_xEndT': 62087.4827671762, 'dSlope_xEndT^2': 13580.830615214854, 'x_EndT_abs^2': -0.000842287765099925, 'ideal_state_770_ty^2': -2237.8985443626666, 'ideal_state_770_tx^2 x_EndT_abs': -16.447718852307414, 'ideal_state_770_tx dSlope_xEndT x_EndT_abs': -13.929546692369508, 'dSlope_xEndT^2 x_EndT_abs': -3.2033475757333143, 'dSlope_xEndT_abs ideal_state_770_ty^2': 3364.6401097763382, 'x_EndT_abs^3': 1.936634587279551e-07}\n",
      "r2 score= 0.8153420697863665\n",
      "RMSE = 37.11860171392724\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(\n",
    "    data, target, test_size=0.2, 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",
    "#     \"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()\n",
    "# lin_reg = Lasso(alpha=0.00001)\n",
    "# lin_reg = LassoCV()\n",
    "# lin_reg = ElasticNet(alpha=0.0001)\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": 12,
   "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": 13,
   "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": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9T2xbebr/d34o9wwSzLR0JN+dC2jrsHyBWV3rULpAz64tsiqLBDO3TcqFGUyyKIu0b89k4S6TVt1FsrklU3Z585vuNinXJkGAn0TalSBZdJuUezcNtEnKtQtumUcqoLz7mTxiJ0iQXOvMQj+eJiVKovhf4vsFEM0/53y/zyGpNvjU832+Ptd1XQEAAAAAAAAX3MSwAwAAAAAAAAAGgUQYAAAAAAAAxgKJMAAAAAAAAIwFEmEAAAAAAAAYCyTCAAAAAAAAMBZIhAEAAAAAAGAskAgDAAAAAADAWCARBgAAAAAAgLFAIgwAAAAAAABjgUQYAAAAAAAAxgKJMAAAAAAAAIwFEmEAAAAAAAAYCyTCAAAAAAAAMBZIhAEAAAAAAGAskAgDAAAAAADAWCARBgAAAAAAgLFAIgwAAAAAAABjgUQYAAAAAAAAxgKJMAAAAAAAAIwFEmEAAAAAAAAYCyTCAAAAAAAAMBZIhAEAAAAAAGAskAgDAAAAAADAWCARBgAAAAAAgLFAIgwAAAAAAABjgUQYAAAAAAAAxgKJMAAAAAAAAIwFEmEALiTHcYYdAgAAAABgxJAIA3BhlEolJRIJTU9Pa3l5edjhAAAAAABGzE+GHQCA8VAqlbSxsaFSqSTbtlWpVOQ4jgzDkGmampmZkWVZCoVCCgaDR86PRCKybVvFYtF7znEc5fN5bWxsKJ/PUwUGAAAAADiRz3Vdd9hBAOdBOp1WLBY783mWZTUlbzrl8/nOfI5hGKpWqy1fi0QiymazZx4zl8u1TFS14jiOVldXtba2duZ5LMtSMBiU3+9XLpdTNps98l7Wk2OlUunI+eFwWJlM5szzAgAAAAAuLpZGAm2KRqOqVqsqFosKh8MnHhsMBpXL5bzje8F1XZXLZaVSKRmGcexxpmkqlUqpXC4fmwSTpEwmo2q12lZiKx6Pq1gsqlqttp0Eqy9RPJwEqyeo6uPVryuXyykajXrXViqVtLa2plgs5iXsDld81cdxXVemabYVFwAAAABgfFERBnQoFospnU4feX4QlUi2bcvv97d8rVqtnpgoO04gEGhZWZVKpRSNRs8UWygUkm3bTc9Ho1GlUqm2xlhbW1MikTjyvGmaKpfLLc85/HlQEQYAAAAAOIyKMKBDxy2T7GT55FmZpinLso48Hw6HO0qCScfHfZYkWD6fl9/vb0qCGYahYrHYdhJMOqhAK5fLR66lUqkce06n1w0AAAAAGB8kwoAOHbcUb1BL9FrN083crc49S3KpVCopFAodOX9ra6tl0q6deHZ2dppioBk+AAAAAKAbJMKADh2XJJqZmRlsICPAtm0tLi4eeT6TyXSUBKurV5NdRNls9sjyUfxVPp8fdggAAAAALiASYcA51Srhdvny5Z6O125SLxKJHKnWikajbTfWP4lpmorH412PMyilUkmxWEx+v18+n0/T09Py+/0KBALeTp3ZbFbLy8unJnvy+bx3/kVPDNm2rXQ6rVAoJJ/Pd6S6cNjG6bMAAAAALrKfDDsAAOdbPp9v2WQ/mUz2bI6VlRVv90nHcUayH5jjOFpeXvZ2uDRN09td1LZtlUollUol7/X6OSepJxgdx1EsFjt2o4DzzLZtRSIR2bY90ktfx+GzAAAAAMYBFWEAutKqyX43TftbMQzjTE37h6Fe7SUd7LRZLpeVyWSUyWRULBZVLpePVMidlkwZh2W2pmmqWCyqWq2OdOXfOHwWAAAAwDggEQagY/l8vmWfq5WVlZ7PFYlEJJ28c+RJ8vm8IpFI05LFUCikdDrddWxra2vecrl4PN4yaWeapnK5XFOl3Gk9wlKplEzTlGmaymQyXcc56kZtOWSjcfssAAAAgIuKpZEAOnZcQqCbBvnHCQaDSiaTZ67McRxHi4uLR5ZvOo6jfD6vfD6vZDLZVWP/1dVV736rCrlG8Xhc79+/19ra2qlJvWAwOFZL8IZVdRWLxZRKpU48Ztw+CwAAAOCioiIMQMdaNQ3vRYP848Tj8TMtubRtW7Ozsy17mB0+LhAIdNQEvZPeVslkUpZlsWvkCEin0z2pCgQAAABwPpAIA9AR27ZbJnL6UQ3WCcdxvAbn4XBYqVRKxWJRuVzu2F5UoVDozMmpw0mwdpNpKysrI90cfhzYtn1qBR8AAACAi4WlkQA6clzCyO/3DziS1uoJqVwud6RKLRgMKhaLtUx8RSIRFYvFtucxTbPpcSKR0NLS0qmVa/UdJUd1F8yLznGcke5JBgAAAKA/qAgD0JHjEmGjtLteqyRYXb15/WGlUulMSyQNw2hKZDmOo0Ag0FZlmeu6JMGGoN43jqWpAAAAwPghEQZcIIlEQj6fr6NbIBA401zHLesblcSOYRin9iszTbNpF8e6Vs+d5PDxtm3L7/f3pPdUqVRSLBY7tdLOcRyl02kFAoGmeevL/6anp+Xz+eT3+7W2tnbkfNu2lUgkvF01/X6/YrFYy8+5Ps/h79BhoVDIG6/x1ssloel0umme6elpBQKBltcoHbyfrfrGNcZ33Hvd7mdR5ziO1tbWFAqFvPe/Hl8ikWgrEWfbttbW1uT3+5tibvwM/H6/twwYAAAAwClcAB2TdORWrVYHMnc0Gm05fy9vpmmeef5cLjeQ6z8sHo83xREMBts+1zCMI9dxVsFgsOX7YVmWWywWzzRWsVh0o9FoU1yGYRw5rlqtuslk0rUsq2nOVCrluu7R96TxFo1GvXFOOs40zWO/04e/A8dJJpNt/Y0Ui8W2P4NMJuO9L5lMpmkM0zSPjb1arbrFYvFITMVisenWOF47n8VhqVTKOz6ZTLrFYtEtl8tuJpNp+rzC4fCRc6vVqhuPx73raIyxWq0e+bwb4yqXy6fGBgAAAIwzKsKACyQajapYLHZ0S6VSZ5rruMqvUalKOUtlWqvKsdN2mjwsl8t5fb8OjxMIBNqu2LFtWxsbG21XHVmWdaThe7FYVCAQUKlUUiaTUblcVrVabdokIJ1Oq1QqKRQKKZ/PK5fLqVqtqlqtKpfLee+fbdtaXV1tOXckEmkrxlbvSzfy+bw398rKStP4lmUpk8lIOoh9eXm56VzDMGRZ1pFNHerPNb521s+iLpFIKBaLyTAMlctlxeNxWZYl0zQVDodVLBa9mLPZrPx+f9N3o1KpKBQKHfleFgoFzc7Oest6y+WyUqmU91k5jkPzfwAAAOA0w87EAeeZRqwiLJlMdjze4WocnVIRdriipn6rVyMN2uGqplaVNsepVxf14jqOe186GbdxrNOqkA5XCR1XmXe4cq2xMqxR43ty3Ny5XK6tCq5qtdrTirBwOHzqd/60cdqNva7dz6KxEuy0z7qx4uu4CsbGGA3DaPm5Hv7+Dur/gwAAAIDziIowAB0Z9Yqwszi886PU+XXE43FVq9Vjq6Dqu1W2M36ruI7TuElBNBo9tj/a4Z0Sj6sEbIzfcZyR/VyP+x42Pt+LpvjtfhaJRMKbPxqNtnWsdFDl1mqThsbrWF9fb/m5Hv6uFQqFtmIFAAAAxhGJMAAdOS4xUC6XBxxJ91olU96/f9/VeJlMRsVi8cgSPOkg6bG4uNjT5FK7S0EbP7ezLB8dpR0W19fXlUwmlclkTk02SYNLzqbTaW+udhJn0Wi06TM46yYNjRrnG6XPCgAAABg1JMIAdOS4iqNWVS2jrlXS4vLly12Pa1nWsf3XSqXSkf5Vg9Dprp6VSqW3gXTBMAzF4/Fje7Kdtb9br+RyOe9+uxVkjX9H5/FvBwAAADhvSIQB6Firaifbtkd2Gd1ZnGVZ4mmi0ajK5fKRJFQ2mx3p5EenSbNBs21b6XRa+Xy+ZSP8QcZxVoe/ZxfhbwcAAAAYZSTCAHTs1q1bLZ8f5eTOcQ4nfXqZCKuPVywWj8xz1t068Vf1XS9TqdSJfdEGpTG
      "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": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAEAAElEQVR4nOz9b2xb550v+n6XlPame1pxSe6Le2Ff1FqM52Tui1trUQrQ7leNyGQDJ0GbipSCyU4CTEzSbrvnAGpMSi4O9gYuYomyY2CfmbYi5Wxg0gkwEmmnRbuB1FxyX+1d7DG55NwLXORMwiUFiM85wB6Ti2znNLcTc90XGq6Q4r9FcvH/9wMQkciHz/MsUrLDr5/n9wiGYRggIiIiIiIiIiIacRP9ngAREREREREREVEvMAgjIiIiIiIiIqKxwCCMiIiIiIiIiIjGAoMwIiIiIiIiIiIaCwzCiIiIiIiIiIhoLDAIIyIiIiIiIiKiscAgjIiIiIiIiIiIxgKDMCIiIiIiIiIiGgsMwoiIiIiIiIiIaCwwCCMiIiIiIiIiorHAIIyIiIiIiIiIiMYCgzAiIiIiIiIiIhoLDMKIiIiIiIiIiGgsMAgjIiIiIiIiIqKxwCCMiIiIiIiIiIjGAoMwIiIiIiIiIiIaCwzCiIiIiIiIiIhoLDAIIyIiIiIiIiKiscAgjIiIiIiIiIiIxgKDMCIiIiIiIiIiGgsMwoiIiIiIiIiIaCwwCCMiIiIiIiIiorHAIIyIiIiIiIiIiMYCgzAiIiIiIiIiIhoLDMKIiIiIiIiIiGgsMAgjIiIiIiIiIqKxwCCMiIiIiIiIiIjGAoMwIhpJuq73ewpEREREREQ0YBiEEdHIUFUV4XAY09PT8Pv9/Z4OERERERERDZjH+j0BIhoPqqpid3cXqqpC0zRks1noug5RFCFJEmZmZiDLMjweD9xud9XzfT4fNE1DOp0279N1HYqiYHd3F4qicBUYERERERERNSQYhmH0exJEwyAWiyEYDLb8PFmWK8KbdgmC0PJzRFFELper+ZjP50MikWi5z2QyWTOoqkXXdWxsbGBra6vlcWRZhtvthtPpRDKZRCKRqHotS+GYqqpVz/d6vYjH4y2PS0RERERERKOLWyOJLAoEAsjlckin0/B6vQ3but1uJJNJs70dDMNAJpNBNBqFKIp120mShGg0ikwmUzcEA4B4PI5cLmcp2AqFQkin08jlcpZDsNIWxZMhWCmgKvVXuq5kMolAIGBem6qq2NraQjAYNAO7kyu+Sv0YhgFJkizNi4iIiIiIiMYXV4QRtSkYDCIWi1Xd34uVSJqmwel01nwsl8s1DMrqcblcNVdWRaNRBAKBlubm8XigaVrF/YFAANFo1FIfW1tbCIfDVfdLkoRMJlPzOSffD64IIyIiIiIiopO4IoyoTfW2SbazfbJVkiRBluWq+71eb1shGFB/3q2EYIqiwOl0VoRgoiginU5bDsGA4xVomUym6lqy2Wzd57R73URERERERDQ+GIQRtaneVrxebdGrNU4nY9d6bivhkqqq8Hg8Vc/f39+vGdpZmc/h4WHFHFgMn4iIiIiIiDrBIIyoTfVCopmZmd5OZABomobFxcWq++PxeFshWElpNdkoSiQSVdtH6XOKovR7CkRERERENIIYhBENqVqB26lTp2ztz2qo5/P5qlZrBQIBy4X1G5EkCaFQqON+ekVVVQSDQTidTgiCgOnpaTidTrhcLvOkzkQiAb/f3zTsURTFfP6oB0OapiEWi8Hj8UAQhKrVhf02Tu8FEREREdEoe6zfEyCi4aYoSs0i+5FIxLYx1tfXzdMndV0fyHpguq7D7/ebJ1xKkmSeLqppGlRVhaqq5uOl5zRSChh1XUcwGKx7UMAw0zQNPp8PmqYN9NbXcXgviIiIiIjGAVeEEVFHahXZ76Rofy2iKLZUtL8fSqu9gOOTNjOZDOLxOOLxONLpNDKZTNUKuWZhyjhss5UkCel0GrlcbqBX/o3De0FERERENA4YhBFR2xRFqVnnan193faxfD4fgMYnRzaiKAp8Pl/FlkWPx4NYLNbx3La2tsztcqFQqGZoJ0kSkslkxUq5ZjXCotEoJEmCJEmIx+Mdz3PQDdp2yHLj9l4QEREREY0qbo0korbVCwQ6KZBfj9vtRiQSaXlljq7rWFxcrNq+qes6FEWBoiiIRCIdFfbf2Ngwv661Qq5cKBTCw4cPsbW11TTUc7vdY7UFr1+rroLBIKLRaMM24/ZeEBERERGNKq4II6K21SoabkeB/HpCoVBLWy41TcPs7GzNGmYn27lcrraKoLdT2yoSiUCWZZ4aOQBisZgtqwKJiIiIiGg4MAgjorZomlYzyOnGarB26LpuFjj3er2IRqNIp9NIJpN1a1F5PJ6Ww6mTIZjVMG19fX2gi8OPA03Tmq7gIyIiIiKi0cKtkUTUlnqBkdPp7PFMaisFUslksmqVmtvtRjAYrBl8+Xw+pNNpy+NIklTxfTgcxvLyctOVa6UTJQf1FMxRp+v6QNckIyIiIiKi7uCKMCJqS70gbJBO16sVgpWUitefpKpqS1skRVGsCLJ0XYfL5bK0sswwDIZgfVCqG8etqURERERE44dBGNEICYfDEAShrZvL5WpprHrb+gYl2BFFsWm9MkmSKk5xLKl1XyMn22uaBqfTaUvtKVVVEQwGm66003UdsVgMLperYtzS9r/p6WkIggCn04mtra2q52uahnA4bJ6q6XQ6EQwGa77PpXFO/gyd5PF4zP7Kb3ZuCY3FYhXjTE9Pw+Vy1bxG4Pj1rFU3rnx+9V5rq+9Fia7r2NragsfjMV//0vzC4bClIE7TNGxtbcHpdFbMufw9cDqd5jZgIiIiIiJqwiCitgGouuVyuZ6MHQgEao5v502SpJbHTyaTPbn+k0KhUMU83G635eeKolh1Ha1yu901Xw9Zlo10Ot1SX+l02ggEAhXzEkWxql0ulzMikYghy3LFmNFo1DCM6tek/BYIBMx+GrWTJKnuz/TJn4F6IpGIpd+RdDpt+T2Ix+Pm6xKPxyv6kCSp7txzuZyRTqer5pROpytu5f1ZeS9OikajZvtIJGKk02kjk8kY8Xi84v3yer1Vz83lckYoFDKvo3yOuVyu6v0un1cmk2k6NyIiIiKiccYVYUQjJBAIIJ1Ot3WLRqMtjVVv5degrEppZWVarZVjzU6aPCmZTJp1v07243K5LK/Y0TQNu7u7llcdybJcVfA9nU7D5XJBVVXE43FkMhnkcrmKQwJisRhUVYXH44GiKEgmk8jlcsjlckgmk+brp2kaNjY2ao7t8/kszbHW69IJRVHMsdfX1yv6l2UZ8XgcwPHc/X5/xXNFUYQsy1WHOpTuK3+s1feiJBwOIxgMQhRFZDIZhEIhyLIMSZLg9XqRTqfNOScSCTidzoqfjWw2C4/HU/VzmUqlMDs7a27rzWQyiEaj5nul6zqL/xMRERERNdPvJI5omGHAVoRFIpG2+zu5GgdNVoSdXFFTupVWI/XayVVNtVba1FNaXWTHddR7Xdrpt7yvZquQTq4Sqrcy7+TKtfKVYeXKX5N6YyeTSUsruHK5nK0rwrxeb9Of+Wb9WJ17idX3onwlWLP3unzFV70VjOVzFEWx5vt68ue3V38GERERERENI64II6K2DPqKsFacPPkRaP86QqEQcrlc3VVQpdMqrfRfa171lB9SEAgE6tZHO3lSYr2VgOXz13V9YN/Xej+H5ffbURTf6nsRDofN8QOBgKW2wPEqt1qHNJRfx87OTs339eTPWiqVsjRXIiIiIqJxxCCMiNpSLxjIZDI9nknnaoUpDx8+7Ki/eDyOdDpdtQUPOA49FhcXbQ2XrG4FLX/fWtk+OkgnLO7s7CASiSAejzcNm4DehbOxWMwcy0pwFggEKt6DVg9pKFc+3iC9V0REREREg4ZBGBG1pd6Ko1qrWgZdrdDi1KlTHfcry3Ld+muqqlbVr+qFdk/1zGaz9k6kA6IoIhQK1a3J1mp9N7skk0nza6sryMp/j4bxd4eIiIiIaNgwCCOittVa7aRp2sBuo2tFK9sSmwkEAshkMlU
      "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\"]) - 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
}