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

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 18,
   "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",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "input_tree = uproot.open({\"/work/guenther/reco_tuner/data/param_data_selected.root\": \"Selected\"})\n",
    "array = input_tree.arrays()\n",
    "array[\"dSlope_fringe\"] = array[\"tx_ref\"] - array[\"tx\"]\n",
    "array[\"z_mag_x_fringe\"] = (array[\"x\"] - array[\"x_ref\"] - array[\"tx\"] * array[\"z\"] + array[\"tx_ref\"] * array[\"z_ref\"] ) / array[\"dSlope_fringe\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAACmaElEQVR4nOz9z28bV77n/79KncUAg5ZK8sUFBAnoqBjP4pvFWJTc+7HIdG8nJuWs0xFpzypQt0Urf4Bt0h19s5qYlJP1tSg7274m5exvpLKzyKrDkgJIIDADU2U2BphNVJ+FhmxRPymyxCpazwdAtCidOudNUQ7AV7/PKcPzPE8AAAAAAADAO24g6AIAAAAAAACAXiAIAwAAAAAAwKVAEAYAAAAAAIBLgSAMAAAAAAAAlwJBGAAAAAAAAC4FgjAAAAAAAABcCgRhAAAAAAAAuBQIwgAAAAAAAHApEIQBAAAAAADgUiAIAwAAAAAAwKVAEAYAAAAAAIBLgSAMAAAAAAAAlwJBGAAAAAAAAC4FgjAAAAAAAABcCgRhAAAAAAAAuBQIwgAAAAAAAHApEIQBAAAAAADgUiAIAwAAAAAAwKVAEAYAAAAAAIBLgSAMAAAAAAAAlwJBGAAAAAAAAC4FgjAAAAAAAABcCgRhAAAAAAAAuBQIwgAAAAAAAHApEIQBAAAAAADgUiAIAwAAAAAAwKVAEAYAAAAAAIBLgSAMAAAAAAAAl8J7QRcA/Of//J/1f//v/9VvfvMb/eu//mvQ5QAAAAAAgAv0v/7X/9Kvv/6q//Sf/pP+z//5Pz1d2/A8z+vpisAhv/nNb7S3txd0GQAAAAAAoIcGBgb066+/9nRNOsIQuEYQNjAwoNHR0XNdW61WO762W6zd+7WDXp+1WZu1WftdXJ+1WZu13921g16ftVmbtVn7rOt/85vfXEB1p+u7IGxra0u2batWq8l1Xb1580aSdOXKFZmmKcuyND09rcHBwYArRbv+9V//VTs7OxodHdX29va5rh0fH+/42m6xdu/XDnp91mZt1mbtd3F91mZt1n531w56fdZmbdZm7bOuD+J4pL4Iwp4/f65SqaSVlRW5rtvWNZZlKRaLKZlM6saNGxdbIAAAAAAAAEIv1EHYkydPlM1m5TiOJOngcWaGYRx7TWNMpVKR4zgqFAqyLEv37t3Tn/70p4svGgAAAAAAAKE0EHQBx3n58qWuXr2qdDp9bAjWeH7c47gxjuMolUrpv/yX/6Iff/yxZ68DAAAAAAAA4RG6jrDFxUXlcrkjoVYsFlM0GtWVK1dkWZZM05QkjYyMSJJqtZokyXVdOY6jN2/eyHEclcvl5nbKn3/+WdFoVJlMRvfv3+/ZawIAAAAAAEDwQhOE1et1zczMaGNjQ9J+8BWPxxWLxTQ5OdnV3G/fvtXTp09VKBRk27ay2azK5bLK5TKH6gMAAAAAAFwSodga+fr1a01MTGhjY0OpVEqVSkUvXrzQ3bt3uw7BJGloaEipVErr6+v6+eefdePGDa2vr2t6elr/+Mc/fHgFAAAAAAAACLvAg7BXr14pGo1qYmJClUpFjx8/1sTExIWtZ1mWSqWSXrx4of/9v/+3pqamVK/XL2w9AAAAAAAAhEOgWyM3Nzc1MzOjRCKhlZWVnq4di8W0ubmpqakpzc7O6m9/+1tP14c/5ufnVa/XL90W1yBfd9C/88v82oNyWX/nl3XtIF3m3/llfu1Buay/88u6dpCCft2X9T0P+vcelMv6O7+sawepn1+34R0+lb6HPvjgA8ViMT1+/DioEuS6rqanp3Xnzh39+c9/DqyOy2x8fFw7OzsaGxvT9vZ20OW0rV/rRn/i7w29wt8aeom/N/QKf2voJf7e0Cv9/LcWZO2BdYTdu3dPlmUFGoJJkmmaWl9fl2VZmpmZ0bVr1wKtBwAAAAAAABcjsDPCrly5omKxGNTyLUzT1NramhzHCboUAAAAAAAAXJDAOsLu3r0b1NLHmpyc9OUOlQAAAAAAAAinwO8a2St37twJugQAAAAAAAAE6NIEYYVCQVtbW0GXAQAAAAAAgIBciiBsbW1Nnufp2bNnQZcCAAAAAACAgARyRtjt27e1ubl54evUajW5rts8BP/x48f685//fOHrojPValXj4+Onjpmfn9f8/HyPKgIAAAAAAOextLSkpaWlU8dUq9UeVXNUIEGYaZoqlUoyDONC1/E8r+W54zh6/fq1rl27dqHrojN7e3va2dk5dUy9Xu9RNWebn59XvV7X4OBg0KXgEuDvDb3C3xp6ib839Ap/a+gl/t7QK2H9W6vX62d+tg+S4R1Oi3rg1atXmpqaajsI8zyv69CsMUcmk9H9+/e7mgv+Gh8f187OjgYGBjQ6OnrqWDrCAAAAAAAIr3Y7wvb29jQ2Nqbt7e0eVbYvkCBMkj744ANtbm4e6do6jmEYbY1r5/pIJKK///3vHc8F/zWCsCD+AQAAAAAAgN4KMgcI7LD8mZkZSVK5XNbe3t6Jj2w2K8/zFI1GVSqVtLu7e+r4w4/19XVNTEwonU5rb2+PEAwAAAAAAOCSCuSMMEmanZ2Vbdu6cePGiWM2Nzd17949pdNpff311x2tE41G9eLFC33wwQcyDEP/83/+z05LBgAAAAAAQB8LtCNseHj41DHZbFaWZXUcgjVYlqW5uTnl83l9//33Xc0FAAAAAACA/hRYECZJL168OPXn6+vrSqfTvqw1PT0tz/OUz+d9mQ8AAAAAAAD9JdAg7Cy2bSsajfoy18jIiCSpVCr5Mh8AAAAAAAD6S6iDMNM0tbm56ctcjQDMdV1f5gMAAAAAAEB/CXUQNj09rcePH3c9z9u3b1UoFCTth2sAAAAAAAC4fEIdhKVSKW1sbOjLL7/seI56va5YLCZJMgxDlmX5VR4AAAAAAAD6SKiDsEQioYmJCS0sLOiTTz7R1tbWua5//vy5JiYmZNt283u3bt3yuUoAAAAAAAD0A8PzPC/oIk7jOI4++OADGYYhSYpGo0qn07IsS5Zl6f3335ckbW1tyXVdOY6jUqmklZUVua4rz/NkGIY8z9Pw8LDevHkT4KvBccbHx7Wzs6OxsTFtb28HXQ4AAAAAALhAQeYA7/V0tQ5YlqX19XVNT09L2r+TZDqdPvO6Rr7XCNAMw1CxWLy4QgEAAAAAABBqod4a2RCNRvXzzz9rcnJSnue19TAMo9kJ5nmeVlZWdOPGjaBfCgAAAAAAAALSF0GYtN8ZtrGxoZWVFUWj0SM/bwRfDY0ALJFIaHd3Vzdv3uxluQAAAAAAAAiZ0G+NPCyRSCiRSGhzc1PlclkbGxtyHEe1Wk2SNDIyIsuyFI/HFYvFNDQ0FHDFAAAAAAAACIPQH5aPdx+H5SNIruvKNM2gywD6wofffuXbXD99+rlvcwEAAKC/BJkD9M3WSADwi23bymQyGh4e1tzcXNDlAAAAAAB6pO+2RrZjcXFRt2/f1u9+97ugSwFwTrZt6+nTp7Jtu7ntudG1ZVmWRkZGFI1Gm9ufD0smk3IcRxsbG83vua6rcrmsp0+fqlwuy3XdHr4iAAAAAEBY9F1H2PPnz3Xnzh1dv35dz58/P3bM1NSUJicn9cknn+iXX37pcYW47AqFQvPmDed5TE1N+bJ+J2sPDw+fOF8ymexoznK53HbNrusqk8k0fw+5XE7lclmO4zRDK9d1Zdu2yuWycrmc4vF4c3wmk1GhUFAymdTq6uqR+efm5vTgwQOtrq4SggEAAADAJdY3Qdhf//pXXblyRclkUoVCQbZtNw/IPyyRSMhxHP3973+XZVn67rvvelwtLrNUKqXd3V1tbGwokUicOjYWi6lUKjXH+8HzPFUqFeXz+VPPvrIsS/l8XpVKRbu7uyeOKxaL2t3dValUOrYD66CFhQVtbGxod3f3zLENjS2KuVyu5fuJRELFYrE5X+N1lUolpVKp5muzbVu5XE7pdLoZgh0OuxrzeJ4ny7LaqgsAAAAA8O4JfRD29u1bXb16VZlMpvlhuJ3z/U3T1MbGhq5du6ZEIqFvvvmmB9UC+0zTVDQaVbFYVCqVOnZMIpFohkt+H9ZuWZZSqdS
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(array[\"z_mag_x_fringe\"], bins=50,\n",
    "        range=[5150,5300], color='#2A9D8F', density=True)\n",
    "plt.xlabel(r\"z$_{Mag}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAAC20UlEQVR4nOz9b2xbd57ne34oO4Fzk5JJpdGAVgKqTE4a6bsPuk3K86Duoy6RUwvcZJPpkPYg6STYizKZZGv+QDMh4wQX6MHCUcjuFjDbs2mTrr7AOB1gZLKqEyQD1ER09aM7wI6l4+pnwVTpOA1Iy8VgIx2rk403iXn2gZoskpLIQ/KQh3/eL4CIZH3545f645gf/X7f47Nt2xYAAAAAAAAw4Wa8bgAAAAAAAAAYBoIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMBYIwAAAAAAAATAWCMAAAAAAAAEwFgjAAAAAAAABMhdNeNwD04tFHH9X9+/d16tQp/fZv/7bX7QAAAAAAAIf++3//73rw4IHOnDmjL7/8cqiP7bNt2x7qIwIuOHXqlKrVqtdtAAAAAACAHs3MzOjBgwdDfUx2hGEsNQZhMzPtT/h+5zvf0WOPPTaMtjqqVCqqVquamZnR/Py81+0MzDQ8z2l4jhLPc5JMw3OUpuN5TsNzlKbjeU7Dc5Sm43lOw3OUpuN5TsNzlKbjeU7Dc5RG83l+8cUX+vu///u2NbXX86dOnRpGS00IwjCWfvu3f1u7u7uS1HFn2L/6V/9Kf/zHfzyErjpbXFzU7u6u5ufntbOz43U7AzMNz3ManqPE85wk0/Acpel4ntPwHKXpeJ7T8Byl6Xie0/Acpel4ntPwHKXpeJ7T8Byl0Xyef/zHf6x/+2//raNaL0YdEYRhrDlJvWdnZ4fUDQAAAAAA0212dlYLCwtta2o72bzgaRD2xBNPePnwTXw+n/7bf/tvXreBLo1S6g0AAAAAwLRbWVnRyspK25raTjYveBqEbW9vy+fzyct5/bXH9/l8nvUAAAAAAACAwWs/ZXwKcNFMAAAAAACA6eD5jDDbthWPxxUMBj15fMuyVCgUPHlsAAAAAAAADI/nQVihUNCPfvQjT3uIx+P64Q9/6GkPAAAAAAAAGCzPj0ZGo1GvW9CFCxe8bgEAAAAAAAAD5vmOsLm5Oa9bAOCylZUVHRwcaHZ21utWBmYanqM0Pc9zGkzL13Ianuc0PMdpMS1fy2l4ntPwHKfFtHwtp+F5TsNzRG98tofT4k+dOqX9/X3PvzHv3bunubk5PXjwwNM+4FztUqsLCwva2dnxuh3HxrVvYNLxswmMHn4ugdHDzyUwesb159LLvj09GjlKV2wcpV4AAAAAAADgPk+DsGw26/luMEk6e/asstms120AAAAAAABggDwNwl5//fWhPt5nn3124seG3QsAAAAAAACGy/OrRg5TKpXyugUAAAAAAAB4ZKqCsM3NTa9bAAAAAAAAgEdOe91ALz777DNZluW43jRN5fP5ru4DAAAAAACAyTIWQdhnn32mbDarcrks0zR7WsO2bfl8Ppc7AwAAAAAAwLgY+SDsypUryuVykg7DLGCcrays6ODgYCSulgrgN/jZBEYPP5fA6OHnEhg9/Fx2z2ePcLr005/+VIlEQpLqu7n6adfn8+nBgweu9AZvLS4uand3VwsLC9rZ2fG6HQAAAAAA4JCXr+lHekfY6uqqpMMAy7ZtBYNBhcNhBYNBSdLjjz/ecY3PP/9clmXp5s2bunfv3kD7BQAAAAAAwOga6SDMMIz6TrCNjQ0tLy/3vFY6ndYTTzzhVmsAAAAAAAAYMzNeN9CO3++XdBhi9ROCSVIwGNS5c+dc6AoAAAAAAADjaKSDsNoRyAsXLriyXj6fd2UdAAAAAAAAjJ+RDsJqu8D29vZcXQ8AAAAAAADTZ6SDsDfffFO2bcswDFfW+8lPfuLKOhgdlUpFi4uLbW9ra2tetwkAAAAAwFRYW1vr+Dq9Uql41t9ID8s/e/as3nnnHWWzWWWzWX3nO9/pa718Pq8f/ehHLnWHUVCtVrW7u9u25uDgYEjdAAAAAAAw3Q4ODjq+TvfSSAdh0uGg/I2NDSUSCf385z/veZ27d++6trMMo2NmZkbz8/Nta2ZnZ4fUDQAAAAAA0212dlYLCwttayqViqrV6pA6auazbdv25JG7tLS0pHv37imTyWhubs7Rffb29mRZlra3t3Xz5k1ZlqUHDx4MuFMMw+LionZ3d7WwsKCdnR2v2wEAAAAAAA55+Zp+5HeESdIvfvELSdL29rZSqVRPa9i2LZ/P52ZbAABgCGxbqnb4td2MT+r0v3m31gEAAMD4Gvkg7NVXX1WhUJAk+Xw+9bKBjQAMAIDxVbWl02+1r/n2qnSqw//u3VoHAAAA42ukrxp5/fp15fP5evjV6ynOMTn9CQAAAAAAgAEa6R1h+Xxe0m92gkWjUcViMfn9/q7nhF27dk2fffbZALsFAABeef+O9FJkdNYBcDyOKAMAvDbSQZhhGPL5fPL7/drc3NS5c+d6Xuvy5ct6/PHHXewOAAAMw/t3Ote8XJJOzUgvnB/8OsC0cTO84ogyAMBrIx2E+f1+3bt3T1euXOkrBKutdf48/6oFAGCc3Lsvvfahs9pXPpCe/l1p9szg1gGmEeEVAGCSjPSMsEjk8GxCMBh0Zb3r16+7sg6A6WZZltctAFPjPUP68htntV98Ld0wBrsOAAAAxttIB2GpVEq2bcs0TVfWG5UdYaZpKhaLOao1DEOpVEqhUEg+n08+n0+hUEiZTGYgL8Zrj1Mqlbq+by6XUywWUyAQqPeZSCRULpdd7xMYNsMwlMlkFAgEdPnyZa/bAabGR592V//xCfVurQPgeE6OHnuxFgAArUY6CIvH4zp//rzW19ddWe9nP/uZK+ucxLKseljV7hYKhTrucrMsS4lEQpFIRIVCoSkMNE1TuVxOgUBAhULBtf4zmUxPoWO5XFYgEFAmk5EkFYtFbW9vK5vNyjAMxWIxxWIxdtFgYGohVSwWUygUqoexgUBAkUhEsVhMmUzmxFC29rPWyLIslUolJRKJ+jq5XI7vY2DI9r9yp96tdYBp5HS+npM6N9cCAKAXIx2ESYehytbWlv76r/+677VWV1dd6Ohk3YRStdDoOJZlKRKJONqVlUqllEqlHD/uScrlsnK5XE/3q4VcyWRSGxsbikajCgaDisfj2t7eVjgcVrlcViQSIUQ4QaFQcBSitt5aw5te9fLYgUDgxPUSiURPa3aze9CyLGUymfrnIZfLqVwuyzTN+veZZVkyDKP+/R2Lxer1mUxGhUJBiUTi2J+1y5cva3V1VaVSie9bwEOBR9ypd2sdYNp0O1/v4P5w1gIAoFcjH4QFg0Fdu3ZNP/rRj/pa5969ezKMwQ78cBq01YKikyQSCZmmqXA4XN9dtb29rWKxqHQ6faS+UCj0dJSxprb7rJ/7BYNB5fP5Y+uKxaKkw51svTzONEgmk9rf39fW1pbi8Xjb2mg0qo2NjXq9G2zb1vb2tvL5vPx+/4l1ta/z9va29vf3T6wrFova39+vB6PtpNNpbW1taX9/v2NtTe2IYmt4G4/H6+H5/v5+/XltbGwomUzWn5thGMrlckqlUvWfndawq7aObduuzSkEpoltSw+q7W92h6vQSdLTT3b3uE+dUO/WOjVuPT9g1Lk5X49ZfQCAUeCz7dH9Z1rjUcZ0Oi2fz9cxJDiOZVm6efOmLMvSgwcP3GyxrlAoKJVKKZ1Od5z/tbS0dGLY0LhONps9tqY
      "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(x=ak.to_numpy(array[\"tx\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"dx/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 21,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAACi9UlEQVR4nOz9X4xbd57n/X1IyR4Z9lCscmN2C1XBtFnjWXUuMiNSmjxpIEi2i5zZABbsaZPSwI6lBOsibaOxCGrHpDW+yoVcZvV03SziEakeIJEfIVsiu9uGDaSjojwPEAR4nlUV5clNvNvDIzdSBebZrKuO2DastSWeXFSTXX/Jwz+Hh+R5vwDC9efL3/nWH5XJD39/fJZlWQIAAAAAAADGnN/tBgAAAAAAAIBBIAgDAAAAAACAJxCEAQAAAAAAwBMIwgAAAAAAAOAJBGEAAAAAAADwBIIwAAAAAAAAeAJBGAAAAAAAADyBIAwAAAAAAACeQBAGAAAAAAAATyAIAwAAAAAAgCcQhAEAAAAAAMATCMIAAAAAAADgCQRhAAAAAAAA8ASCMAAAAAAAAHgCQRgAAAAAAAA8gSAMAAAAAAAAnkAQBgAAAAAAAE8gCAMAAAAAAIAnEIQBAAAAAADAEwjCAAAAAAAA4AkEYQAAAAAAAPAEgjAAAAAAAAB4AkEYAAAAAAAAPIEgDAAAAAAAAJ5AEAYAAAAAAABPIAgDAAAAAACAJxCEAQAAAAAAwBMIwgAAAAAAAOAJx91uAOjGk08+qQcPHujYsWP6gz/4A7fbAQAAAAAANv3n//yf9ejRI504cUJfffXVQK/tsyzLGugVgT44duyY6vW6220AAAAAAIAu+f1+PXr0aKDXZEYYRlIjCPP7/ZqamnL8etVqdaDXs4OeRq8fiZ7soid7hq2nYetHoie76MmeYetp2PqR6Mkuehq9fiR6soue7Bm2ngbdT+N6x44dc/xa+xGEYST9wR/8gTY3NzU1NaWNjQ3HrzczMzPQ69lBT6PXj0RPdtGTPcPW07D1I9GTXfRkz7D1NGz9SPRkFz2NXj8SPdlFT/YMW0+D7qdxPTe2OmKzfAAAAAAAAHiCqzPCnn32WTcvv4fP59N/+k//ye02AAAAAAAA4BBXg7BKpSKfzyc39+tvXN/n87nWAwAAAAAAAJzn+aWRHJoJAAAAAADgDa5vlm9ZluLxuEKhkCvXN01T+XzelWsDAAAAAABgcFwPwvL5vF599VVXe4jH4/qLv/gLV3sAAAAAAACAs1xfGhmNRt1uQWfPnnW7BQAAAAAAADjM9Rlhk5OTbrcAtLWwsKBaraZAIOB2K030NJqG8XtET/YMY0/DZhi/R/RkzzD2NIyG7fs0bP1I9DTKhu37NGz9SPQ0yobx+zRsPQ1bP07yWS7uFn/s2DFtb2+7/o2+f/++Jicn9ejRI1f7gH0zMzPa3NzU9PS0NjY23G4HQ4rfE4w7fscxzvj9xjjj9xvjjN9v2OHm74mrSyOH6cTGYeoFAAAAAAAA/edqEJbNZl2fDSZJJ0+eVDabdbsNAAAAAAAAOMjVIOzNN98c6PU+//zzIz836F4AAAAAAAAwWK6fGjlIqVTK7RYAAAAAAADgEtdPjRyktbU1t1tAn1WrVc3MzLSsWVhY0MLCwoA6AgAAAADAu5aXl7W8vNyyplqtDqibg0YyCPv8889lmqbtesMwlMvlOroPRkO9Xtfm5mbLmlqtNqBuAAAAAADwtlqt1vZ5uptGIgj7/PPPlc1mVSqVZBhGV2NYliWfz9fnzuA2v9+vqampljXDcCADAAAAAABeEAgEND093bKmWq2qXq8PqKO9hj4Iu3z5spaWliTthFnAblNTU9rY2HC7DQyphYUF1Wo1wlCMLX7HMc74/cY44/cb44zfb9jZnmhmZsa1WWM+a4jTpZ/97GdKJBKS1JzN1Uu7Pp9Pjx496ktvcFfjH8309DRBGAAAAAAAI8TN5/RDPSNscXFR0k6AZVmWQqGQwuGwQqGQJOnpp59uO8YXX3wh0zR18+ZN3b9/39F+AQAAAAAAMLyGOggrl8vNmWCrq6uam5vreqx0Oq1nn322X60BAAAAAABgxPjdbqCVYDAoaSfE6iUEk6RQKKRnnnmmD10BAAAAAABgFA11ENZYAnn27Nm+jJfL5foyDgAAAAAAAEbPUAdhjVlgW1tbfR0PAAAAAAAA3jPUQdjf/M3fyLIslcvlvoz305/+tC/jAAAAAAAAYPQMdRB28uRJvfvuu1pZWdFvfvObnsdjaSQAAAAAAIB3DXUQJu1slB8Oh5VIJHoa5969e32bWQYAAAAAAIDRc9ztBuxYXV3VmTNn9OyzzyqTyWhyctLW/ba2tmSapiqVim7evOlwlwAAAAAAABhmIxGEffLJJ5KkSqWiVCrV1RiWZcnn8/WzLQDALpYl1a3WNX6fxJ9iAAAAAG4Z+iDs9ddfVz6flyT5fD5ZVptnWYcgAAMA59Ut6fjbrWseXpGO8ScZAAAAgEuGeo+wa9euKZfLNcOvbkKwXu4HAAAAAACA8THUM8Iapzw2ZoJFo1HFYjEFg8GO9wm7evWqPv/8cwe7BQAAAAAAwDAb6iCsXC7L5/MpGAxqbW1NzzzzTNdjzc/P6+mnn+5jdwAAAAAAABglQ700MhgMSpIuX77cUwjWGOv06dN96AoAAAAAAACjaKiDsEgkIkkKhUJ9Ge/atWt9GQeAt5mm6XYLfWVZ0qN66xtbLQIAAAAYB0O9NDKVSun27dsyDKMv4zEjDEC3yuWyVlZWlM/nFY1GVSgU3G6pbzjtEQAAAIBXDHUQFo/Hdfr0aa2srOiv//qvex7v5z//uX74wx/2oTMAw6QRUpXLZRmG0TwkIxgMKhQKaXJyUuFwWLFYTNFo9MD9E4mEDMPQ+vp682OmaapUKmllZUWlUmnsZoEBAAAAgBcNdRAmSYVCQX/0R3+kX/ziF/rLv/zLnsZaXFwkCMNQyufzSqVSHd8vHA7vCW+65fN1PtUnGAxqe3v70M8lEgkVi8WOx1xdXT00qDqMaZpaXFzU0tJSy5pyuSxJKpVKzdpwOKxoNKrZ2Vmtrq6qWCwqHA7vue/8/LwMw2jeH4NlWTsz1Vrx+6QufnUBAAAAeNhQ7xEm7ewPdvXqVb366qs9jXP//n2e0GJoJZNJbW9va319XfF4vGVtNBrV6upqs74fLMtSpVJRLpdrHlJxmFAopFwup0qlcmQIJu0E2Nvb27aCrXQ6rfX1dW1vb9sOwTKZjCYmJg6EYPF4XIVCoTle4+taXV1VMplsfm3lcllLS0tKpVLNwG7/jK/GOJZl9W2fQtjXWK7Z6tYuKAMAAACA/YZ6RtjPf/5zSdJ3vvMdTUxM6Nlnn20bEhzGNE3dvHmz3+0BfRUMBhUOh1UoFJRKpZTP5w/UNIIeJ4RCISWTyeZMqcOsr6+3DMp2CwaDikajikajikQihwbRuVxOyWTSdo+GYSgWix3YNzCZTCqXyx16n1AopFAopGg0qlwup6WlJWUyGdvXlHbCx8N+HgAAAACA0TLUQdg777yju3fvNt+3LKvlMqhWLMvqavkX4IajgrBulk92KhQKKRwOHwiu4vG47RBsv1QqdWjvnYRgpVJJsVhsz8eCwaBu3759YFljK+l0WvF4XJFIZM8ssK2trSPv0+3XDQAAAAAYLkO9NPL8+fOyLEuWtbP+hSALXnHUUrxBLdE77Dq9XPuw+3YSLpXL5b6EYLv7uXfv3p4e2AwfAAAAAMbfUAdhjRkkPp+vGYh1ewNGyVEh0eTk5GAbGQKGYWhubu7AxwuFQlchWEMwGOzbHmvDpFgsHlg6it8plUputwAAAADARUO9NPLkyZMKh8O6e/eustmswuFwV0GAYRh655139Omnn/a/SWAMHfbv7Omnn+7reHb/LScSiQOztRp7mfUqFAopnU53veR6kMrlsnK5nEqlkgzDUDAY1OTkpILBoEKhkC5cuCBp57TLbDbbctlpqVRSIpHQ5OSkcrmc/uUPev9eDivDMFQqlVQoFJoh2DC9OLL/Z9GP32sAAAAARxvqIEySLly4oNn
      "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(x=ak.to_numpy(array[\"ty\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"dy/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -300, 300, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"x\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -300, 300, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"y\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#bins = np.linspace( -1.0, 1.0, 50 )\n",
    "#sns.regplot(x=ak.to_numpy(array[\"dSlope_out\"]), y=ak.to_numpy(array[\"z_mag_x\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/work/guenther/reco_tuner/env/tuner_env/envs/tuner/lib/python3.10/site-packages/IPython/core/pylabtools.py:151: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.\n",
      "  fig.canvas.print_figure(bytes_io, **kw)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAACrjklEQVR4nOz9XWxbZ5rme1+UncDZSdGk0hhAIIEuU51GCvug26Q9B7WPpkR2HUyyk66QNlCZJEB3iUqya3oA9YSMK8eJiuwuAe/b86YtumaAcTrA2GTVJEgayIR0+mgPsMcSnT4LporLKUAEgQFKWmYnO0Y+uN4DNVmkPsjFj8XFj/8PEFq0bj7rlimlmpfv51key7IsAQAAAAAAADNuwe0GAAAAAAAAgHEgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFwgCAMAAAAAAMBcIAgDAAAAAADAXCAIAwAAAAAAwFw47XYDwCAefvhh3b9/X6dOndK/+Bf/wu12AAAAAACATf/rf/0vffPNNzpz5ow+//zzsV7bY1mWNdYrAiOwsLAgfnQBAAAAAJheCwsL+uabb8Z6TSbCMJXaQ7BAIOBiJ/2p1WpqNBpaWFjQ0tKS2+3YRt/jN6290/d40fd40fd40fd40fd40fd40fd40fd4TXvfp06dGvu1CcIwlRYWFlq/7Lu7u263Y1swGFS1WtXS0hJ9j8G09i1Nb+/0PV70PV70PV70PV70PV70PV70PV70PV7T3rcbRx1xWD4AAAAAAADmgqsTYY899pibl+/g8Xj0P//n/3S7DQAAAAAAADjE1SCsUqnI4/G4euh58/oej8e1HgAAAAAAAOC8ud8ayZ0HAQAAAAAA5oPrh+VblqV4PK5QKOTK9U3TVC6Xc+XaAAAAAAAAGB/Xg7BcLqcf/ehHrvYQj8f1/e9/39UeAAAAAAAA4CzXt0ZGo1G3W9DFixfdbgEAAAAAAAAOc30ibHFx0e0WMIW+9a1v6d69e/rWt77lditzYX19XfV6XV6v1+1W+jKtfUvT2/u09j2tpvXvm77Ha1r7nlbT+vdN3+M1rX1Pq2n9+6bv8ZrWvtE/j+XiafGnTp3S/v6+6z9o9+7d0+Lior755htX+4B9wWBQ1WpVgUBAu7u7brdj27T2Dcw6fjeBycPvJTB5+L0EJs+0/l662berWyMn6Y6Nk9QLAAAAAAAARs/VICyTybg+DSZJZ8+eVSaTcbsNAAAAAAAAOMjVIOyVV14Z6/U+/fTTE7827l4AAAAAAAAwXq4flj9Oa2tr+m//7b+53QZGqFarKRgMdq1ZX1/X+vr6mDoCAAAAAGB+bW5uanNzs2tNrVYbUzdHzVUQtr297XYLGLFGo6Fqtdq1pl6vj6kbAAAAAADmW71e7/k+3U1TGYR9+umnMk3Tdr1hGNra2urrOZgOCwsLWlpa6lozCefQAQAAAAAwD7xerwKBQNeaWq2mRqMxpo46TUUQ9umnnyqTyahUKskwjIHWsCxLHo9nxJ3BbUtLS1N1i1gAAAAAAGaZneOJgsGga1NjEx+EXblyRdlsVtJBmAVMs/X1ddXrdabUgAnD7yYwefi9BCYPv5fA5OH3sn8ea4LTpV/84hdKJBKS1JrmGqZdj8ejb775ZiS9wV3N9DgQCDARBgAAAADAFHHzPf1ET4RtbGxIOgiwLMtSKBRSOBxWKBSSJD366KM91/jtb38r0zR18+ZN3bt3z9F+AQAAAAAAMLkmOggrl8utSbBisaiVlZWB10qlUnrsscdG1RoAAAAAAACmzILbDXTj8/kkHYRYw4RgkhQKhXTu3LkRdAUAAAAAAIBpNNFBWHML5MWLF0ey3tbW1kjWGZZhGIrFYq6umc1mFYvF5Pf75fF4tLy8rEQioVKpNHAPTqwJAAAAAAAwKhMdhDWnwPb29ka6nlNM05TH4+n5sby83Ar5xr1mqVSS3+9XOp2WJOXzeVUqFWUyGZXLZcViMcViMZmmafv7dmJNAAAAAACAUZvoIOwnP/mJLMtSuVweyXo///nPR7LOSXK5nO3aZmg0zjVLpVIrkEomkyoWi4pGowqFQorH46pUKgqHwyqVSopEIraCKyfWBAAAAAAAcILHsizL7Sa6yWazymQy+vTTT/Wtb31rqLUuXryo27dvj6izo/x+v62gJxqNqlgsjnVN0zR17tw5maapUCikSqVybJ1hGFpeXnZtTbvcvNUqAAAAAAAYnJvv6Sf6rpHSwUH5xWJRiURCH3zwwcDr3L17d2STZcfJ5XIyTVOpVKrnWV0XLlwY+5qJRKIVqHWbHGtOchUKBZVKJeVyOSWTybGtCQAAAAAA4JSJnwhrunDhgu7du6d0Oq3FxUVbz9nb25NpmqpUKrp586ZM09Q333zjSH/NiaeTpqLcXLN9IkuS9vf3W3fkPE6hUFAikZB0cOfO/f39sazZDybCAAAAAACYTkyE9fDRRx9JOgiE1tbWBlrDsix5PJ5RttVSKBRkGMZI70o5yjUzmUzr82g02jWwkqR4PN763DRNFQqFjj9zak0AAAAAAHAyy5IaPcaZFjySQ/HHTJj4IOyll15qHRjv8Xg0yACbUwFY08bGhnw+ny5dujSRa7YfuB8Oh209JxQKyTAMSdKNGzeOhFZOrAkAAAAAAE7WsKTTr3Wv+fp16RRB2Ikm+q6R165d09bWViv8GnQXp5O7P8vlssrlskzTlN/v1/LystbW1lQoFCZizcPnol28eNHW89rDrcPXdWJNAAAAAAAAp030RFhzW2BzEiwajSoWi8nn8/V9TtjVq1f16aefjrzHw4fEG4ahXC7XmpiKx+O6cuWK7ampUa9ZKpU6HodCIVs9HK4rl8ut6zmxJgAAAAAAgNMmOggrl8vyeDzy+Xza3t7WuXPnBl5rdXVVjz766Ai7OwioDodChxUKBRUKBSWTSVvnfY16zdu3b3c87nWWV9Phv6vt7e1WaOXEmgAAAAAAAE6b6CDM5/Pp3r17unLlylAhWHOt8+fPj6izA6FQSFtbW607U5ZKpdYZWIflcjltb29rZ2dnrGsefu6g01vtd650Yk0AAAAAAACnTXQQFolE9NFHH9kOWnq5du3aSNZpl0wmOx6bpqlcLqeNjQ2ZptnxtXK5rFgspmKxOLY1TwrR+tV+XSfWBKaJaZq2JyEBAAAAAJNjog/LX1tbk2VZIwteRj0Rdhyfz6dUKqX9/X3l8/kjb5ZLpZKy2ezY1hw0bDp8jb29PUfXBCZduVxWOp2W3+/X6uqq2+0AAAAAAAbgsZy8peIIRCIRLSwsHDmXahC//OUv9YMf/GAEXdlnmqZWVlY67rTo8/m0v78/ljU9ns57ptp9uUulkmKxWOtxNBptTZ05sWa/gsGgqtWqFhYWtLS0NNAa7dbX17W+vj70OnBHuVzWjRs3VC6XZRhG6yYZPp9PoVBIi4uLCofDisViikajR56fSCRkGEbHNmPTNFUqlXTjxg2VSqWOADgejyufz4/jWwMAAACAlm8a0unXutd8/bp0yqGxp83NTW1ubg69Tq1WU6PRUCAQ0O7u7gg6s2+it0ZKUj6f1x/8wR/ov/7X/6o//dM/HWqtjY2NsQdhPp9POzs7ikQireCq+Qb7uDfko17T5/ONZAti+zSXE2sOqtFoqFqtDr1OvV4feo1h5HI5ra2t9f28cDjc89w5Ow6Hm3Z0C3QTiYQKhULfaxaLRdu/F6ZpamNjo+uEpWmard+R9snJcDisaDSq5eVlFYtFFQqFIzduWF1dlWEYHYEzAAAAAMyzer0+kvf
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( -1, 1, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"dSlope_fringe\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=bins, fit_reg=None, x_estimator=np.mean)\n",
    "plt.xlabel(\"$\\Delta$dx/dz\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.05, 0.95, 'LHCb'),\n",
       " expsuffix: Custom Text(0.05, 0.955, 'Simulation'))"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/work/guenther/reco_tuner/env/tuner_env/envs/tuner/lib/python3.10/site-packages/IPython/core/pylabtools.py:151: UserWarning: This figure includes Axes that are not compatible with tight_layout, so results might be incorrect.\n",
      "  fig.canvas.print_figure(bytes_io, **kw)\n"
     ]
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAADwAElEQVR4nOz9bWxb953nf38o2WnSpDKlDIK6UqY1NRn0j6LNmJSv5Or1qBW5vUFypTMh7SZGkuI/EZkExQ5GuyXjFthnjUq21WCwQGLSnQdx1k1tst3JlaDNWnTm2R+Ztcg0LbIIdsrjzEaKimAiHatJk8Y2z/VAQ5akJPLw5vBGfL8Aorr58ne+unFcfvz7fY/LsixLAAAAAAAAwB430usGAAAAAAAAgG4gCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFAgCAMAAAAAAMBQIAgDAAAAAADAUCAIAwAAAAAAwFDY1+sGgFbceOON+uCDDzQ6Oqpbbrml1+0AAAAAAACb3n77bV27dk3XX3+93nvvva5e22VZltXVKwIdMDo6qmKx2Os2AAAAAABAi0ZGRnTt2rWuXpMdYRhIpSBsZGREBw8e3Pb5tbW1up/vlX7tS+rf3uiref3aG301p1/7kvq3N/pqXr/2Rl/N69fe6Kt5/dobfTWnX/uS+rc3+mpev/Zmp69SzejoaJe7IwjDgLrlllu0urqqgwcPamVlZdvnp6am6n6+V/q1L6l/e6Ov5vVrb/TVnH7tS+rf3uiref3aG301r197o6/m9Wtv9NWcfu1L6t/e6Kt5/dqbnb5KNb0YdcSwfAAAAAAAAAyFnu4Iu+2223p5+Soul0v/+3//7163gSatra1pampqx4+X/ndxcVHz8/Pdbg0AAAAAgKHzu9/9TtLur9dLn+uVngZhhUJBLpdLvZzXX7q+y+XqWQ9oXbFY1Orqat3Pb25udrEjAAAAAACGVynjafR6vVeGfkYYN80cbHaG5Y+NjfWgMwAAAAAAhk9po5GdYfm90PMgzLIsBYNBeTyenlzfNE2lUqmeXBvtszMsn2ORAAAAAAB0x8c+9jFtbm7aGpbfCz0PwlKplB5++OGe9hAMBvWlL32ppz0AAAAAAADAWT2/a6Tf7+91Czpy5EivWwAAAAAAAIDDer4jbGJiotctYA+an5/X5uZm380H69e+pP7urR/18/erX3ujr72jX79n9NW8fu6tH/Xz96tfe6OvvaNfv2f01bx+7q0f9fP3q19769e+SlxWD6fFj46OamNjo+ffnMuXL2tiYkLXrl3raR+wr3SeeHJyctczx9gb+FljWPG7j2HG7z+GGb//GFb87g+XXv68e3o0sp/u2NhPvQAAAAAAAKDzehqExePxnu8Gk6QDBw4oHo/3ug0AAAAAAAA4qKdB2Le+9a2uXu+NN97Y9XPd7gUAAAAAAADd1fO7RnZTJBLpdQsAAAAAAADokaEKwpaXl3vdAgAAAAAAAHpkX68baMUbb7wh0zRt1xuGoWQy2dRzAAAAAAAAsLcMRBD2xhtvKB6PK5vNyjCMltawLEsul6vDnQEAAAAAAGBQ9H0QduLECSUSCUlbYRaA4TI/P6/Nzc2+uMMs0E387mOY8fuPYcbvP4YVv/voFpfVx+nST3/6U4VCIUkq7+Zqp12Xy6Vr1651pLd2GIahSCSipaUlW7WVu+Hcbrc8Ho9mZmYUi8Xk8Xiavn4+n1cymazaYefxeBQMBnXixAm53e6m15SkRCKhpaUlLS8vyzRNeTweeb1eRSIR+f3+ltbczdTUlFZXVzU5OamVlZWOrg0AAAAAAJzTy9f0fR2EzczMKJ/Py+VyybKscrBSCn9uvvnmhmu88847Mk1T586d0+XLlx0NwkzT1Pj4uK3acDisZDJZtyaRSCgWi9Wticfjikajtvubm5tTJpOpW5dMJhUOh22tKUnZbFahUEimacrv95cDunw+r1gsJsMw5Pf7lU6nWw7ZahGEAQAAAAAwmAjCdjEyMlLeCXb+/HnNzs62vJZhGLrtttscDcLsBFclhUKh7m6uQCCgbDYrt9stv98vj8cjwzCUz+e3zUlLp9MKBoN1r2eapnw+n+0Za3aCOmkrBAsEAnWf4/P5lM/n5fF4lMvlOhKGEYQBAAAAADCYevmafqSrV2tSKTCJRqNthWDS1tG/Q4cOdaCr3S0sLNiqKwVbu4nFYspms4rH49rY2FA6nVY8Hlc6nVahUFA8Hq+qLx0frScUCskwDHm93vI6hUJB6XR6xx1lqVSq4c4x0zTL1/Z4PLsGZ+l0WtJWGGmnVwAAAAAAACf09bB8j8ejV155RUeOHOnIenZ2OLUqlUrJNE1Fo9HyDqndzMzM7Po5wzDKs7Z2m6sVjUZVKBSUSqXKH8vn8/J6vbv2ls1mFY1Gt4VopdlgkUhEoVBI+Xy+/Lm5ubm6O81KxyEl1d0JV7pGJpNRNptVKpVq6uglAAAAAABAJ/T10chYLKYf/OAHSiaTevjhh3vdTl3T09OSto48tiMUCunIkSMN537VziOrNytsenpaHo+n4XB+wzDKX0fJboFcbe3GxkbdI4+ZTKa8G8ztdmtjY6NuL41wNBIAAAAAgMHE0chdfPvb35ZlWVW7lNrxox/9qCPr1MpkMjIMw/Z8sHpKu8oaKd09svL9nZRmipWOJ9bj8Xi27Rjb7XtfWef3+xvO/arcWWaaZsNjlwAAAAAAAJ3W10HYgQMH9L3vfU9nz57V7373u7bXc+po5MLCgtxut44ePdr2Wo12bVVaX18vv73bccuzZ88qHA7bHlBfu/vrnXfe2bGu8ljmbkcya1UGd2fPnrX1HAAAAAAAgE7p6yBM2pqH5fV62x6yfunSpY7tLKuUz+eVz+fLRxWnp6cViUQc3/FkmmZ5Ppff7981jDp27Ni2XV711K5Te1RS2r5LzO4Mt8q12REGAAAAAAC6ra+H5ZcsLS1pZmZGt912m2KxmCYmJmw9b319XaZpqlAo6Ny5c470Vnsc0jAMpVKp8o6pYDCoEydO2N41ZVfp6/F4PHWPPTZ73VK4VrLT3S2z2WzDmp3U1tUb8A8AAAAAANBpAxGEvfTSS5K2BtFHIpGW1rAsSy6Xq5NtyTCMbaFQrUwmo0wmo3A43LGjmaZpKhKJyOv1Kp1O2z72aIdhGFXv7zQo/+LFi1Xv273+zTffXPX+8vLyng/CLMtSsUHNiNTx300AAAAAALBd3wdhjz76aHl3lcvlUis3uXQqZPB4PEomk+VdZ9lsdluQVJJKpbS8vKxcLtfWNQ3DUCAQkNvt1oULFzoagklb4VRJOBzetYdKre4Ia/cOm4OgKGnsuRfr1mze82WNdqcdAAAAAACGWl/PCDt16pSSyWQ5/GolBGvneXaEw2FFo1Elk0kVCgVtbGwoHo/vGFDl83kFAoGWr5XJZDQ9PS3DMMozyRKJRBvdb1e5a223u2DuFvY1q/YYJgAAAAAAgJP6ekdYKZQp7QTz+/3l3VDNzgk7efKk3njjDQe73eJ2uxWNRhWNRpXJZDQ3N1cV+GSzWSUSCUWjUVvrmaapVCqlZDK5YwAVi8V08eLFunPC7DIMozwIPx6P77rTq9UAqzYcrLzrZavW1tY0NTXV9jrz8/Oan59vex0AAAAAAPaqxcVFLS4utr3O2tpaB7ppTV8HYfl8Xi6XS263W8vLyzp06FDLa83NzW2bUeW0YDAov9+v2dnZqjstLiws2A7CstmsCoWC/H7/rjPJMplMU+Habkp3l/R4PG2vZUcndoQVi0Wtrq62vc7
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#import matplotlib.pyplot as plt\n",
    "bins = np.linspace( -2000, 2000, 50 )\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l0\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T1X1\")\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l4\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T2X1\")\n",
    "sns.regplot(x=ak.to_numpy(array[\"x_l8\"]), y=ak.to_numpy(array[\"z_mag_x_fringe\"]), x_bins=50, fit_reg=None, x_estimator=np.mean, label=\"T3X1\")\n",
    "plt.legend()\n",
    "plt.xlabel(\"x [mm]\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "['tx^2' 'tx dSlope_fringe' 'ty^2' 'dSlope_fringe^2']\n",
      "intercept= 5205.144186525624\n",
      "coef= {'tx^2': -320.7206595710594, 'tx dSlope_fringe': 702.1384894815535, 'ty^2': -316.36350963107543, 'dSlope_fringe^2': 441.59909857558097}\n",
      "r2 score= 0.9604900589467942\n",
      "RMSE = 8.772908410819978\n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso, Ridge\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",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, 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",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"dSlope_fringe\",\n",
    "    \"tx^2\",\n",
    "    \"tx dSlope_fringe\",\n",
    "    \"ty^2\",\n",
    "    \"dSlope_fringe^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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",
    "\n",
    "lin_reg = LinearRegression()#Lasso(alpha=0.01)\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))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(exptext: Custom Text(0.0, 1, 'LHCb'),\n",
       " expsuffix: Custom Text(0.0, 1.005, 'Simulation'))"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAABMIAAAOWCAYAAAANzz7PAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjYuMCwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy89olMNAAAACXBIWXMAAA9hAAAPYQGoP6dpAADXPElEQVR4nOz9T2xb+X3v/7+ONLdIcVvpSL47G4h1OL7odxWLlO+iu4zImS5apBmT8iBF20Ut0g76uwulJqXJOrZJz2hzf0hMyrP4pWhQi7LToF1MTcrZfRexeOzsgjvhkQKMd9fSkdqiQRvr/Ba6ZEiJoij+PRKfD4CoKH34OW+KtBq+5v35fAzP8zwBAAAAAAAA59zIoAsAAAAAAAAA+oEgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMAAAAAAAAQ4EgDAAAAAAAAEOBIAwAAAAAAABDgSAMOINc1x10CQAAAAAAnDkEYcAZYdu2UqmUJiYmND8/P+hyAAAAAAA4c94ZdAHAeWDbth4/fizbtuU4jra3t+W6rkzTlGVZmpycVDAYVCQSUTgcPvL4WCwmx3FUKpWq33NdV8ViUY8fP1axWKQLDAAAAACADhme53mDLgKDl8vllEgkTv24YDBYF960yzCMUz/GNE3t7Ow0/FksFtPa2tqp5ywUCg2DqkZc19W9e/eUyWROfZ1gMKhwOKxAIKBCoaC1tbUjv8tKOGbb9pHHR6NR5fP5U18XAAAAAIBhxtJISJLi8bh2dnZUKpUUjUabjg2HwyoUCtXx3eB5nsrlsrLZrEzTPHacZVnKZrMql8vHhmCSlM/ntbOz01KwlUwmVSqVtLOz03IIVlmieDgEqwRUlfkqz6tQKCgej1efm23bymQySiQS1cDucMdXZR7P82RZVkt1AQAAAACA49ERhoYSiYRyudyR7/ejE8lxHAUCgYY/29nZaRqUHScUCjXsrMpms4rH46eqLRKJyHGcuu/H43Fls9mW5shkMkqlUke+b1mWyuVyw8ccfj3oCAMAAAAA4PToCENDxy2TbGf55GlZlqVgMHjk+9FotK0QTDq+7tOEYMViUYFAoC4EM01TpVKp5RBMOuhAK5fLR57L9vb2sY9p93kDAAAAAIDfIghDQ8ctxevXEr1G1+nk2o0ee5pwybZtRSKRI49fX19vGNq1Us/m5mZdDWyGDwAAAABAbxGEoaHjQqLJycn+FuIDjuNodnb2yPfz+XxbIVhFpZvsPFpbWzuyfBS/VSwWB10CAAAAAAwlgjD4UqPA7cKFC12dr9VQLxaLHenWisfjLW+s34xlWUomkx3P0y+2bSuRSCgQCMgwDE1MTCgQCCgUClVP6lxbW9P8/PyJYU+xWKw+/rwHQ47jKJfLKRKJyDCMI92FgzZMrwUAAACA4fbOoAsA/KxYLDbcZD+dTnftGktLS9XTJ13X9eV+YK7ran5+vnrCpWVZ1dNFHceRbduybbv688pjmqkEjK7rKpFIHHtQwFnmOI5isZgcx/H10tdheC0AAAAAQKIjDGiq0Sb7nWza34hpmqfatH8QKt1e0sFJm+VyWfl8Xvl8XqVSSeVy+UiH3ElhyjAss7UsS6VSSTs7O77u/BuG1wIAAAAAJIIw4FjFYrHhPldLS0tdv1YsFpPU/OTIZorFomKxWN2SxUgkolwu13FtmUymulwumUw2DO0sy1KhUKjrlDtpj7BsNivLsmRZlvL5fMd1+p3flkPWGrbXAgAAAMDwYmkkcIzjAoFONsg/TjgcVjqdPnVnjuu6mp2dPbJ803VdFYtFFYtFpdPpjjb2v3fvXvXrRh1ytZLJpN68eaNMJnNiqBcOh4dqCd6guq4SiYSy2WzTMcP2WgAAAAAYXnSEAcdotGl4NzbIP04ymTzVkkvHcTQ1NdVwD7PD40KhUFuboLezt1U6nVYwGOTUSB/I5XJd6QoEAAAAgPOCIAxowHGchkFOL7rB2uG6bnWD82g0qmw2q1KppEKhcOxeVJFI5NTh1OEQrNUwbWlpydebww8Dx3FO7OADAAAAgGHD0kiggeMCo0Ag0OdKGqsEUoVC4UiXWjgcViKRaBh8xWIxlUqllq9jWVbd/VQqpbm5uRM71yonSvr1FMzzznVdX+9JBgAAAACDQkcY0MBxQZifTtdrFIJVVDavP8y27VMtkTRNsy7Icl1XoVCopc4yz/MIwQagsm8cS1MBAAAA4CiCMJwZqVRKhmG0dQuFQqe61nHL+vwS7JimeeJ+ZZZl1Z3iWNHoe80cHu84jgKBQFf2nrJtW4lE4sROO9d1lcvlFAqF6q5bWf43MTEhwzAUCASUyWSOPN5xHKVSqeqpmoFAQIlEouHrXLnO4ffQYZFIpDpf7a2bS0JzuVzddSYmJhQKhRo+R+ng99lo37ja+o77Xbf6WlS4rqtMJqNIJFL9/VfqS6VSLQVxjuMok8koEAjU1Vz7GgQCgeoyYAAAAADomAccQ9KR287OTl+uHY/HG16/mzfLsk59/UKh0Jfnf1gymayrIxwOt/xY0zSPPI/TCofDDX8fwWDQK5VKp5qrVCp58Xi8ri7TNI+M29nZ8dLptBcMBuuumc1mPc87+jupvcXj8eo8zcZZlnXse/rwe+A46XS6pX8jpVKp5dcgn89Xfy/5fL5uDsuyjq19Z2fHK5VKR2oqlUp1t9r5WnktDstms9Xx6XTaK5VKXrlc9vL5fN3rFY1Gjzx2Z2fHSyaT1edRW+POzs6R17u2rnK5fGJtAAAAANAMHWE4M+LxuEqlUlu3bDZ7qmsd1/nll66U03SmNeocO+mkycMKhUJ136/D84RCoZY7dhzH0ePHj1vuOgoGg0c2fC+VSgqFQrJtW/l8XuVyWTs7O3WHBORyOdm2rUgkomKxqEKhoJ2dHe3s7KhQKFR/f47j6N69ew2vHYvFWqqx0e+lE8VisXrtpaWluvmDwaDy+bykg9rn5+frHmuapoLB4JFDHSrfq/3ZaV+LilQqpUQiIdM0VS6XlUwmFQwGZVmWotGoSqVStea1tTUFAoG698b29rYikciR9+XGxoampqaqy3rL5bKy2Wz1tXJdl83/AQAAAHRu0Ekc/Es+6whLp9Ntz3e4G0cndIQd7qip3CrdSP12uKupUafNcSrdRd14Hsf9XtqZt3auk7qQDncJHdeZd7hzrbYzrFbt7+S4axcKhZY6uHZ2drraERaNRk98z580T6u1V7T6WtR2gp30Wtd2fB3XwVhbo2maDV/Xw+/ffv0NAgAAAHA+0REGNOD3jrDTOHzyo9T+80gmk9rZ2Tm2C6pyWmUr8zeq6zi1hxTE4/Fj90c7fFLicZ2AtfW7ruvb1/W492Ht97uxKX6rr0UqlapePx6PtzRWOuhya3RIQ+3zWFlZafi6Hn6vbWxstFQrAAAAADRCEAY0cFwwUC6X+1xJ5xqFKW/evOlovnw+r1KpdGQJnnQQeszOznY1XGp1KWjt63aa5aN+OmFxZWVF6XRa+Xz+xLBJ6l84m8vlqtdqJTiLx+N1r8FpD2moVXs9P71WAAAAAM4egjCggeM6jhp1tfhdo9DiwoULHc8bDAaP3X/Ntu0j+1f1Q7unem5vb3e3kA6YpqlkMnnsnmyn3d+tWwqFQvXrVjvIav8dncV/OwAAAADOH4Iw4BiNup0cx/HtMrrTOM2yxJPE43GVy+UjIdTa2pqvw492Q7N+cxxHuVxOxWKx4Ub4/azjtA6/z87Dvx0AAAAAZxtBGHCMGzduNPy+n8Od4xwOfboZhFXmK5VKR65z2tM68VuVUy+z2WzTfdH6pTb
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace( 5150,5300, 30 )\n",
    "ax = sns.regplot(x=y_test, y=abs(y_test-y_pred_test), x_bins=bins, fit_reg=None, x_estimator=np.mean, label=\"bla\")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(y_test, bins=30,\n",
    "        range=[5150,5300], 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)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "median_z_mag_x = np.median(array[\"z_mag_x_fringe\"])\n",
    "print(median_z_mag_x)\n",
    "params_per_layer = [[] for _ in range(12)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_array(name, intercept, coef):\n",
    "    coef = [str(c)+\"f\" for c in coef if c != 0.0]\n",
    "    intercept = str(intercept) + \"f\"\n",
    "    code = f\"constexpr std::array {name}\"\n",
    "    code += \"{\" + \", \".join([intercept]+list(coef)) +\"};\"\n",
    "    return code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "array[\"x_diff_straight_l0\"] = array[\"x_l0\"] - array[\"x\"] - array[\"tx\"] * ( array[\"z_l0\"] - array[\"z\"])\n",
    "array[\"x_l0_rel\"] = array[\"x_l0\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    #\"x_l0_rel\",\n",
    "    \"x_diff_straight_l0\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"x_l0_rel\",\n",
    "    \"tx^2\",\n",
    "    #\"tx x_l0_rel\",\n",
    "    \"tx x_diff_straight_l0\",\n",
    "    \"ty^2\",\n",
    "    #\"x_l0_rel^2\"\n",
    "    \"x_diff_straight_l0^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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()#Lasso(alpha=0.004)\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))\n",
    "print(format_array(\"zMagnetParams_l0\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[0] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l1_rel\"] = array[\"x_l1\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l1_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l1\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[1] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l2_rel\"] = array[\"x_l2\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l2_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l2\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[2] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l3_rel\"] = array[\"x_l3\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l3_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l3\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[3] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l4_rel\"] = array[\"x_l4\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l4_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l4\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[4] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l5_rel\"] = array[\"x_l5\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l5_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l5\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[5] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l6_rel\"] = array[\"x_l6\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l6_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l6\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[6] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l7_rel\"] = array[\"x_l7\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l7_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l7\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[7] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l8_rel\"] = array[\"x_l8\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l8_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l8\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[8] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l9_rel\"] = array[\"x_l9\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l9_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l9\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[9] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l10_rel\"] = array[\"x_l10\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l10_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l10\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[10] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "array[\"x_l11_rel\"] = array[\"x_l11\"] / 3000\n",
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"x_l11_rel\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform( X_train ) \n",
    "X_test_model = poly.fit_transform( X_test ) \n",
    "\n",
    "lin_reg = Lasso(alpha=0.01)\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.get_feature_names_out(input_features=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))\n",
    "print(format_array(\"zMagnetParams_l11\", lin_reg.intercept_, lin_reg.coef_))\n",
    "params_per_layer[11] = [lin_reg.intercept_] + list(lin_reg.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = [\n",
    "    \"tx\", \n",
    "    \"ty\", \n",
    "    \"dSlope_fringe\",\n",
    "]\n",
    "target_feat = \"z_mag_x_fringe\"\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data, target, test_size=0.2, random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=2, 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",
    "keep = [\n",
    "    #\"tx\",\n",
    "    #\"ty\",\n",
    "    #\"dSlope_fringe\",\n",
    "    \"tx^2\",\n",
    "    \"tx dSlope_fringe\",\n",
    "    \"ty^2\",\n",
    "    \"dSlope_fringe^2\"\n",
    "]\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\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",
    "\n",
    "lin_reg = LinearRegression()#Lasso(alpha=0.01)\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))\n",
    "print(format_array(\"zMagnetParams_dSlope\", lin_reg.intercept_, lin_reg.coef_))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.optimize\n",
    "def parabola(x, a,b,c):\n",
    "     return a*x**2 + b * x + c\n",
    "params_1 = np.array([p[1] / params_per_layer[0][1] for p in params_per_layer])\n",
    "x = [array[f\"z_l{n}\"][0] - array[\"z_ref\"][0] for n in range(12)]\n",
    "print(params_1)\n",
    "print(x)\n",
    "plt.plot(x, params_1, 'o')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "params_3 = np.array([params_per_layer[0][3] / p[3] for p in params_per_layer])\n",
    "x = [array[f\"z_l{n}\"][0] - array[\"z_ref\"][0] for n in range(12)]\n",
    "print(params_3**2)\n",
    "print(x)\n",
    "plt.plot(x, params_3, 'o')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.metrics import mean_squared_error\n",
    "feautures = [\"tx\", \"ty\", \"dSlope\"]\n",
    "data = np.column_stack([ak.to_numpy(array[feat]) for feat in features])\n",
    "target = ak.to_numpy(array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    data,\n",
    "    target,\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")\n",
    "poly = PolynomialFeatures(degree=2, include_bias=False)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "lin_reg = LinearRegression() # or Lasso if regularisation is needed\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\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))"
   ]
  }
 ],
 "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.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}