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

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 62,
   "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",
    "    {\n",
    "        \"/work/cetin/LHCb/reco_tuner/data/tracking_losses_ntuple_B_BJpsi_def_selected.root\": \"Selected\"\n",
    "    }\n",
    ")\n",
    "array = input_tree.arrays()\n",
    "\n",
    "array[\"dSlope_yEndT\"] = array[\"ideal_state_9410_ty\"] - array[\"ideal_state_770_ty\"]\n",
    "array[\"dSlope_yEndT_abs\"] = abs(array[\"dSlope_yEndT\"])\n",
    "\n",
    "array[\"dSlope_xEndT\"] = array[\"ideal_state_9410_tx\"] - array[\"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\"]\n",
    "    - array[\"ideal_state_9410_x\"]\n",
    "    - array[\"ideal_state_770_tx\"] * array[\"ideal_state_770_z\"]\n",
    "    + array[\"ideal_state_9410_tx\"] * array[\"ideal_state_9410_z\"]\n",
    ") / array[\"dSlope_xEndT\"]\n",
    "\n",
    "array[\"z_mag_corr\"] = array[\"z_mag_xEndT\"] - array[\"match_zmag\"]\n",
    "\n",
    "sel_array = array[\n",
    "    (array[\"z_mag_xEndT\"] < 5700)\n",
    "    & (array[\"z_mag_xEndT\"] > 5000)\n",
    "    # & (array[\"match_chi2\"] < 5)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC6HUlEQVR4nOzdT0xjZ77n/8+hezHS6IKhRiMhkDocp2aTxRSG6v0t7Ny7TcqGrHODnepVi+7CRfZTVXZ1UFbTZVPJ+oJd1duesqne38Ku6kU2k/hAJJClmYk5uDXSbJrzW/Czr80fY/zvHMP7JVmN4fHzfA8m3e1Pvs9zDMdxHAEAAAAAAADX3IjbBQAAAAAAAACDQBAGAAAAAACAG4EgDAAAAAAAADcCQRgAAAAAAABuBIIwAAAAAAAA3AgEYQAAAAAAALgRCMIAAAAAAABwIxCEAQAAAAAA4EYgCAMAAAAAAMCNQBAGAAAAAACAG4EgDAAAAAAAADcCQRgAAAAAAABuBIIwAAAAAAAA3AgEYQAAAAAAALgRCMIAAAAAAABwIxCEAQAAAAAA4EYgCAMAAAAAAMCNQBAGAAAAAACAG4EgDAAAAAAAADcCQRgAAAAAAABuBIIwAAAAAAAA3AgEYQAAAAAAALgRCMIAAAAAAABwIxCEAQAAAAAA4EYgCAMAAAAAAMCNQBAGAAAAAACAG4EgDAAAAAAAADcCQRgAAAAAAABuhF+6XQDQa//xP/5H/b//9//0i1/8Qv/5P/9nt8sBAAAAAAAN/tf/+l/6+9//rv/wH/6D/u///b8DXdtwHMcZ6IpAn/3iF7/Q8fGx22UAAAAAAIAWRkZG9Pe//32ga9IRhmunFoSNjIxocnLy0vHlcvlK4zt13dYZ5FrXbZ1BrnXd1hnkWlyT99cZ5FrXbZ1BrnXd1hnkWtdtnUGudd3WGeRaXJP31xnkWtdtnUGudd3W6fVatbl+8Ytf9Ki6K3CGzO7urvPy5UtnY2PDefbsmfPo0SPn0aNHzrNnz5yNjQ1ne3vbOTo6crtMuGhqasqR5ExNTfVl/KDq8vo6g1zruq0zyLWu2zqDXItr8v46g1zruq0zyLWu2zqDXOu6rTPIta7bOoNci2vy/jqDXOu6rTPIta7bOr1ea5B1nzYUHWGvXr1SLpfT1taWbNtu6zWmaSoYDCoSiejevXv9LVBSMplULpfTzs6ObNuWaZoKBAKKxWIKBoNDt7ab1wMAAAAAANAPnr5r5IsXL3T79m1FIhGl02kdHh7KcRw5LY41q/28VCopnU4rFArp9u3b+vbbb/tSYz6f1/j4uOLxuCQpk8moVCopkUioWCwqFAopFAq1HeC5vXYv5xwfH5dhGB09YrFYJ78SAAAAAACAC3kyCHvz5o1u376tWCwmy7Ik6Uz4VQu8Tj/OG2NZlqLRqP7Lf/kv+utf/9qzOvP5fD0UikajyuVyCgaDMk1T4XBYpVJJgUBA+Xxec3NzPQ3D+rF2L+fMZrNdXW8oFOr4tQAAAAAAAOfx3NbItbU1JZPJM6FWMBhUIBDQrVu3ZJqmfD6fJGliYkKSVKlUJEm2bcuyLP3888+yLEv5fL4eyPz4448KBAKKx+N6/PhxV3Xatq1IJCLpZBtmKpU6d1wmk5Hf75dlWYpEIsrlcl2t26+1ez3nRa9vVzgc7ur1AAAAAAAAp3kmCKtWq1pYWFChUJB0EnyFQiEFg0HNzs52NffR0ZE2NzeVTqdVLBaVSCSUz+eVz+c1Ojra0ZyRSKQesNW2EZ6n1k2VzWaVz+eVTqcVjUY7WrOfa/dyzloAaZqm4vG4gsFgPbBsZW5uTpZlcQYZAAAAAADoC09sjXz//r1mZmZUKBQUjUZVKpX0+vVrPXz4sOsQTJLGxsYUjUa1s7OjH3/8Uffu3dPOzo7m5+f1t7/97crz1YKemsXFxZbjl5aW6l+3CpncWrvXc6ZSKZmmWX8/ax18rR6VSqW+DbbWmQYAAAAAANBLrgdh7969UyAQ0MzMjEqlkp4/f66ZmZm+rWeapnK5nF6/fq3//b//t+bm5lStVq80RyKRqH8dDAbr2zQv0rjNz7ZtZbPZK63X77V7PWc6nVYmk7l0nkaNc1wWxAEAAAAAAHTC1a2Ru7u7WlhYUDgc1tbW1kDXDgaD2t3d1dzcnBYXF/XnP/+57dem0+n614FAoK3XmKZZ73ja3Nzs+Aysfqzd6zkzmUzb89Rsbm5Kai+I67WVlRVVq9WOt8l6zSCvZ1BrXbd1Buk6/u6u4zUNynX83V3HaxqU6/a7u45/d4NyHX93122dQeLvwfvrDNJ1/N1dx2salOtyTYZz+lT6Afrwww8VDAb1/Plzt0qQbduan5/XgwcP9Lvf/e7S8cViUXNzc/XnmUymrVArEok0dT118mvvx9puXk+NbdsaHx+XdLKtstsz1Kanp3VwcKCpqSnt7+93NVcvebUuuIO/BzTi7wE1/C2gEX8PaMTfAxrx94CaYf1bcLNu17ZGPnr0SKZpuhqCSZLP59POzo7+23/7b3r//v2l4xvP0pJOOqPacXpcsVhsu8Z+ru3m9dQ0dqSxLRIAAAAAAPSLa0HYrVu3lMlk3Fq+ic/n0/b2dn2rXytv374989p23Lp1q+n5zs5O2/X1c203r6emti0yEAgMfFskAAAAAAC4OVw7I+zhw4duLX2u2dnZtu5QeTos67SDqlQqtV9cH9d283qkk22RtW6yxrtRAgAAAAAA9Jrrd40clAcPHvRknna6xtph27Yn1nbzeiQ13SSh0xsIAAAAAAAAtMPVu0YOUjqdVjwe1wcffNDVPJ0GPqe3/FUqFU+s7eb1SCeH40sn2yLb7UZrV7lc1vT0dNfzrKysaGVlpQcVAQAAAAAwvNbX17W+vt71POVyuQfVdOZGBGHb29tyHEcvX75s686Qg9BpAOXVtTuZs9/bIo+Pj3VwcND1PNVqtQfVAAAAAAAw3KrVak8+Z7vJlSDsyy+/1O7ubt/XqVQqsm27vv3v+fPnXQdhPp+vJ0FSJ4fC92NtN6+n39siR0ZGNDk52fU8o6OjPagGAAAAAIDhNjo6qqmpqa7nKZfLOj4+7kFFV+dKEObz+ZTL5WQYRl/XcRyn6bllWXr//r3u3LnT8ZwTExM9CY4mJiY8sbab11PbFmmaZs+3RUrS5OSk9vf3ez5vp1ZWVlStVgnWIIm/BzTj7wE1/C2gEX8PaMTfAxrx94CaQf8t9OrooOnpadc6y1wJwpaWlpRMJtse7zhOR6FZ42tqodjW1lZXQVgnnU/S2a2DnXaE9Xptt66ncVvkTTkkn3PG0Ii/BzTi7wE1/C2gEX8PaMTfAxrx94Aa/hauzpW7Rs7OztY7gBzHufRhGEZb4y561DiOo0wm01Xt8/PzTc/b7aY6fZi83+/3xNpuXU/jtsh+nA8GAAAAAABwmitBmCQtLCxIkvL5vI6Pjy98JBIJOY6jQCCgXC6nw8PDluNPP3Z2djQzM6NYLKbj42P98MMPXdU9NzfX9Lx2/thlSqVS0/NgMOiJtd26nlog6fP5FAgErvRaAAAAAACATrgWhC0uLioQCOjevXsXjtnd3dWjR48Ui8W0s7OjhYUFjY2NXWmdQCCg169fK5VK6Te/+U23ZZ/poGo3OGrstPL5fB2didWPtd24Htu2lc/nJUnRaLTt1wEAAAAAAHTD1Y6w8fHxlmMSiYRM09Qf//jHrtYyTVPLy8tKpVL6y1/+0tVcgUCg6Tyst2/ftvW6nZ2d+tenwyc313bjetgWCQAAAAAA3OBaECZJr1+/bvnznZ0dxWKxnqw1Pz8vx3HqdyrsxuLiYv3r2oHvl2kcF4/HPbX2oK+HbZEAAAAAAMANrgZhlykWiz0LSiYmJiRJuVyu67kaw7naFr9WGseYptnR+WD9XHuQ19O4LbIxgAMAAAAAAOg3TwdhPp9Pu7u7PZmrFoC1e1fEVgKBQFP4k81mW45vvFPlZd1T6XRa8Xj8wrO6+rF2P6/ntMZtkZFI5EqvBQAAAAAA6IbhOI7jdhEX+fjjj3V4eNj2uVUXOTo6qp9HNj4+rp9//rnr2iz
      "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()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACglUlEQVR4nOz9T2xbZ5rv+/5IO4FzkqZJp9GAIA0qVKWRGnWbtPegzqhL5K7BiZF0h7SB+MYB7i2RSVCnB8IOGSc4Q0ehukrAQR+kTbr6AtfZAbZNVu0EyQWyLTo9O4Ntia6eGbuayylAAoENRKLZCWIkDtcZqMiI+kMukmtxkVzfDyCUZD16+dCSK+bP7/u8PtM0TQEAAAAAAABTzu92AwAAAAAAAMAoEIQBAAAAAADAEwjCAAAAAAAA4AkEYQAAAAAAAPAEgjAAAAAAAAB4AkEYAAAAAAAAPIEgDAAAAAAAAJ5AEAYAAAAAAABPIAgDAAAAAACAJxCEAQAAAAAAwBMIwgAAAAAAAOAJBGEAAAAAAADwBIIwAAAAAAAAeAJBGAAAAAAAADyBIAwAAAAAAACeQBAGAAAAAAAATyAIAwAAAAAAgCcQhAEAAAAAAMATCMIAAAAAAADgCQRhAAAAAAAA8ASCMAAAAAAAAHgCQRgAAAAAAAA8gSAMAAAAAAAAnkAQBgAAAAAAAE8gCAMAAAAAAIAnEIQBAAAAAADAEwjCAAAAAAAA4AnH3W4AGMSTTz6phw8f6tixY/qLv/gLt9sBAAAAAAAW/c//+T/1/fff68SJE/r6669H+tg+0zTNkT4iYINjx46p2Wy63QYAAAAAABiQ3+/X999/P9LHZEcYJlIrCPP7/ZqZmXG7HctqtdpE9t0vLzxPLzxHiec5TbzwHCVvPE8vPEfJG8/TC89R8sbz9MJzlLzxPL3wHCVvPE8vPEdpcp9nq+9jx46N/LEJwjCR/uIv/kJbW1uamZnR5uam2+1YNjc3N5F998sLz9MLz1HieU4TLzxHyRvP0wvPUfLG8/TCc5S88Ty98BwlbzxPLzxHyRvP0wvPUZrc59nq241RRwzLBwAAAAAAgCcQhAEAAAAAAMATXD0a+eyzz7r58B18Pp/+x//4H263AQAAAAAAAIe4GoRVq1X5fD65eXFl6/F9Pp9rPQAAAAAAAMB5nj8a6WYIBwAAAAAAgNFx/dZI0zSVSCQUDoddefx6va5CoeDKYwMAAAAAAGB0XA/CCoWCfvGLX7jaQyKR0M9//nNXewAAAAAAAICzXA/CYrGY2y3o7NmzbrcATJWlpSU1Gg0FAgG3W3GMF56j5J3n6QVe+V564Xl64Tl6hVe+l154nl54jl7hle+lF56nF54jBuMzXRyS5ff7Va/XXf/BfPDggU6dOqXvv//e1T5g3dzcnLa2tjQ7O6vNzU2327FsUvsGph1/NoHxw59LYPzw5xIYP5P659LNvl0dls9NjQAAAAAAABgVV4OwcbqxcZx6AQAAAAAAgP1cDcJyuZzrxyIl6eTJk8rlcm63AQAAAAAAAAe5GoS9+eabI328L7744sjPjboXAAAAAAAAjJarQdiopdNpt1sAAAAAAACASzwVhK2vr7vdAgAAAAAAAFxy3O0GBvHFF1+oXq9brjcMQ/l8vq+vAQAAAAAAwHSZiCDsiy++UC6XU7lclmEYA61hmqZ8Pp/NnQH9WVpaUqPRGItLIgD8gD+bwPjhzyUwfvhzCYwf/lz2z2eapul2E91cvnxZKysrknbDrGH4fD59//33drQFl83NzWlra0uzs7Pa3Nx0ux0AAAAAAGCRm6/px3pG2G9/+1vlcrl2AMaOLgAAAAAAAAxqrI9GLi8vS9oNwEzTVDgcViQSUTgcliQ9/fTTPdf48ssvVa/XdfPmTT148MDRfgEAAAAAADC+xvpopN/vb+8Cu3XrlhYWFgZeyzAMPfvssxyNnBKtbZR+v18zMzNda5eWlrS0tDSizgAAAAAA8K7V1VWtrq52ranVamo2m64cjRzrHWHBYFAPHjxQJpMZKgSTpHA4rGeeecamzjAums2mtra2utY0Go0RdQMAAAAAgLc1Go2er9PdNNZBWDgc1t27d3X27Flb1svn87asg/FhZUcYt2cAAAAAADAagUBAs7OzXWtaO8LcMNZB2MLCgu7evavt7W3b1sN0mZmZ4dZIAAAAAADGhJXxRK1xR24Y61sj3377bZmmqUqlYst6v/nNb2xZBwAAAAAAAJNnrIOwkydP6r333tONGzf07//+70Ovx9FIAAAAAAAA7xrrIEySMpmMIpGIksnkUOvcv3/ftp1lAAAAAAAAmDxjPSOsZW1tTWfOnNGzzz6rbDarU6dOWfq67e1t1et1VatV3bx50+EuAQAA3GWaUtPsXuP3ST7faPpBf/j+AQDgvIkIwj7//HNJUrVaVTqdHmgN0zTl428NAABgijVN6fg73WseXZGO8VciW9kVYPH9AwDAeWMfhL3++usqFAqSJJ/PJ9Ps8beMQxCAAQAAwCkEWAAATI6xnhF27do15fP5dvg1SAg2zNcBAABMmw/vut0BAACAe8Z6R1jrlsfWTrBYLKZ4PK5gMNj3nLCrV6/qiy++cLBbAAAAd1kJuV4tScf80sXTzvcDAAAwbsY6CKtUKvL5fAoGg1pfX9czzzwz8FqLi4t6+umnbewOAABgfDx4KL3xsbXa1z6Szv1ECpxwtCXs8eFd6VJ0fNYBAMCrxvpoZDAYlCRdvnx5qBCstdbp0/zTJwAAGC+mKX3f7P5mZcrDBxXp6++sPeZX30rXK8P1jR9Y3YnXq86udQAAwNHGekdYNBrV559/rnA4bMt6165ds2UdAN5Wr9fbQT0ADMuuQeuf3OvvcT+9J/3yp/19DQ6yayceO/oAABiNsd4Rlk6nZZqmDMOwZT12hAEYVKVSUTabVSgU0uLiotvtAMABO984W4/D2bUTjx19AACMxlgHYYlEQqdPn9aNGzdsWe93v/udLesAGC+tkCoej2t+fl6hUEg+n0+hUEjRaFTxeFzZbFblcvnQr08mk4pGOweu1Ot1lUolJZPJ9jorKyuq1+sjeEYA0L/QE87WTxu7jqQOshPPyXUAAEB3Y300UpKKxaJ+/OMf67/+1/+qv/3bvx1qreXlZf3d3/2dTZ0B9ikUCkqn031/XSQS0cbGxtCP7/P1OG9ziGAwqJ2dnUM/l0wmVSqV+l5zbW1NsVjMUm29Xtfy8rJWVla61lQqu/9kXi6X27WRSESxWEzz8/NaW1tTqVRSJBLp+NrFxUUZhtH+egBwk5UB6eeek279wfqazz83XE+Tzq4jqXbtxGNHHwAAozHWO8IkKRwO6+rVq/rFL34x1DoPHjzgBS3GViqV0s7OjjY2NpRIJLrWxmIxra2ttevtYJqmqtWq8vl819lX4XBY+Xxe1Wr1yBBM2g2wd3Z2LAVbmUxGGxsb2tnZsRyCtY4o7g/BEomEisVie73W81pbW1MqlWo/t0qlopWVFaXT6XZgt3+nV2sd0zRtm1MIAIexa0D6KxHpycesPeZTj0uXIt1r7NoxNcmsfG/s2onHjj4AAEZjrHeEtY4y/vmf/7lCoZCeffbZniHBYer1um7evGl3e4CtgsGgIpGIisWi0um0CoXCgZpW0OOEcDisVCrV3il1mI2NDctD4oPBoGKxmGKxmKLR6KFBdD6fVyqVstyjYRiKx+MH5gamUinl8/lDvyYcDiscDisWiymfz2tlZUXZbNbyY0q74eNh3w8AGJadA9JPnpDef2E3NOvl6ou9B63btWPKNHfX6sbvkwbYnDwUqwHkMb90scuYWbt24rGjDwCA0RjrIOzdd9/V3bs//C3FNM2ux6C6MU1zoONfgBuOCsIGOT7Zr3A4rEgkciC4SiQSA9+UmE6nD+29nxCsXC4rHo93/FowGNTt27cPHGvsJpPJKJFIKBqNduwC297ePvJruCESgFMGGZDe7abHi6d7B2HXE92DnX5YObJpV6BmJzsDyFci0lufWfs+dtuJZ9c6AACgu7E+Gnn+/HmZpinzT/vuCbLgFUcdxRvVEb3DHme
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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_xEndT\"]),\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}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACnmUlEQVR4nOzdT2wj6X3n/w/Vk2ACJ+qiJgggSEDcxUzgnLJdVO8ht7XIzSWDJB6yG4hh+xCL9BjBHoSYnLbPHpl0rEuApMn2XsYwsGrSjg37kAyrJ6ffZVuqcW7GxqweAxIILDBStWwjA8cj/g5a0qTEP0Wy+E/1fgEEJPHhU19K6hnxw+/zPJFms9kUAAAAAAAAcMOtzLsAAAAAAAAAYBYIwgAAAAAAABAKBGEAAAAAAAAIBYIwAAAAAAAAhAJBGAAAAAAAAEKBIAwAAAAAAAChQBAGAAAAAACAUCAIAwAAAAAAQCgQhAEAAAAAACAUCMIAAAAAAAAQCgRhAAAAAAAACAWCMAAAAAAAAIQCQRgAAAAAAABCgSAMAAAAAAAAoUAQBgAAAAAAgFAgCAMAAAAAAEAoEIQBAAAAAAAgFAjCAAAAAAAAEAoEYQAAAAAAAAgFgjAAAAAAAACEAkEYAAAAAAAAQoEgDAAAAAAAAKFAEAYAAAAAAIBQIAgDAAAAAABAKBCEAQAAAAAAIBQIwgAAAAAAABAKBGEAAAAAAAAIBYIwAAAAAAAAhAJBGAAAAAAAAEKBIGwOXNdVMpmceJ5YLKZIJKJqtTrS4xzHUTabbT8+EokoFospn8/L87yx6ykWi0omk4pGo+050+m0bNsee04AAAAAAICgEIQFyPO8drA06BaLxWSa5kTXyufzcl135PrS6bTi8bjK5XLX413XVbFYVDQaVblcHmle27YVjUaVz+clSZVKRfV6XYVCQY7jKJlMKplMThSyAQAAAAAATOqleRdwk4wSILVCo3HYtq1isTjSYzzPUzwe9xWeZbNZHR0dqVQq+aql1d2WyWS6HmOaplKplOLxuGzbVjwe19HRkQzDGKl2AAAAAACAINARFqC9vT1f4xKJxNgdYa2urlGl02m5rivLstodW/V6XZVKRblc7tr4crk8dMllZy2mafYNziqViqTLrrNxagcAAAAAAAgCHWEBKZfL8jxPuVxu6P5fW1tbY19nZ2dHa2trkuR7qWG5XJZt28rlcioUCl33tbq2stms0um0HMfpulYqleo7bzqdbtcwqMOtdY1qtSrbtlUul5XJZHzVDgAAAAAAEJRIs9lszruImyAWi0mS6vX61K5RLpfbyxa3t7fbIVSlUhkYWLX2JKvVagPnd123/TxaarWaEonE0LFnZ2cDlzxWq9V2N5hhGDo7OxtYCwAAAAAAQNBYGhmAarUq13Un2vdrGNd1lc1mlcvlZFmW78c5jiPXddvLEwcxTfNax1hnh1inznGJRGLovl+dQZ3neSOfdAkAAAAAADApgrAA7O3tyTAM3b9/f2rXSKfTsizrWlA1zMHBgTKZjO8N6q92f33wwQc9x3UeDOA3mOvcF+3g4MDXYwAAAAAAAILCHmETchyn3TUVjUZlmqYSiYSSyeTA5YqjyOfzchxnrGWXDx48GGlj/quh1tWlktL1LrF79+75nrt1auWkHWEf+9jH9OGHH+rWrVv6vd/7vYnmAgAAAAAAs/N//+//1UcffaSXX35Zv/jFL2Z6bYKwCV1dDum6rsrlcrtjKpVK6eHDhyMtZ+zkOI6KxaJKpdJYJ02Oet2rG/D3uqZt20PH9HJ1nOM4Y39fPvzwQ11cXOji4kInJydjzQEAAAAAAObnww8/nPk1CcIm4LrutVDoqmq1qmq1qkwmo1KpNPI1tre3lUqlZnbKYqtjq6XXRvnPnj3r+tzvsstXXnml6/PDw8Oxg7Bbt27p4uJCKysrWl9fH2uOUTQajZlezw9qWr56JGryi5r8WbSaFq0eiZr8oiZ/Fq2mRatHoia/qGn56pGoyS9q8mfRapp1Pa3r3bp1a+rXuoogbAKmaapUKsnzPNXrddm2fS1IaimXyzo8PNTR0ZHv+VunLD5+/DiQev04PDxsf9wvfLv6HMftCJvkhM3f+73f08nJidbX13V8fDz2PH5tbm7O9Hp+UNPy1SNRk1/U5M+i1bRo9UjU5Bc1+bNoNS1aPRI1+UVNy1ePRE1+UZM/i1bTrOtpXW8eWx0RhE3oaljkeZ7K5bL29vauLTN0HEfJZFK1Wm3ovK1Oslqt5rvjKgidXWv9TsHsF/aN6ur3ZxyNRkObm5sTz7O7u6vd3d2J5wEAAAAA4Kba39/X/v7+xPM0Go0AqhnPXIOwV199dZ6X7xKJRPR//s//mXgewzCUy+WUy+VUrVa1s7PTFfjYtq1isahcLtd3Ds/zlE6nlclkei5NnBbXddsb4RcKhb6dXuMGWFcDvdPT07Hm6RTUHmHn5+cTzwEAAAAAwE12fn6+9Pt0zzUIq9frikQiajabc6uhdf1IJBL43KlUSolEQtvb210nLe7t7Q0Mwra3t9vLLmepUChIulzCOKi+oATRERbU+uXV1dWJ5wAAAAAA4CZbXV3VxsbGxPO09gibh9AvjZx2CGcYho6OjhSPx9thmOd5sm27Z7dXsViU4zgj7SUWBMdxVC6XZRjG0KWbhmEEEmIFseRzUdZTAwAAAABw0wW1rVBrj7B5mHsQ1mw2lUqlfG+4HrTWnl7T9vjxY8Xj8fbntVrtWhDmOI7y+bwKhcLYpymOa2dnR5L09OnToT+LtbW1QIKwtbW1iecAAAAAAADwa+5BWLlc1uc///m51pBKpfSnf/qnU72GZVlKJBKybVtS7w3n0+m0LMuaybLETtlsVo7jqFKp+Argxu3kuhqezfIQAAAAAAAAgLkHYbPcDL6fe/fuzeQ6yWSyHYRdVSwW5bquEomE0un00Lk6Q6W9vT0dHBy0P3/w4IFSqZSvmsrlssrlskqlku/HbG1tde155nmer1Dr6ub4sVjM1/UAAAAAAACCMPcgLEzL4zqXHF593h988IEk9Q3KBnEcpyuYMk3TV6hl27ay2axKpZIymYzv63Uu8ZQuu9v8dJLV6/WuzxchBPVrd3dX5+fnC7WpPjUtp0X8HlGTP4tY06JZxO8RNfmziDUtokX7Pi1aPRI1LbNF+z4tWj0SNS2zRfw+LVpNi1bPNEWaczyy8datWzo7O5v7N/rFixdaW1vTRx99NNXrOI7TDpFyuVz7lEZJyufzKhaLgVzn6tyDaikUCiMvxex8HpJUqVR8BW/ZbLa9H5thGDo7Oxvpup1aG+ttbGywWT764vcENx2/47jJ+P3GTcbvN24yfr/hxzx/T1ZmerUr5pjBXTOLWg4PD9sfJ5PJrvsKhYKazabvW2d3WaVS6bpvWAjmuq62t7eVy+XG2o/MsqyupZDPnj3z9bjO57+1tTXydQEAAAAAACYx1yCsUCjMvRtMkm7fvj00PApC59LAeS0LdF1X8XhcmUzG93N2Xfdat9r9+/fbH3cuyxykc1w+n/f1GAAAAAAAgKDMNQj70pe+NNPrvf/++33vm0Ut1WpVkmZ+KmSL53lKJpO6f//+SMFfOp2+Ftxls9n2x372NescY5rmUu0PBgAAAAAAboa5b5Y/S9lsVv/yL/8yl2tXq1W5rivDMPTw4cOZX9/zPMXjcZmmqXw+L9d1hz7Gdd1259bVzfAty1IikWgHXNVqdeA+YZVKpf0x3WAAAAAAAGAeQhWEde5RNSnbtpVOp+V5nhKJhAqFQt+TE13X1c7OjiTp6dOnXftrzcr29rZc15XruorFYiM9tlQq9f16a669vb2+QZjnee1N8hOJxEgnVAIAAAAAAARlKYOw999/X57n+R7vuq5KpdJIjxmmUqm057Ntu73v1tXQqBWYra2t6ejoqGuT+1mJx+O+9/HqpV9wZZqmKpWK0um0HMdRsVjsuexze3tb0uVJkZ2dYQAAAAAAALO0FEHY+++/r0KhINu2fS3p66XZbCoSiQRWUzqdbnc5tZTLZT158kSJREJra2s6PDyU4zjK5XJ6+PDhXDr
      "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": 66,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC1HUlEQVR4nOz9TWxbd57mfV+kXYGDStOk0miAEIGqkJMbqVWPSXsWtesSidokSKZC2kCCJMCNEpnKUzMLYkLGlcX0LByF7Cpuep60Radv4HaQhU1Wd4JkURPR6d0Az0Sia3bBVPM4DUggMEBJx+wEMZKYfBYaski9kIfkIQ9fvh9AgGT99D+/Y73Y59L/xdVqtVoCAAAAAAAAFpzb6QYAAAAAAACAaSAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUiAIAwAAAAAAwFIgCAMAAAAAAMBSIAgDAAAAAADAUjjrdAPAKL7//e/rwYMHOnPmjP7iL/7C6XYAAAAAAIBF//t//289fPhQ586d01dffTXVa7tarVZrqlcEbHDmzBk1m02n2wAAAAAAACNyu916+PDhVK/JjDDMpXYQ5na75ff7J369er0+1etZQU/z149ET1bRkzWz1tOs9SPRk1X0ZM2s9TRr/Uj0ZBU9zV8/Ej1ZRU/WzFpP0+6nfb0zZ85M/FpHEYRhLv3FX/yF9vb25Pf7tbu7O/HrBQKBqV7PCnqav34kerKKnqyZtZ5mrR+JnqyiJ2tmradZ60eiJ6voaf76kejJKnqyZtZ6mnY/7es5sdURm+UDAAAAAABgKTg6I+zJJ5908vI9XC6X/tf/+l9Ot4Eh1et1BQKBvjXpdFrpdHpKHQEAAAAAsLwKhYIKhULfmnq9PqVujnM0CKvVanK5XHJyv/729V0ul2M9YHTNZlN7e3t9axqNxpS6AQAAAABguTUajYHP6U5a+j3CODRzvlnZyM/j8UypGwAAAAAAlpvH49Hq6mrfmvZm+U5wPAhrtVqKx+MKBoOOXN80TRWLRUeujfHNysaCAAAAAADA2vZE7c3yneB4EFYsFvXzn//c0R7i8bh++tOfOtoDAAAAAAAAJsvxUyOj0ajTLejSpUtOtwAAAAAAAIAJc3xG2MrKitMtAAOl02k1Go2Z2m+MnubTLP4d0ZM1s9jTrJnFvyN6smYWe5pFs/b3NGv9SPQ0z2bt72nW+pHoaZ7N4t/TrPU0a/1Mkqvl4G7xZ86c0cHBgeN/0ffv39fKyooePnzoaB+wrr2eeHV1lT3CcCq+TrDo+BrHIuPrG4uMr28sMr6+YYWTXyeOLo2cpRMbZ6kXAAAAAAAA2M/RICyXyzk+G0ySzp8/r1wu53QbAAAAAAAAmCBHg7DXX399qtf74osvTn3ftHsBAAAAAADAdDl+auQ0pVIpp1sAAAAAAACAQ5YqCNve3na6BQAAAAAAADjkrNMNjOKLL76QaZqW6w3D0Obm5lAfAwAAAAAAgMUyF0HYF198oVwup0qlIsMwRhqj1WrJ5XLZ3BkAAAAAAADmxcwHYVevXlU+n5d0GGYBgFXpdFqNRmMmTqcFJoGvcSwyvr6xyPj6xiLj6xuzztWa4XTpt7/9rRKJhCR1ZnON067L5dLDhw9t6Q3OCgQC2tvb0+rqqnZ3d51uBwAAAAAAWOTkM/1Mzwjb2NiQdBhgtVotBYNBhcNhBYNBSdLjjz8+cIw//vGPMk1Tt2/f1v379yfaLwAAAAAAAGbXTAdh1Wq1MxNsa2tLa2trI4+VyWT05JNP2tUaAAAAAAAA5sxMB2Fer1f3799XJpMZKwSTpGAwqCeeeMKmzjAr6vW6AoFA35p0Oq10Oj2ljgAAAAAAWF6FQkGFQqFvTb1en1I3x810EBYMBnX37l1dunTJlvE2NzdtGQezo9lsam9vr29No9GYUjcAAAAAACy3RqMx8DndSTMdhK2trenu3bva39+3bTwsFrfbLb/f37eG00oAAAAAAJgOj8ej1dXVvjX1el3NZnNKHfWa6VMj79+/L5/Pp1dffVXvvPPO2OO9++67+vnPf25DZ3Aap0YCAAAAADCfnHymd0/1akM6f/683n77bd26dUv/+q//OvZ4LI0EAAAAAABYXjMdhEmHpz2Gw2ElEomxxrl3756q1apNXQEAAAAAAGDezPQeYW1bW1u6ePGinnzySWWzWa2srFj6uP39fZmmqVqtptu3b0+4SwAAAAAAAMyyuQjCPv30U0lSrVZTKpUaaYxWqyWXy2VnWwAAAJiiVktqDtjd1u2S+C8fAAA4zcwHYb/4xS9ULBYlSS6XS6Ps7U8ABgAAMP+aLensm/1rvrsmneG/fgAA4BQzHYTduHGjs8H9qCGYpJE/DgAAAONhFhcAAJglMx2EHQ3BotGoYrGYvF7v0PuEXb9+XV988cUEuwUAAMBRzOICAACzZKaDsGq1KpfLJa/Xq+3tbT3xxBMjj7W+vq7HH3/cxu4AAAAwa96/K70ccboLAAAwq9xON9CP1+uVJF29enWsEKw91oULF2zoCgAAAE54/+7gmlfK1uoAAMBymukgLBI5/HVeMBi0ZbwbN27YMg6A5WaaptMtAMDSuf9Aeu1Da7WvfiA1Hky0HQAAMKdmOghLpVJqtVoyDMOW8ZgRBmBU1WpV2WxWPp9P6+vrTrcDAAvFygyu96rSV99aG+/Lb6Sb1fF6AgAAi2mmg7B4PK4LFy7o1q1btoz3D//wD7aMA2C2tEOqWCymUCgkn88nl8sln8+nSCSiWCymbDarSqVy4scnEonODNQ20zRVLpeVSCQ64+TzeWaDAcCQ7FrO+NHnw1334yHrAQDAcpjpIEySSqWSdnZ29I//+I9jj7WxsWFDR4D9isWiXC7X0C9Hw5tRjXJtn8936niJRGKkMU8Lqk5imqay2Wzn7yGfz6tSqcgwjE5YZZqmqtWqKpWK8vm8YrFYpz6bzapYLCqRSKhcLh8bf319XRsbGyqXy4RfADAiO5czHnw93LWHrQcAAMth5oOwYDCo69ev6+c///lY49y/f1/VKnPkMZuSyaQODg60s7OjeDzetzYajWpra6tTb4dWq6VarabNzc3OIRUnCQaD2tzcVK1W08HBwal1pVJJBwcH2traUjQa7XvtTCajnZ0dHRwcDKxtay9RzOfzPX8ej8c74fnBwUHnvra2tpRMJjv3Vq1Wlc/nlUqlOiHY0bCrPU6r1bJtn0IAWDZ2Lmf0PTrctYetBwAAy8HVarVaTjdxmu6ljJlMRi6Xa2BIcBLTNHX79m2ZpqmHDx/a2SIcEggEtLe3p9XVVe3u7jrdju1SqZSKxeKxP28HPZNkGIZCodCJ7zs4OOgblJ0mEomcGERvbm4qmUwO1VssFju2b2AymdTm5qalMfL5vLLZ7LE/DwaDqtVqJ37M0c/HND4PALAIfvr/SJ/8YYj6J6Xf/d8nv++//nfpP3xkfay/fUb65Y+t1wMAgOlx8pn+7FSvNqS33npLd+/+acOIVqt1bAaIVa1WSy6Xy67WMCPq9boCgUDfmnQ6rXQ6PaWO7HFaEJZKpSZ+7WAwqHA4fCy4isfjI4Vg0mHfJ/U+TAhWqVQUi8V6/szr9erOnTsKh8OWx8lkMorH44pEIj2zwPb390/9mFHvGwCWnZ3LGV8KS2/8ztoMs8cekV62/k8DAACwUaFQUKFQ6FtTr9en1M1xMx2EXb58ufMw3t5DCOjWbDa1t7fXt6bRaEypG/ucthRvWkv0gsHgsSBsnGuf9LHDhEvVatWWEKy7n3v37umJJ57o2U8MAGA
      "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(\"dy/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACcq0lEQVR4nOzdTWxbd5rn+9+hnMBB0jSpNAYQJGAqVGfgWnWbtGdRuymRVbOIkXSHtAcxbAMXEZkEdWehnpBxZTE9C0chq0qbuciYdPVdOG1gbLJqEiSL6pBO7+5ibB1ndsZUk04BEggMEItmJYiR2Dx3oSZL1At5SPHw8OX7AQjo5dH/PJLoxPz5/3+OYVmWJQAAAAAAAGDCedxuAAAAAAAAABgGgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATIUjbjcA9OPZZ5/Vo0ePNDMzo3/1r/6V2+0AAAAAAACb/s//+T968uSJjh49qm+++Wao1zYsy7KGekVgAGZmZtRoNNxuAwAAAAAA9Mnj8ejJkydDvSY7wjCWmkGYx+PR3Nyc49erVqtDvZ4d9DR+/Uj0ZBc92TNqPY1aPxI92UVP9oxaT6PWj0RPdtHT+PUj0ZNd9GTPqPU07H6a15uZmXH8WntYGLpyuWyFw+FDrxMIBCxJVj6f7/lr0+m0FQ6HLZ/PZ0myAoGAFY1GrWKx2Hc/Tqx5kPn5eUuSNT8/P/C1R+F6dtBTd6PWj2XRk130ZM+o9TRq/VgWPdlFT/aMWk+j1o9l0ZNd9NTdqPVjWfRkFz3ZM2o9TdNrbIblD1CtVpNhGF0fi4uLCgQCh7pWKpVSpVLp+etKpZL8fr9SqZQkKZ/Pq1wuK51OyzRNRSIRRSIR1Wo1V9cEAAAAAAAYNI5GDlAul7Nd2wyN+lEqlZTJZPr6ukgkIkmKx+PKZrOtzwUCAUWjUYVCIZVKJYVCIa2vr8vn8w19TQAAAAAAACe4GoS9+OKLbl6+jWEY+t//+38fao3V1VVbdeFwuO8dYbVaTbFY7FBfFwgE2gKrnfL5vBYXF1WpVBSLxVQsFoe6JgAAAAAAgFNcDcLK5bIMw5Dl4o0rm9c3DONQ6+RyOdVqNSWTydYOqYOcPHmy7+ssLy9rdnZWkno6ahiLxVr1nXajNXdxFQoFlUol5XI5xePxoa0JAAAAAADglKk/GjmoEC6dTisQCCidTg9kvf3kcjkVCgWtr69raWnJ9tdVKhWVSqXW+2fOnOlYf/bsWRUKBUnbAdd+oZUTawIAAAAAADjJ9SDMsixFo9FDD4/vV61W62m2134KhYIqlcqBRwMHoVKpKJFIKJlMKhgM9vS1O8O5cDjcdUZXNBptvV2r1VQoFNo+5tSaAAAAAAAATnI9CMvlcnr99ddd7SEajeqnP/1p31+/uroqn8/XdVfUYcRiMQWDwb52nO0M+uyGaIFAoHVXyhs3buwJrZxYEwAAAAAAwEmuB2HhcNjtFnTq1Km+v9Y0TZmmKUny+/0KBAIKh8OKRCIDC3pSqZRM01S5XO6rv53sfq/BYLAVWjWPNDq55qhbWVlRvV6X1+t1u5UWehpPo/gzoid7RrGnUTOKPyN6smcUexpFo/ZzGrV+JHoaZ6P2cxq1fiR6Gmej+HMatZ5GrR8nGZaLk+o9Ho9qtZrrP+iHDx9qdnZWT5486flrI5FI26ys3aLRqC5dutTzccYm0zQVCoWUzWbb5mr5/f7WoPp8Pn9g6JbJZNoG2a+vr9vqJZVKKZPJ7Pt1TqzZq4WFBW1ubmp+fl4bGxt9rYHJx/MEk47nOCYZz29MMp7fmGQ8v2GHm88Tz1Cvtsth79Tott0D4/dTKBQUCoWUSCT6usbS0pKi0Wjfw+Vv377d9n63WV5Nzz//fNv7d+7ccXRNAAAAAAAAp7l6NNLFzWh79NNLIBBQNptVrVZTuVxWqVRqHf3bLZfL6c6dO1pfX7e9fiwWkyRdvXq1596advdj96YEu+t2Hst0Yk0AAAAAAACnuRqEpdNp149FStKxY8f6GkIvac9OreZdKFdXV1tHF5tM01QkElGxWOy6bqFQUKFQULFYtL3jaj8HBXO92vm9OLFmv6rVqhYWFg69zsrKilZWVg69DgAAAAAAk2ptbU1ra2uHXqdarQ6gm/64GoS9/fbbQ73el19+qR/84AeO9uLz+ZRMJpVMJlUoFLS8vNwW+JRKJWUyGSWTyQPXqNVqisViisfjh76ZQL9h0+7w7cGDB46u2a9Go6HNzc1Dr1Ov1w+9BgAAAAAAk6xerw/kNbibXL9r5DAlEgn94z/+49CuF41GFQ6HtbS01HanxdXV1Y5B2NLSUuvY5agYxO4tJ9b0eDyam5s79DqjsDMRAAAAAIBR5vV6NT8/f+h1qtWqGo3GADrq3VQFYW4MZ/f5fFpfX1coFGqFYbVaTaVSad/dXplMRqZp9jRLrNv1BxE47dzN5cSa/Zqbm+NOJAAAAAAADMGgxgo17xrphrEMwr788suegphKpdIaau+Wq1evKhQKtd4vFot7gjDTNJVKpZROpxUMBgdy3dnZ2YF837Ozs46uCQAAAAAA4LSxCMK+/PJLpdPpjndl7MayLBmGMeDO7AsGgwqHwyqVSpL2Hzgfi8UUDAY7HpvsVb+7rnYHXbt3hA16TcAJKysrqtfrHH3FxOI5jknG8xuTjOc3JhnPb4y6kQ/CLl26pEwmI2k7zBpnkUikFYTtlslkVKlUFA6HFYvFuq61M1RaXV3VjRs3Wu+fPXtW0WhUknTy5Mm2+WS1Ws1WALV7kP3i4mLrbSfWBJzAnUAx6XiOY5Lx/MYk4/mNScbzG6NupIOw3/zmN0qn05IkwzBkGMZYh2GBQKD19u5jgV999ZUkHRiUdWKaZlswFQgEWkHYzuOY0vZONDvHLsvlctv7O49xOrEmAAAAAACA00Y6CFtdXZWkVgAWCAQUDAZbgdLzzz/fdY2vvvpKtVpNN2/e1MOHDx3tt5udQdiwjgWePHmy7X27odXOHWc+n6+tdyfWBAAAAAAAcNpIB2GmabbmehWLRS0tLfW9VjKZ1Isvvjio1vqy866VkUik7XPpdLq1+82OxcXF1pyxfD7f2gG2WzAYbLvL4+3btw+sPajX3cGXE2sCAAAAAAA4zeN2A500d00lk8lDhWDS9m6sF154YQBd9W/n0cBhHgs8c+ZM6+2dRyg72VmXSqWGsiYAAAAAAICTRjoIax6dO3Xq1EDWy2azA1mnX4VCQZIGeldIOxKJROttOzPIdtYEAoF9Qzsn1gQAAAAAAHDSSAdhzV1gu+82eNj13FAoFFSpVOTz+XTp0qWhXjsYDLYFT81A7iD5fL719kE7t5xYEwAAAAAAwEkjHYT9/Oc/l2VZto/edfPrX/96IOtI2zuc/H6/DMNQJBLp2GOlUtHy8rIk6datW0MblL/Tzt1wzZsQ7KdWqymXy0naPr4Zj8eHuiYAAAAAAIBTRjoIO3bsmN5//33duHFDf/zjHw+93iCPRubz+daw+FKppFAo1HZcsKn5udnZWZXLZVt3V3RCIBBo7coyTVOZTGbfuuauOZ/P17aLa1hrAgAAAAAAOGWkgzBpe55WMBhULBY71Dr3798f2M4ySfv2k8vl5Pf7FYvFlEgkFAqFFIlEFI/Htb6+3pp55pZoNKpisSifz6dUKqVYLCbTNFWr1VqBnWmaCgaDun//vq2da06sCQAAAAAA4ATDsizL7SbsOHnypB4+fKhUKqXZ2VlbX/PgwQPVajWVy2X
      "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": 68,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOVCAYAAACLW0xhAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC5qklEQVR4nOz9X4xb933ve39IOYZ84lDkuNjAYAZITNaFc9WKlHqRc9UM2WzgsWA3JqUNe1sG9pMhHT/ZveCuSSvGQdMLeUJml8BBexyRSgscuQaORCa1YT+Aq6HcuwLP1gyV3hk75bIDzIDABqxZ4rZhwZbJ52I2GXL+LC6Si//fL2AQjeY7v/VdmlEsfub3+y5Xs9lsCgAAAAAAAJhz7kk3AAAAAAAAAIwDQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFgJBGAAAAAAAABYCQRgAAAAAAAAWAkEYAAAAAAAAFsJDk24AGMTXv/513b9/XydOnNC/+3f/btLtAAAAAAAAm/7H//gf+uqrr3Ty5El99tlnY722q9lsNsd6RcABJ06cUKPRmHQbAAAAAABgQG63W1999dVYr8mOMMykVhDmdru1vLzc9+fXarWhPn9Qk7ruol57Ee95Ua+9iPc8yWsv4j1P8tqLeM+TvPYi3vMkr72I9zzJay/iPS/qtRfxnid57UW852Gv3frcEydOjKg7C01gBq2srDQlNVdWViby+YOa1HUX9dqLeM+Leu1FvOdJXnsR73mS117Ee57ktRfxnid57UW850leexHveVGvvYj3PMlrL+I9D3vtSfbNjjDMtFqtptXVVcuaZDKpZDI5po4AAAAAAFhcuVxOuVzOsqZWq42pm8MIwjDTGo2Gdnd3LWvq9fqYugEAAAAAYLHV6/Wer9MnaaJB2BNPPDHJy3dxuVz67//9v0+6DfTJzllkj8czpm4AAAAAAFhsHo9HKysrljWtGWGTMNEgrFqtyuVyqTnBB1e2ru9yuSbWAwa3vLysnZ2dSbcBAAAAAABkbzzR6urqxHaNuSdy1SkyyRAOAAAAAAAA4zPxGWHNZlPRaFR+v38i1zdNU4VCYSLXBgAAAAAAwPhMPAgrFAr6wQ9+MNEeotGovve97020BwAAAAAAAIzWxIOwcDg86RZ09uzZSbeAMUsmk6rX6ws1SH+S9zypay/iPU8aX+vFufakLOKf9yLe8yQt6p/3In5/T9Ii/nkv6rUnZVH/vBfx79Ykzep9u5oTHJLldrtlmubE/9Du3bunpaUlffXVVxPtA/a1BuutrKzM1LD8We0bs4XvM4wa32MYB77PMGp8j2Ec+D7DqM3q99gk+57osHye1AgAAAAAAIBxmWgQNk1PbJymXgAAAAAAAOC8iQZhmUxm4sciJenUqVPKZDKTbgMAAAAAAAAjNNEg7JVXXhnr9T7++ONjPzbuXgAAAAAAADBeEw3Cxi2RSEy6BQAAAAAAAEzIQgVhW1tbk24BAAAAAAAAE/LQpBsYxMcffyzTNG3XG4ahfD7f1+cAAAAAAABgvsxEEPbxxx8rk8moXC7LMIyB1mg2m3K5XA53BvQnmUyqXq9PxUMiML/4PsOo8T2GceD7DKPG9xjGge8zjBrfY/1zNZvN5qSbsHLp0iVls1lJ+2HWMFwul7766isn2sKEra6uand3V263W8vLy5a1yWRSyWRyTJ0BAAAAALC4crmccrmcZU2tVlOj0dDKyop2dnbG1Nm+qd4R9stf/lKZTEbSfojlcrmGDsMwXxqNhnZ3dy1r6vX6mLoBAAAAAGCx1ev1nq/TJ2mqg7CNjQ1Jagdgfr9fwWBQfr9fkvTYY4/1XOOTTz6RaZq6ceOG7t27N9J+MX52doSxRRQAAAAAgPHweDxaWVmxrGntCJuEqT4a6Xa723O9bt68qbW1tYHXMgxDTzzxBEcj50TraOQktlECAAAAAIDBTfI1vXusV+uT1+uVJKVSqaFCMEny+/16/PHHHegKAAAAAAAAs2iqg7DWEcizZ886sl4+n3dknWEZhqFIJGK7NpFIKBAIyOVyyefzKRQKKZFIDPwEzUnq594BAAAAAACcNNVBWGsX2N27dx1db1RM02wP9bd6CwQC7ZDPSjabVSAQUKFQaIdepmmqUqmoUCgoEAi0n6jZi8/ns9XbUW+JRGLs9w4AAAAAAOC0qQ7CfvzjH6vZbKpSqTiy3i9+8QtH1jlOoVCwXZtOpy0/HolElE6n5fV6FY1GlUqlFI1GD4VI6XRapVLJcq1SqSTTNG33dlQvvTh57wAAAAAAAKMw1cPypf1dUZlMRh9//LG+8Y1vDLXW2bNndfv2bYc6O8zn89kKnMLhsDY3N4/9eDqdbt93KpU69PFsNnsoTLL6MkYiEZXL5Z59HcfOt4hT924Xw/IBAAAAAJhNk3xN/9BYrzaAVCqlzc1NxWIxvf/++wOv89FHHzm2s+wohUJBpmkqlUr13EF15syZYz9mGIay2aw2NzcVDoePrEmlUqpWq127sCqVioLB4JHrlctl+f1+pdNphcNhLS0t9byfUCgkwzCO7aGTU/cOAAAAAAAwSlO/I6zlzJkzunfvntLptK0gR9qfLWaapqrVqm7cuCHTNPXVV1+NpL9AICBJqlarQ60Ti8V09uzZI3eCdTJNUz6fr/3+cbvHWkcnt7e320/h7MUwjPb95PN5xeNxy3qn7r0f7AgDAADzptmUGj3+Ze52SS7XePoBAGBU2BHWwwcffCBpP2ixM7j9KM1mU64R/auhVCrJMAxHnkrZ2lnVi9frld/vbw/RPy7kKhQKunXrlu0QTFLXzLHz58/3rHXq3gEAABZZoyk99Jp1zYPL0gmCMAAABjb1QdgPf/jD9hFAl8tla17VQaMKwFo2Njbk9Xp7hkZ29DM/q/NpmscdOSwWi0cembRy/fp1SfvzvHoFaE7eOwAAwCxiJxcAALNjqoOwq1evtncaDRqCSfaGvQ+qUqm0Z4/5fD75/X6Fw2FFIhFFo9GRXdc0zfZw+nA4fGzYZWfG18F1W/cTi8Usayd17wAAzBMnQxQCmclgJxcAALNjqoOwgyFYK2Txer19zwm7cuWKPv74Y8d7PPj0RsMwVCgU2rvYotGoLl261PeurF5u3LghSfL7/SoWi46t2zmAv9cur0ndOwAA88TJEIVABgAAwNpUD8t3u91yuVzyer3a2trS448/PvBapmnqsccec3RYfudQ+V7i8bhjc7Rag/KDwaCKxaL8fr8j60r7T4tsPYFye3v72LpJ3XsLw/IBAPPiq4bN8Mo93rVgn1N/7nbW+b+j0sVQf/0BADBtGJZ/DK/Xq3v37unSpUtDhWCttU6fPu1QZ/v8fr/y+Xz7yZTlcrk9vP6gQqGgra0ty3DJDsMw2rvi+h2C30vnscgLFy5Y1k7i3gEAAGbVW3d6B1hv3em9zoul/UDteWf/WQsAwMKY6p8HhkL7/1pwasfT1atXHVmnUzweVyqVUj6fV7Va1d7enjKZzJEBVaVSUSQSGfhapVJJgUBAhmG0d4Vls9khuu/WOm4pydaMr3HeOzBNWvP5AGAW2Qlb0B+7AZZV3b370svv2LveS29L9fv2agEAQLepPhpZKpV0/vx5ZbNZ/cVf/MWk2+lbqVTS+vr6oRfNmUxGqVTK1hqmaapQKCifzx+74yoajToyJ8zusUg7nLh3K61tlG63W8vLy0Ovl0wmlUwmh14H86tSqej69esqFAoKh8OOzuYDsNicPM54bXs/cOnlH86zo8gp9+5LK69Ln33Zu/bRh6XdS5Ln5OGP/e2/SP/5Xfv
      "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": 69,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACe5UlEQVR4nOzd349ieX7f/xc1k2gjx/ShxopUAsnbB0/kXK0aqnPxvfMWrG9iJd6BbimW7Ytsg8f6KhfEC9vra08teJebSPFArW8cWUo37CZRcuEMp8d/QFcxmzsrXk7PSoWQIm1xGsfKKvFwvhf1hYUqfkNxoM7zIZUExYfPefOju+e85v35nIDruq4AAAAAAACAe+7A6wIAAAAAAACAbSAIAwAAAAAAgC8QhAEAAAAAAMAXCMIAAAAAAADgCwRhAAAAAAAA8AWCMAAAAAAAAPgCQRgAAAAAAAB8gSAMAAAAAAAAvkAQBgAAAAAAAF8gCAMAAAAAAIAvEIQBAAAAAADAFwjCAAAAAAAA4AsEYQAAAAAAAPAFgjAAAAAAAAD4AkEYAAAAAAAAfIEgDAAAAAAAAL5AEAYAAAAAAABfIAgDAAAAAACALxCEAQAAAAAAwBcIwgAAAAAAAOALBGEAAAAAAADwBYIwAAAAAAAA+AJBGAAAAAAAAHyBIAwAAAAAAAC+QBAGAAAAAAAAXyAIAwAAAAAAgC8QhAEAAAAAAMAXCMIAAAAAAADgC+96XQCwil/4hV/Qz372M73zzjv6R//oH3ldDgAAAAAAWND//J//U1988YW+9KUv6W//9m+3euyA67ruVo8IbMA777yjfr/vdRkAAAAAAGBFBwcH+uKLL7Z6TDrCsJcGQdjBwYGOjo7u9FidTmdrx9q14/v12F4f36/H9vr4HJvP3C/H9vr4fj2218f367G9Pr5fj+318f16bK+P79dje338fT324LnvvPPOHVU3gwvsoXA47Epyw+HwvTrWrh3fr8f2+vh+PbbXx+fYfOZ+ObbXx/frsb0+vl+P7fXx/Xpsr4/v12N7fXy/Htvr4+/rsb2sm83yAQAAAAAA4AueLo18//33vTz8mEAgoP/xP/6H12UAAAAAAADgjngahLVaLQUCAbke7tc/OH4gEPCsBgAAAAAAANw93y+N9DKEAwAAAAAAwPZ4ftVI13WVSqVkmqYnx3ccR9Vq1ZNjAwAAAAAAYHs8D8Kq1aq+8Y1veFpDKpXSr//6r3taAwAAAAAAAO6W50sjE4mE1yXo8ePHXpcAAAAAAACAO+Z5R9jh4aHXJQA7LZfLqdfrKRgMcmwfHd8rXr9uv37nvH7fveTX993rz9zr43vF69ft1++c1++7l/z8vvv5tXvJr++715+518f3yr6+7oDr4W7x77zzjrrdrudv2tu3b3V4eKgvvvjC0zqwuEgkona7rYODAx0dHc0cm8vllMvl1j5WOBzW5eXlyvMAi+I7h23i+4Zt4zuHbeM7h23i+4Zt28XvXLlcVrlcnjmm0+mo3+97UrenHWG7dMXGXaoFi+v3+2q32zPH9Hq9LVUDAAAAAIC/9Xq9uefpXvI0CCsWi553g0nSgwcPVCwWvS4DK1ikI2wXvmMAAAAAAPhBMBhUOByeOWbQEeYFT4Owb37zm1s93ueff64vf/nLO1ELNuPo6Ghn2j8BAAAAAPC7RbYnGizp9ILnV43cpmw263UJAAAAAAAA8IivgrDz83OvSwAAAAAAAIBHPF0auarPP/9cjuMsPN62bVUqlaWeAwAAAAAAgPtlL4Kwzz//XMViUZZlybbtleZwXVeBQGDDlQEAAAAAAGBf7HwQ9vz5c5VKJUnXYRawbblcTr1ej6tPYmv4zmGb+L5h2/jOYdv4zmGb+L5h2/jOLS/g7nC69IMf/EDpdFqSht1c65QbCAT0xRdfbKQ2eGtwhYlwOMxVIwEAAAAA2CNentPvdEfY6emppOsAy3VdmaapWCwm0zQlSe+9997cOX7605/KcRy9fPlSb9++vdN6AQAAAAAAsLt2OghrNpvDTrBGo6GTk5OV58rn83r//fc3VRoAAAAAAAD2zIHXBcxiGIak6xBrnRBMkkzT1MOHDzdQFQAAAAAAAPbRTgdhgyWQjx8/3sh8lUplI/MAAAAAAABg/+x0EDboAru6utrofAAAAAAAAPCfnQ7Cvv3tb8t1XTWbzY3M9/3vf38j8wAAAAAAAGD/7HQQ9uDBA33nO9/Rixcv9Dd/8zdrz8fSSAAAAAAAAP/a6SBMut4oPxaLKZ1OrzXPmzdvNtZZBgAAAAAAgP3zrtcFLKLRaOj4+Fjvv/++CoWCDg8PF3re1dWVHMdRq9XSy5cv77hKAAAAAAAA7LK9CMI+/fRTSVKr1VI2m11pDtd1FQgENlkWAAAAAADA1riu1HdnjzkISMQf0+18EPbhhx+qWq1KkgKBgFx3zic+AQEYAAAAAADYd31XevcPZ4/5uz+S3iEGmWqng7Czs7PhBverhmCSVn4edl+n01EkEpk5JpfLKZfLbakiAAAAAAD8q1wuq1wuzxzT6XS2VM1tOx2E3QzBEomEksmkDMNYep+wjz/+WJ9//vkdVgsv9Pt9tdvtmWN6vd6WqgEAAAAAwFt//pn0O3Hvjt/r9eaep3tpp4OwZrOpQCAgwzB0fn6uhw8frjzXs2fP9N57722wOuyCg4MDHR0dzRwTDAa3VA0AAAAAAHfnzz+bP+Z369I7B9JvPbr7eiYJBoMKh8Mzx3Q6HfX7/S1VNC7g7vC6wcPDQ719+1bFYlF/8Ad/sPZ8x8fHOj8/30Bl8FokElG73VY4HNbl5aXX5QAAAAAAcKfe/kwKfyT97f+dP/Yf/n2p/VwKfunu61qFl+f0B1s92pLi8etePtM0NzLf2dnZRuYB4G+O43hdAgAAAACf+ffNxUIwSfpf/0f6s+bd1rOvdjoIy2azcl1Xtm1vZL5HjzzqCwSw95rNpgqFgkKhkJ49e+Z1OQAAAAB85r/81XLj/+uS4/1ip/cIS6VSevTokV68eLGRpZE//OEP9fWvf30DlQHYJc1mUy9evFCz2ZRt28OLZBiGIdM0dXh4qFgspmQyqUQicev56XRatm3r4uJi+DvHcWRZll68eCHLsugCAwAAAOCp7v++2/F+sdNBmCTVajX9yq/8iv7jf/yP+s3f/M215jo9PSUIw06qVqvKZrNLPy8Wi42FN6sKBAJLP8cwDHW73YmPpdNp1ev1pedsNBoTg6pJHMfR6empSqXSzDHN5nU/sGVZw7GxWEyJRELRaFSNRkP1el2xWGzsuc+ePZNt28PnAwAAAICXQv/gbsf7xU4vjZSu9wf7+OOP9Y1vfGOted6+fcsJLXZWJpNRt9vVxcWFUqnUzLGJREKNRmM4fhNc11Wr1VKlUpFhGFPHmaapSqWiVqs1NQSTrgPsbre7ULCVz+d1cXGhbre7cAg2WKJ4MwRLpVKq1WrD+Qavq9FoKJPJDF9bs9lUqVRSNpsdBnY3O74G87iuu7F9CgEAAABgVb/xq8uN/2dLjveLnb5q5A9/+MPh7Xw+r0AgMDckmMRxHL18+VKO4+iLL77YZInwyH2/amQ2m1W1Wr31+0HQc5ds21Y0Gp34WLfbnRmUTROPxycG0ZVKRZlMZqnaksnkrX0DM5mMKpXKQnOUSiUVCoVbvzdNU61Wa+Jzbn4e2/gcAAAAAGAUV43cjJ1eGvnRRx/ps88+G953XXfmMqhZXNddafkX4IVpQdgqyyeXZZqmYrHYreAqlUqtFIJJ13VPqn2ZEMyyLCWTybHfGYahV69e3VrWOEs+n1cqlVI8Hh/rAru6upr6nFVfNwAAAABsyoMvSf/un0u/u8AuNB//i90Nwby200sjnzx5Itd1NWhaI8iCX0xbiretJXqTjrPOsSc9d5lwqdlsbiQEG63nzZs3YzWwGT4AAACAXfdbj+aP+bPUYuP8aqeDsEEHSSAQGAZiq/4A+2RaSHR4eLjdQnaAbds6OTm59ftarbZSCDZgGMbG9ljbJfV6/dbSUfycZVlelwAAAADcqX9JCDbTTi+NfPDggWKxmD777DM
      "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": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACseklEQVR4nOz9X4xb933n/78OpRhy41DkuFhgQAKNOHXhXnVFjvYivWqGbBaojaQ1KS3in21gG3Ec/7K94K6HUYzvpTPmpCXwRfeXilT2C/yUGlhpmKwDp0Cqobx3C/xWM1R6Z2wzlAMMQWCBiEesDQuJxfO7mCVDzh/y8M85h3+eD4DIUPPm57xHnFF8XvP5Y1iWZQkAAAAAAACYcz6vGwAAAAAAAADcQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFQBAGAAAAAACAhUAQBgAAAAAAgIVAEAYAAAAAAICFcNbrBoBRfP7zn9fjx4915swZ/at/9a+8bgcAAAAAANj0v//3/9aTJ0907tw5ffLJJ65e27Asy3L1isAE+Hw+8a0LAAAAAMDsMgxDrVbL1WsyIwwzqTsEC4VCHnYynHq9rlarJZ/Pp+XlZa/bsY2+3TervdO3u+jbXfTtLvp2F327i77dRd/uom93zWrftVpNkjyZ4EIQhpnk8/k6P+wHBwdet2NbOBxWrVbT8vIyfbtgVvuWZrd3+nYXfbuLvt1F3+6ib3fRt7vo21307a5Z7fvMmTOde3q3sVk+AAAAAAAAFoKnM8Kee+45Ly/fwzAM/a//9b+8bgMAAAAAAAAO8TQI29/fl2EYnm563r6+YRie9QAAAAAAAADnLfzSSE4eBAAAAAAAWAyeb5ZvWZaSyaQikYgn1zdNU8Vi0ZNrAwAAAAAAwD2eB2HFYlHf+MY3PO0hmUzqK1/5iqc9AAAAAAAAwFmeL42Mx+Net6BLly553QIAAAAAAAAc5vmMsKWlJa9bwAz6whe+oEePHukLX/iC160shEwmo2azKb/f73UrQ5nVvqXZ7X1W+55Vs/r3Td/umtW+Z9Ws/n3Tt7tmte9ZNat/3/Ttrlnte1Z5eU9vWB7uFn/mzBk1Gg3Pv9EePXqkpaUlPXnyxNM+YF84HFatVlMoFNLBwYHX7dg2q30D846fTWD68HMJTB9+LoHpM6s/l1727enSyGk6sXGaegEAAAAAAMDkeRqE5XI5z2eDSdL58+eVy+W8bgMAAAAAAAAO8jQIe/PNN1293kcffXTq59zuBQAAAAAAAO7y/NRIN62vr3vdAgAAAAAAADyyUEHY7u6u1y0AAAAAAADAI2e9bmAUH330kUzTtF1frVZVKBSGeg0AAAAAAADmy0wEYR999JFyuZzK5bKq1epIY1iWJcMwJtwZvFav1xUOh/vWZDIZZTIZlzoCAAAAAGBx5fN55fP5vjX1et2lbo6b+iDs2rVr2traknQYZgHdWq2WarVa35pms+lSN4NlMhk1m82pOC0VwG/xswlMH34ugenDzyUwfabx57LZbA68T/eSYU1xuvSjH/1IqVRKkjqzucZp1zAMPXnyZCK9wVvhcFi1Wk0+n0/Ly8t9a5kRBgAAAACAO+zOCGu1WgqFQjo4OHCps0NTHYStrq6qUqnIMAxZlqVIJKJoNKpIJCJJevbZZweO8atf/Uqmaer27dt69OgRQdicaAdhXvzQAAAAAACA0Xl5Tz/VSyPbIZgk7ezsaG1tbeSxNjY29Nxzz02qNQAAAAAAAMwYn9cN9BMIBCQdhljjhGCSFIlEdOHChQl0BQAAAAAAgFk01UFYewnkpUuXJjJeoVCYyDgAAAAAAACYPVMdhLVngT18+HCi43mtWq0qkUhMbLxKpaJSqaStrS2Vy+WJjTtsn1tbW0okEgoGgzIMQysrK0qlUhPtCQAAAAAAYFRTHYR95zvfkWVZqlQqExnvBz/4wUTGOY1pmjIMY+BjZWWlM9tt1OtsbW0pFovJMAytra3p3r17ikajWl1ddb3PcrmsYDCobDYrSdre3tb+/r5yuZwqlYoSiYQSiYRM0xz5awYAAAAAABjXVAdh58+f1zvvvKNbt27pX/7lX8Yez+mlkcVi0XZtOzQahmmaWl9f7wmddnZ21Gg0lMvlFI/HO/uqudVnuVzuhFzpdFo7OzuKx+OKRCJKJpPa399XNBpVuVxWLBYjDAMAAAAAAJ4xLMuyvG5ikEQioTNnzuhnP/vZyGM8ePBAv//7v68nT55MsLNewWDQVtATj8e1s7Mz1NjFYlHZbLYzfqFQUDqdHqHLyfVpmqYuXLgg0zQViUS0v79/Yl21WtXKyoqtMe3y8qhVAAAAAAAwOi/v6c+6erUR7ezsaHV1Vc8995yy2ayWlpZsve7hw4cyTVP7+/u6ffu2oz0Wi0WZpqmNjY2B+2rZWb7YLZvNamtrS9LhAQI7OzsjL62cZJ+pVKoTqPWbOdaeHVYqlVQul1UsFkcO8QAAAAAAAEY1EzPCPvjgA21sbKhSqcgwjJHGsCxLhmE4NiOsPePptFlRo0qlUiqVSpKkQCCgBw8e2Fr+eJpJ9dk9y0uSGo1G375KpZJSqZSkw6+j0WiMdf1ZmRFmWVJrwE+Yz5BG/LYGAAAAAGDmMCOsj29+85udPa0Mw9Aoud2o4ZldpVJJ1Wp14nuQJRKJnhMX9/b2xgrBJtlnLpfrfGxnb7JkMtn52DRNlUqlnj+bVy1LOvtW/5rP3pbOEIQBAAAAAOC4qd4s/8aNGyoUCp3wa9TJa05Petvc3FQgENDly5cnNmaxWOwJwQqFwlgnTUqT7bN7w/1oNGrrNd3937p1a+weAAAAAAAAhjHVM8LaM5faM8Hi8bgSiYQCgcDQ+4Rdv35dH3300cR7rFQqqlQqkg43oY9EIp0+R53xVK1Wtb6+3nkejUbH3lNrkn22x2m7dOmSrddFo1FVq1VJ6iz3BAAAAAAAcMtUB2HtPcECgYB2d3d14cKFkce6evWqnn322Ql2d+joJvHValXFYrEzYyqZTOratWu2Z01J6uyl1Xbt2rWp6rN7ppok2zPVjtZVKpWh/l4AAAAAAADGMdVLI9v7Tl27dm2sEKw91sWLFyfQ1W9Vq9VjodBRpVJJsVisZ4bXoDG7Z1wFAoGx99KadJ/37t3reW5337KjQeTu7q6t1wEAAAAAAEzCVAdhsVhMkv0ZR4PcuHFjIuO0RSIRFQoF5XI5pdPpvn0Wi8XO19PP0Y3s4/G4pN+euriysiLDMBQMBhWLxbS1tSXTNF3ts728sXt8O47WTfqEzVn17n2vO8CwBv3MAQAAAACm01QvjVxfX9fdu3ePBS+jmvSMMEnH9u4yTVPFYlGbm5vHbpYrlYoSiYR2dnZOHa97E3pJWlpaUiwWO7Yvl2manX2/stmstre3+84cm2Sfk3o/FiFMsBNyvVaSzviklyf/7YkJqlQqunXrlorFouLxuLa3t71uCQAAAAAwpKmeEZZMJnXx4sWJnTD44x//eCLj9BMIBLSxsaFGo6Ht7e1jywbL5bK2trZOfG21Wj0WDt2+fVvr6+tqNBqyLEuWZWl/f/9YsJVKpYbagH6cPkcNsI5e4+HDhyONMysePZbe+Im92tffk5qPHW1nrrUD4UQioZWVFQWDwZ6Zk4lEQtls9tQlwqlU6thMSNM0OzMxh5mBCQAAAACYXoZlWZbXTfRTrVb1+7//+/rRj36kP//zPx9rrEuXLh3b38pppmlqbW3t2L5fjUbjWG25XFYikeipe/Dgwal7cBWLxZ49vQbVT6pPwzB6ntv
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "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_xEndT\"]),\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}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAACEDElEQVR4nOz9T2yj6X0v+P6oqhht2FFR6oMBBAlImkof+OAuMkVVzSK7sciTTXzzp8XqQXxtb1xS2sjMBTTHkstezJlFtywl1uZc5LRYPovrIMBUi/YkSC6QabH77GYxXWL77IKZiN0BShAwgEssOYaNuFu8i4oYSSVRlMT/7+cDECBLPz7vj+JbquJXz/O8qXq9Xg8AAAAAGHIjvW4AAAAAALpBEAYAAABAIgjCAAAAAEgEQRgAAAAAiSAIAwAAACARBGEAAAAAJIIgDAAAAIBEEIQBAAAAkAiCMAAAAAASQRAGAAAAQCIIwgAAAABIBEEYAAAAAIkgCAMAAAAgEQRhAAAAACSCIAwAAACARBCEAQAAAJAIgjAAAAAAEkEQBgAAAEAiCMIAAAAASARBGAAAAACJIAgDAAAAIBEEYQAAAAAkgiAMAAAAgEQQhAEAAACQCIIwAAAAABJBEAYAAABAIgjCAAAAAEgEQRgAAAAAiXCz1w3AVXzuc5+Ln//851Gv12NkpHme+6u/+qvx+c9/vkudAQAAQHL94z/+Y/z0pz9tWnN4eBipVCo++9nPxs9+9rMudfZcql6v17t6RGiDGzduxOHhYa/bAAAAAK5oZGQkPv30064e04wwBtLxIKwTM8L29vbi8PAwRkZGYmJi4sp9OlbnjzWMr8mxBuc4jjVYxxrG1+RYg3WsYXxNjjU4x3GswTrWML4mxxqc41z3WK3OCIt4/tm+6+owgCYnJ+sRUZ+cnBzI8R1r8I7jWIN1rGF8TY41OMdxLMfq9XEca7CONYyvybEG5ziONVjHGqbX1M3XcprN8gEAAABIhJ4ujXz11Vd7efgTUqlU/J//5//Z6zYAAAAA6JCeBmE7OzuRSqWi3sP9+o+On0qletYDAAAAAJ2X+KWRvQzhAAAAAOienl81sl6vx9zcXGQymZ4cv1arRbFY7MmxAQAAAOiengdhxWIxvv71r/e0h7m5ufjt3/7tnvYAAAAAQGf1fGlkLpfrdQtx9+7dXrcAAAAAQIf1fEbY+Ph4r1uAFywuLsbBwUGMjo46Vp8b1u/fsB6rW4b1+zesx+qWYf3+Deuxuqlbr2tY3yvnxWAcZ5iP1U3Oi8E5Vjc5LwZLqt7D3eJv3LgR+/v7Pf/GPnv2LMbHx+PTTz/taR+0bmpqKnZ3d2NycjKePHnS63Y4g/eIszgvOIvzgrM4LziL84KzOC84i/Oiv/Xy/enp0sh+umJjP/UCAAAAQPv1NAhbXV3t+WywiIhbt27F6upqr9sAAAAAoIN6GoR985vf7OrxPv7443O/1u1eAAAAAOiunl81spsWFhZ63QIAAAAAPZKoIOzx48e9bgEAAACAHrnZ6wau4uOPP45ardZyfbVajY2NjUs9BwAAAIDhMhBB2Mcffxyrq6tRLpejWq1eaYx6vR6pVKrNnQEAAAAwKPo+CHvw4EGsra1FxPMwCxgMi4uLcXBw0BdXhqV/OC84i/OCszgvOIvzgrM4LziL84LzpOp9nC798Ic/jEKhEBHRmM11nXZTqVR8+umnbemN3pqamord3d2YnJyMJ0+e9LodAAAAoEW9/Ezf1zPCVlZWIuJ5gFWv1yOTyUQ2m41MJhMRES+//PKFY/zkJz+JWq0W77zzTjx79qyj/QIAAADQv/o6CKtUKo2ZYFtbWzE7O3vlsZaWluLVV19tV2sAAAAADJiRXjfQTDqdjojnIdZ1QrCIiEwmE6+88kobugIAAABgEPV1EHa0BPLu3bttGW9jY6Mt4wAAAAAwePo6CDuaBfb06dO2jgcAAABA8vR1EPbtb3876vV6VCqVtoz3/e9/vy3jAAAAADB4+joIu3XrVnz3u9+NR48exU9/+tNrj2dpJAAAAEBy9XUQFvF8o/xsNhuFQuFa43z00Udtm1kGAAAAwOC52esGWrG1tRV37tyJV199NZaXl2N8fLyl5z19+jRqtVrs7OzEO++80+EuAQAAADqnXo84rDevGUlFpFLd6WcQDUQQ9v7770dExM7OTiwsLFxpjHq9HilnAgAAADCgDusRN7/TvOaTNyNuiD/O1fdB2BtvvBHFYjEiIlKpVNTrF0SfZxCAAQAAANDXQdjDhw8bG9xfNQSLiCs/j/63t7cXU1NTTWsWFxdjcXGxSx0BAABAcq2vr8f6+nrTmr29vS5186K+DsJOh2C5XC7y+Xyk0+lL7xP29ttvx8cff9zBbumFw8PD2N3dbVpzcHDQpW4AAAAg2Q4ODi78nN5LfR2EVSqVSKVSkU6n4/Hjx/HKK69ceaz79+/Hyy+/3Mbu6AcjIyMxMTHRtGZ0dLRL3QAAAECyjY6OxuTkZNOavb29ODw87FJHJ/V1EJZOp+PZs2fx4MGDa4VgR2Pdvn27TZ3RLyYmJuLJkye9bgMAAACI1rYnmpqa6tmssZGeHLVFMzMzERGRyWTaMt7Dhw/bMg6QbLVardctAAAAcAV9HYQtLCxEvV6ParXalvHMCAOuqlKpxPLycoyNjcX9+/d73Q4AAABX0NdLI+fm5uL27dvx6NGj+Hf/7t9de7wf/ehH8Qd/8Adt6AzoJ5VKJR49ehSVSiWq1WrjIhnpdDoymUyMj49HNpuNfD4fuVzuhecXCoWoVquxvb3d+LNarRblcjkePXoU5XLZLDAAAIAh0NdBWETE5uZm/MZv/Eb8r//r/xq///u/f62xVlZWBGH0pWKxGAsLC5d+XjabPRHeXFUqlbr0c9LpdOzv75/5tUKhEKVS6dJjbm1tnRlUnaVWq8XKykqsra01ralUKhERUS6XG7XZbDZyuVxMT0/H1tZWlEqlyGazJ557//79qFarjecDAAAw+Pp6aWTE8/3B3n777fj6179+rXGePXvmAy19a35+Pvb392N7ezvm5uaa1uZyudja2mrUt0O9Xo+dnZ3Y2NiIdDp9bl0mk4mNjY3Y2dk5NwSLeB5g7+/vtxRsLS0txfb2duzv77ccgh0tUTwdgs3NzcXm5mZjvKPXtbW1FfPz843XVqlUYm1tLRYWFhqB3ekZX0fj1Ov1tu1TCAAAQG+l6vV6vddNnOdHP/pR4/7S0lKkUqkLQ4Kz1Gq1eOedd6JWq8Wnn37azhbpkaMrTExOTg7lVSMXFhaiWCy+8OdHQU8nVavVmJ6ePvNr+/v7TYOy88zMzJwZRG9sbMT8/Pylesvn8y/sGzg/Px8bGxstjbG2thbLy8sv/Hkmk4mdnZ0zn3P6/ejG+wAAAHDap4cRN7/TvOaTNyNu9Pm0p15+pu/rpZFvvfVWfPjhh43H9Xq96TKoZur1+pWWf0EvnBeEXWX55GVlMpnIZrMvBFdzc3NXCsEinvd9Vu+XCcHK5XLk8/kTf5ZOp+O99957YVljM0tLSzE3NxczMzMnZoE9ffr03Odc9XUDAADQX/o6I7x3717U6/U4mrQmyCIpzluK160lemcd5zrHPuu5lwmXKpVKW0Kw4/189NFHJ3qwGT4AAMDw6+sg7GgGSSqVagRiV73BIDkvJBofH+9uI32gWq3G7OzsC3++ubl5pRDsSDqdbtsea/2kVCq9sHSUf1Eul3vdAgAA0EN9vTTy1q1bkc1m48MPP4zV1dXIZrNXCgKq1Wq89dZb8eMf/7j9TcIQOuvv2csvv9zW8Vr9u1woFF6YrTU/P9/yxvrNZDKZWFpauvKS626qVCqxsbER5XI5qtVqpNPpGB8fj3Q6HZlMJl5//fWIeH61y9XV1abLTsvlchQKhRgfH4+NjY22fC/7VbVajXK5HJubm40QrJ9+OZKk9wIAAPpBXwdhERGvv/56TE9Pxze/+c0rj3H79u147bXXEjmbBgZZuVw+c5P91dXVth3jwYMHjSCsVqv13X5gtVot7t+/37i6ZSaTaVw0pFqtRqVSiUql0vj60XO
      "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": 79,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intercept= 5303.132333467949\n",
      "coef= {'ideal_state_770_tx': -13.441950438394027, 'dSlope_xEndT': -4.369407701653064, 'dSlope_xEndT_abs': 230.37763264120008, 'x_EndT_abs': 0.026704166547390083, 'ideal_state_770_ty': -14.30702118941663, 'ideal_state_770_tx^2': 77928.66766333766, 'ideal_state_770_tx dSlope_xEndT': 69687.8060631645, 'ideal_state_770_tx dSlope_xEndT_abs': -86.35194840522867, 'ideal_state_770_tx x_EndT_abs': 0.00970580581706756, 'ideal_state_770_tx ideal_state_770_ty': -180.61816411073835, 'dSlope_xEndT^2': 7579.721527263535, 'dSlope_xEndT dSlope_xEndT_abs': -49.35311026671389, 'dSlope_xEndT x_EndT_abs': 0.015857700483615466, 'dSlope_xEndT ideal_state_770_ty': -35.74149120660231, 'dSlope_xEndT_abs^2': 7579.723402132595, 'dSlope_xEndT_abs x_EndT_abs': -0.22195273681292957, 'dSlope_xEndT_abs ideal_state_770_ty': 1.9491174288661794, 'x_EndT_abs^2': -0.0009034030507007965, 'x_EndT_abs ideal_state_770_ty': 0.0034286491846703187, 'ideal_state_770_ty^2': -2489.6827607187365, 'ideal_state_770_tx^3': 2357.797471515544, 'ideal_state_770_tx^2 dSlope_xEndT': 1960.401925113787, 'ideal_state_770_tx^2 dSlope_xEndT_abs': -20252.830218161023, 'ideal_state_770_tx^2 x_EndT_abs': -13.500325039413802, 'ideal_state_770_tx^2 ideal_state_770_ty': -766.130872351457, 'ideal_state_770_tx dSlope_xEndT dSlope_xEndT_abs': -19474.231043469852, 'ideal_state_770_tx dSlope_xEndT x_EndT_abs': -11.598943144027315, 'ideal_state_770_tx dSlope_xEndT ideal_state_770_ty': -626.8199943500564, 'ideal_state_770_tx dSlope_xEndT_abs^2': 573.5326109010182, 'ideal_state_770_tx dSlope_xEndT_abs x_EndT_abs': 0.001095488438740151, 'ideal_state_770_tx dSlope_xEndT_abs ideal_state_770_ty': 330.5306955155975, 'ideal_state_770_tx x_EndT_abs^2': -3.111920124211576e-05, 'ideal_state_770_tx x_EndT_abs ideal_state_770_ty': 0.06268520395759442, 'ideal_state_770_tx ideal_state_770_ty^2': 95.20495106050294, 'dSlope_xEndT^3': 42.93032905404232, 'dSlope_xEndT^2 dSlope_xEndT_abs': -4365.747480960424, 'dSlope_xEndT^2 x_EndT_abs': -1.2192562284269424, 'dSlope_xEndT^2 ideal_state_770_ty': -134.19763413582538, 'dSlope_xEndT dSlope_xEndT_abs^2': 42.930335617557944, 'dSlope_xEndT dSlope_xEndT_abs x_EndT_abs': -0.023890507178523725, 'dSlope_xEndT dSlope_xEndT_abs ideal_state_770_ty': 106.84806568209468, 'dSlope_xEndT x_EndT_abs ideal_state_770_ty': -0.009451904933143932, 'dSlope_xEndT ideal_state_770_ty^2': -8.505381537855747, 'dSlope_xEndT_abs^2 x_EndT_abs': -1.2193799254865163, 'dSlope_xEndT_abs x_EndT_abs^2': 0.0002552425918335738, 'dSlope_xEndT_abs x_EndT_abs ideal_state_770_ty': 0.01563137413267504, 'dSlope_xEndT_abs ideal_state_770_ty^2': 2033.0314554223403, 'x_EndT_abs^3': 1.5428826370868625e-07, 'x_EndT_abs^2 ideal_state_770_ty': 6.223748596200707e-06, 'x_EndT_abs ideal_state_770_ty^2': 0.5895518247828866}\n",
      "r2 score= 0.8237497474212265\n",
      "RMSE = 34.7702514276615\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",
    "    # \"dSlope_yEndT\",\n",
    "    # \"dSlope_yEndT_abs\",\n",
    "    \"ideal_state_770_ty\",\n",
    "    # \"ideal_state_9410_ty\",\n",
    "    # \"ideal_state_770_x\",\n",
    "    # \"x_EndVelo_abs\",\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=False)\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\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_xEndT_abs\",  # keep\n",
    "    \"x_EndT_abs\",  # keep\n",
    "    \"ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx^2\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT_abs\",\n",
    "    \"ideal_state_770_tx x_EndT_abs\",\n",
    "    \"ideal_state_770_tx ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT^2\",  # keep\n",
    "    \"dSlope_xEndT dSlope_xEndT_abs\",\n",
    "    \"dSlope_xEndT x_EndT_abs\",\n",
    "    \"dSlope_xEndT ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT_abs^2\",\n",
    "    \"dSlope_xEndT_abs x_EndT_abs\",\n",
    "    \"dSlope_xEndT_abs ideal_state_770_ty\",\n",
    "    \"x_EndT_abs^2\",\n",
    "    \"x_EndT_abs ideal_state_770_ty\",\n",
    "    \"ideal_state_770_ty^2\",  # keep\n",
    "    \"ideal_state_770_tx^3\",\n",
    "    \"ideal_state_770_tx^2 dSlope_xEndT\",\n",
    "    \"ideal_state_770_tx^2 dSlope_xEndT_abs\",\n",
    "    \"ideal_state_770_tx^2 x_EndT_abs\",\n",
    "    \"ideal_state_770_tx^2 ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT dSlope_xEndT_abs\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT x_EndT_abs\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT_abs^2\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT_abs x_EndT_abs\",\n",
    "    \"ideal_state_770_tx dSlope_xEndT_abs ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx x_EndT_abs^2\",\n",
    "    \"ideal_state_770_tx x_EndT_abs ideal_state_770_ty\",\n",
    "    \"ideal_state_770_tx ideal_state_770_ty^2\",\n",
    "    \"dSlope_xEndT^3\",\n",
    "    \"dSlope_xEndT^2 dSlope_xEndT_abs\",\n",
    "    \"dSlope_xEndT^2 x_EndT_abs\",\n",
    "    \"dSlope_xEndT^2 ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT dSlope_xEndT_abs^2\",\n",
    "    \"dSlope_xEndT dSlope_xEndT_abs x_EndT_abs\",\n",
    "    \"dSlope_xEndT dSlope_xEndT_abs ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT x_EndT_abs ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT ideal_state_770_ty^2\",\n",
    "    \"dSlope_xEndT_abs^2 x_EndT_abs\",\n",
    "    \"dSlope_xEndT_abs x_EndT_abs^2\",\n",
    "    \"dSlope_xEndT_abs x_EndT_abs ideal_state_770_ty\",\n",
    "    \"dSlope_xEndT_abs ideal_state_770_ty^2\",\n",
    "    \"x_EndT_abs^3\",\n",
    "    \"x_EndT_abs^2 ideal_state_770_ty\",\n",
    "    \"x_EndT_abs ideal_state_770_ty^2\",\n",
    "]\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_tx^2\",\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",
    "\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)\n",
    "# lin_reg = LassoCV(max_iter=2000)\n",
    "# lin_reg = ElasticNet(alpha=0.1)\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": 77,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['ideal_state_770_tx x_EndT_abs^2',\n",
       " 'dSlope_xEndT x_EndT_abs^2',\n",
       " 'x_EndT_abs^3',\n",
       " 'x_EndT_abs^2 ideal_state_770_ty']"
      ]
     },
     "execution_count": 77,
     "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": [],
   "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": 60,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMEAAAOWCAYAAADm+IXMAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjguMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy81sbWrAAAACXBIWXMAAA9hAAAPYQGoP6dpAADZ2UlEQVR4nOz9T3Bb2X3nf39AeVJJTUJeUtlJVW5etH6VnUWAyuK3swh0e5GUxy2A6nJqkoVFQPLzq1nQFkB2FpNZRBKgbm2mHlsA1RvniWtEUOq4koUtgOr1zyKu1LtU3Likq6SliEsmqXElad5nwQEM8A+IPxfABfl+VaFMkOee+yUBqsxPf885Add1XQEAAAAAAACn2NiwCwAAAAAAAAD6jRAMAAAAAAAApx4hGAAAAAAAAE49QjAAAAAAAACceoRgAAAAAAAAOPUIwQAAAAAAAHDqEYIBAAAAAADg1CMEAwAAAAAAwKlHCAYAAAAAAIBTjxAMAAAAAAAApx4hGAAAAAAAAE49QjAAAAAAAACceoRgAAAAAAAAOPUIwQAAAAAAAHDqEYIBAAAAAADg1CMEAwAAAAAAwKlHCAYAAAAAAIBTjxAMAAAAAAAApx4hGAAAAAAAAE49QjAAAAAAAACceoRgAAAAAAAAOPUIwQAAAAAAAHDqEYIBAAAAAADg1CMEAwAAAAAAwKlHCAYAAAAAAIBTjxAMAAAAAAAApx4hGAAAAAAAAE49QjAAkiTHcYZdAgAAAAAAfUMIBpxhlmUpnU5rcnJSCwsLwy4HAAAAAIC++dqwCwBwNMuy9PjxY1mWJdu2tb29LcdxZBiGTNPU1NSUQqGQotGoIpHIoevj8bhs21a5XK5/znEclUolPX78WKVSie4vAAAAAMCZEXBd1x12ETid8vm8kslkx9eFQqGm4KZbgUCg42sMw1C1Wj3ya/F4XGtrax3PWSwWjwypjuI4ju7evatsNtvxfUKhkCKRiILBoIrFotbW1g79LGvBmGVZh66PxWIqFAod3xcAAAAAgFHAckj0TSKRULVaVblcViwWazk2EomoWCzWx3vBdV1VKhXlcjkZhnHsONM0lcvlVKlUjg3AJKlQKKharbYVaqVSKZXLZVWr1bYDsNqyxIMBWC2cqs1X+76KxaISiUT9e7MsS9lsVslksh7WHez0qs3juq5M02yrLgAAAAAATgM6wTAwyWRS+Xz+0OcH0YFk27aCweCRX6tWqy1DsuOEw+EjO6pyuZwSiURHtUWjUdm23fT5RCKhXC7X1hzZbFbpdPrQ503TVKVSOfKag68HnWAAAAAAgNOMTjAMzHFLI7tZMtkp0zQVCoUOfT4Wi3UVgEnH191JAFYqlRQMBpsCMMMwVC6X2w7ApP3Os0qlcuh72d7ePvaabr9vAAAAAABGESEYBua45XeDWpZ31H16ufdR13YSLFmWpWg0euj69fX1IwO7durZ3NxsqoGN7wEAAAAA2EcIhoE5LiCampoabCE+YNu25ubmDn2+UCh0FYDV1LrITpu1tbVDy0XxW6VSadglAAAAAIDvEYLhzDgqbDt//ryn87Ub6MXj8UNdWolEou1N9FsxTVOpVKrneQbBsiwlk0kFg0EFAgFNTk4qGAwqHA7XT+NcW1vTwsLCiUFPqVSqX3/aQyHbtpXP5xWNRhUIBA51FA7bWXotAAAAAIyOrw27AOCsKZVKR26on8lkPLvH8vJy/ZRJx3F8t/+X4zhaWFion2Jpmmb9BFHbtmVZlizLqn+9dk0rtWDRcRwlk8ljDwQYZbZtKx6Py7ZtXy91PQuvBQAAAIDRQycYMGBHbajfywb9RzEMo6MN+get1uUl7Z+mWalUVCgUVCgUVC6XValUDnXFnRSknIVltaZpqlwuq1qt+rrb7yy8FgAAAABGDyEYMEClUunIva2Wl5c9v1c8HpfU+oTIVkqlkuLxeNNSxWg0qnw+31Nd2Wy2vkQulUodGdaZpqlisdjUHXfSnmC5XE6maco0TRUKhZ5qHAV+WwLZ6Ky9FgAAAABGA8shgQE6LhDoZTP840QiEWUymY67chzH0dzc3KElm47jqFQqqVQqKZPJdL2J/927d+sfH9UV1yiVSunt27fKZrMnhnmRSORMLbsbVrdVMplULpdrOeasvRYAAAAARgOdYMAAHbVJuBeb4R8nlUp1tMzStm1NT08fuWfZwXHhcLjjTc+72csqk8koFApxOqQP5PP5njsBAQAAAGBYCMGAAbFt+8ggpx9dYN1wHKe+oXksFlMul1O5XFaxWDx2/6loNNpROHUwAGs3RFteXvb1RvBngW3bJ3buAQAAAICfsRwSGJDjwqJgMDjgSo5WC6SKxeKh7rRIJKJkMnlk6BWPx1Uul9u6h2maTc/T6bTm5+dP7FarnRzpx5MuzwLHcXy9BxkAAAAAtINOMGBAjgvB/HSS3lEBWE1ts/qDLMtqu6PLMIymEMtxHIXD4ba6yVzXJQAbgtoecSxHBQAAADDqCMFwpqXTaQUCga4e4XC4o3sdt5zPL8GOYRgn7k9mmmbTiY01R33uOAfH2ratYDDoyV5TlmUpmUye2F3nOI7y+bzC4XDTfWtL/iYnJxUIBBQMBpXNZg9db9u20ul0/eTMYDCoZDJ55Gtcu8/B989B0Wi0Pl/jw8tloPl8vuk+k5OTCofDR36P0v7P86g94hrrO+5n3e5rUeM4jrLZrKLRaP3nX6svnU63FcLZtq1sNqtgMNhUc+NrEAwG68t+AQAAAJwxLjBAkg49qtXqQO6dSCSOvL+XD9M0O75/sVgcyPd/UCqVaqojEom0fa1hGIe+j05EIpEjfxahUMgtl8sdzVUul91EItFUk2EYh8ZVq1U3k8m4oVCo6Z65XM513cM/j8ZHIpGoz9NqnGmax76fD77+x8lkMm39fpTL5bZ//oVCof5zKRQKTXOYpnls7dVq1S2Xy4dqKpfLTY/G+dp5LQ7K5XL18ZlMxi2Xy26lUnELhULT6xWLxQ5dW61W3VQqVf8+GmusVquHXu/GuiqVyom1AQAAADg96ATDmZZIJFQul7t65HK5ju51XMeXXzpSOulIO6pj7KQTJRsVi8X6Pl8H5wiHw2136ti2rcePH7fdbRQKhQ5t7l4ulxUOh2VZlgqFgiqViqrVatNhAPl8XpZlKRqNqlQqqVgsqlqtqlqtqlgs1n92tm3r7t27R947Ho+3VeNRP5delEql+r2Xl5eb5g+FQioUCpL2a19YWGi61jAMhUKhQ4c31D7X+LVOX4uadDqtZDIpwzBUqVSUSqUUCoVkmqZisZjK5XK95rW1NQWDwab3xvb2tqLR6KH35MbGhqanp+vLeCuVinK5XP21chyHjf4BAACAs2bYKRzOFvmsEyyTyXQ938FOHJ3QCXawm6b2qHUiDdrBjqajumyOU+ss6vX7OO5n0s2cjXOd1H10sDvouG68gx1rjR1hjRp/Hsfdu1gsttW5Va1WPe0Ei8ViJ77fT5qn3dpr2n0tGjvATnqtGzu9jutabKzRMIwjX9eD791B/fsDAAAAYPjoBAMGxO+dYJ04eMqj1N33kUqlVK1Wj+1+qp1I2c7cR9V0nMbDCBKJxLF7oR08EfG47r/G+h3H8e1retx7sPHzXmyA3+5rkU6n6/dPJBJtjZX2u9uOOoyh8ftYWVk58nU9+F7b2Nhoq1YAAAAAo48QDBiQ44KBSqUy4Ep6d1SY8vbt267nKhQKKpfLh5bdSfuBx9zcnKfBUrtLPxtfs06Wi/rpJMWVlRVlMhkVCoUTgyZpcKFsPp+v36ud0CyRSDS9Bp0cxnBQ4/389FoBAAAA6C9CMGBAjus2Oqqjxe+OCi3Onz/f05yhUOjYvdYsyzq0X9UgdHty5/b2treF9MAwDKVSqWP3YOtkLzcvFYvF+sftdo41/g6N4u8NAAAAgOEiBAMG6KhOJ9u2fbt8rhOdLEdsJZFIqFKpHAqg1tbWfB18dBuYDZpt28rn8yqVSkduej/IOjp18D12Gn5vAAAAAAwOIRgwQNevXz/y834Od45zMPTxKgSrzVUulw/do9MTOfFbtdMtc7lcy33
      "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=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=50, range=[5150, 5700], color=\"#2A9D8F\", alpha=0.8, align=\"left\")\n",
    "ax.set_xlabel(r\"z$_{Mag}$ [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": [
    "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(data,\n",
    "                                                    target,\n",
    "                                                    test_size=0.1,\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",
    "#     \"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,\n",
    "         bins=30,\n",
    "         range=[-200, 200],\n",
    "         color=\"#2A9D8F\",\n",
    "         alpha=0.8,\n",
    "         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
}