{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "ename": "SyntaxError",
     "evalue": "incomplete input (4052887367.py, line 38)",
     "output_type": "error",
     "traceback": [
      "\u001b[0;36m  Cell \u001b[0;32mIn[16], line 38\u001b[0;36m\u001b[0m\n\u001b[0;31m    #   & (array[\"match_chi2\"] < 5)]\u001b[0m\n\u001b[0m                                    ^\u001b[0m\n\u001b[0;31mSyntaxError\u001b[0m\u001b[0;31m:\u001b[0m incomplete input\n"
     ]
    }
   ],
   "source": [
    "import uproot\n",
    "import awkward as ak\n",
    "import matplotlib.pyplot as plt\n",
    "import seaborn as sns\n",
    "import numpy as np\n",
    "import mplhep\n",
    "\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "input_tree = uproot.open({\n",
    "    \"/work/cetin/LHCb/reco_tuner/data/tracking_losses_ntuple_B_BJpsi_def_selected.root\":\n",
    "    \"Selected\"\n",
    "})\n",
    "array = input_tree.arrays()\n",
    "\n",
    "array[\"dSlope_yEndT\"] = array[\"ideal_state_9410_ty\"] - array[\n",
    "    \"ideal_state_770_ty\"]\n",
    "array[\"dSlope_yEndT_abs\"] = abs(array[\"dSlope_yEndT\"])\n",
    "\n",
    "array[\"dSlope_xEndT\"] = array[\"ideal_state_9410_tx\"] - array[\n",
    "    \"ideal_state_770_tx\"]\n",
    "array[\"dSlope_xEndT_abs\"] = abs(array[\"dSlope_xEndT\"])\n",
    "array[\"x_EndT_abs\"] = abs(array[\"ideal_state_9410_x\"])\n",
    "array[\"x_EndVelo_abs\"] = abs(array[\"ideal_state_770_x\"])\n",
    "\n",
    "array[\"y_EndT_abs\"] = abs(array[\"ideal_state_9410_y\"])\n",
    "array[\"y_EndVelo_abs\"] = abs(array[\"ideal_state_770_y\"])\n",
    "\n",
    "array[\"z_mag_xEndT\"] = (\n",
    "    array[\"ideal_state_770_x\"] - array[\"ideal_state_9410_x\"] -\n",
    "    array[\"ideal_state_770_tx\"] * array[\"ideal_state_770_z\"] +\n",
    "    array[\"ideal_state_9410_tx\"] *\n",
    "    array[\"ideal_state_9410_z\"]) / array[\"dSlope_xEndT\"]\n",
    "\n",
    "array[\"z_mag_corr\"] = array[\"z_mag_xEndT\"] - array[\"match_zmag\"]\n",
    "\n",
    "sel_array = array[(array[\"z_mag_xEndT\"] < 5800)\n",
    "                  & (array[\"z_mag_xEndT\"] > 5000)]\n",
    "                #   & (array[\"match_chi2\"] < 5)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 62388 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(\n",
    "    sel_array[\"z_mag_xEndT\"],\n",
    "    bins=100,\n",
    "    # range=[5100, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    density=True,\n",
    ")\n",
    "plt.xlabel(r\"z$_{Mag}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 56032 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.35, 0.35, 35)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(5100, 5700)\n",
    "plt.xlabel(\"$t_x$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_tx_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57292 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dx/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 57944 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"$t_y$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_ty_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54048 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 53904 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_x\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 54752 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"x_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 59312 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 51)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_x$\")\n",
    "plt.ylabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_deltatx_dist.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 47860 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(-0.5, 0.5, 100)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_yEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_y$\")\n",
    "plt.ylabel(\"$z_{Mag}$ [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "intercept= 5308.975458144673\n",
      "coef= {'dSlope_xEndT_abs': 84.02382370039953, 'x_EndT_abs': 0.03552632399369171, 'ideal_state_770_tx^2': 68529.10928948055, 'ideal_state_770_tx dSlope_xEndT': 60538.43896560162, 'dSlope_xEndT^2': 13284.852569734045, 'x_EndT_abs^2': -0.0008174639168529617, 'ideal_state_770_ty^2': -2163.1533980289814, 'ideal_state_770_tx^2 x_EndT_abs': -15.687680488172175, 'ideal_state_770_tx dSlope_xEndT x_EndT_abs': -13.288664637206063, 'dSlope_xEndT^2 x_EndT_abs': -3.077894760025427, 'dSlope_xEndT_abs ideal_state_770_ty^2': 3221.0446732918413, 'x_EndT_abs^3': 1.8433768191744582e-07}\n",
      "r2 score= 0.8143147645331705\n",
      "RMSE = 36.0572267645116\n"
     ]
    }
   ],
   "source": [
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.linear_model import LinearRegression, Lasso, Ridge, ElasticNet, LassoCV\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.metrics import mean_squared_error\n",
    "import numpy as np\n",
    "\n",
    "features = [\n",
    "    \"ideal_state_770_tx\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_xEndT_abs\",\n",
    "    \"x_EndT_abs\",\n",
    "    \"ideal_state_770_ty\",\n",
    "]\n",
    "\n",
    "target_feat = \"z_mag_xEndT\"\n",
    "order = 3\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(sel_array[feat]) for feat in features])\n",
    "target = ak.to_numpy(sel_array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(data,\n",
    "                                                    target,\n",
    "                                                    test_size=0.2,\n",
    "                                                    random_state=42)\n",
    "\n",
    "poly = PolynomialFeatures(degree=order, include_bias=True)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_tx^2\",\n",
    "#     \"dSlope_xEndT^2\",\n",
    "#     \"dSlope_xEndT_abs\",\n",
    "#     \"x_EndT_abs\",\n",
    "# ]\n",
    "\n",
    "keep = [\n",
    "    # \"ideal_state_770_tx\",\n",
    "    # \"dSlope_xEndT\",\n",
    "    # \"ideal_state_770_ty\",  # no\n",
    "    # \"ideal_state_770_tx x_EndT_abs\",  # maybe no\n",
    "    # \"ideal_state_770_tx ideal_state_770_ty\", # dont keep\n",
    "    # \"dSlope_xEndT_abs x_EndT_abs\",  # dont keep\n",
    "    # \"ideal_state_770_tx^2 dSlope_xEndT_abs\", # dont keep\n",
    "    # \"ideal_state_770_tx dSlope_xEndT dSlope_xEndT_abs\",  # dont keep\n",
    "    # \"dSlope_xEndT^2 dSlope_xEndT_abs\", # no\n",
    "    \"dSlope_xEndT_abs\",  # keep\n",
    "    \"x_EndT_abs\",  # keep\n",
    "    \"ideal_state_770_tx^2\",  # do keep\n",
    "    \"ideal_state_770_tx dSlope_xEndT\",  # do keep\n",
    "    \"dSlope_xEndT^2\",  # keep\n",
    "    \"x_EndT_abs^2\",  # do keep\n",
    "    \"ideal_state_770_ty^2\",  # keep\n",
    "    \"ideal_state_770_tx^2 x_EndT_abs\",  # do keep\n",
    "    \"ideal_state_770_tx dSlope_xEndT x_EndT_abs\",  # do keep\n",
    "    \"dSlope_xEndT^2 x_EndT_abs\",  # do keep\n",
    "    \"dSlope_xEndT_abs ideal_state_770_ty^2\",  # maybe keep\n",
    "    \"x_EndT_abs^3\",  # do keep\n",
    "]\n",
    "\n",
    "remove = [i for i, f in enumerate(poly_features) if f not in keep]\n",
    "\n",
    "X_train_model = np.delete(X_train_model, remove, axis=1)\n",
    "X_test_model = np.delete(X_test_model, remove, axis=1)\n",
    "poly_features = np.delete(poly_features, remove)\n",
    "# print(poly_features)\n",
    "\n",
    "lin_reg = LinearRegression(fit_intercept=True)\n",
    "# lin_reg = Lasso(alpha=0.00001)\n",
    "# lin_reg = LassoCV()\n",
    "# lin_reg = ElasticNet()\n",
    "# lin_reg = Ridge(alpha=0)\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features, lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['x_EndT_abs^3']"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "koeffs = dict(zip(poly_features, lin_reg.coef_))\n",
    "\n",
    "remove = []\n",
    "for itr in koeffs.items():\n",
    "    if abs(itr[1]) < 1e-4:\n",
    "        remove.append(itr[0])\n",
    "remove"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['dSlope_xEndT_abs',\n",
       " 'x_EndT_abs',\n",
       " 'ideal_state_770_tx^2',\n",
       " 'ideal_state_770_tx dSlope_xEndT',\n",
       " 'dSlope_xEndT^2',\n",
       " 'x_EndT_abs^2',\n",
       " 'ideal_state_770_ty^2',\n",
       " 'ideal_state_770_tx^2 x_EndT_abs',\n",
       " 'ideal_state_770_tx dSlope_xEndT x_EndT_abs',\n",
       " 'dSlope_xEndT^2 x_EndT_abs',\n",
       " 'dSlope_xEndT_abs ideal_state_770_ty^2',\n",
       " 'x_EndT_abs^3']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ls_koeffs = []\n",
    "for itr in koeffs.items():\n",
    "    ls_koeffs.append(itr[0])\n",
    "ls_koeffs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 88640 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(5150, 5700, 50)\n",
    "ax = sns.regplot(\n",
    "    x=y_test,\n",
    "    y=y_test - y_pred_test,\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax.set_ylim(-150, 175)\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    y_test,\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "ax.set_ylabel(\n",
    "    r\"$\\left\\langle z_{\\mathrm{Mag}}-z_{\\mathrm{Mag}}^{\\mathrm{pred}}\\right\\rangle$ [mm]\"\n",
    ")\n",
    "ax2.set_ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_regression_plot.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 90604 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "bins = np.linspace(5150, 5700, 50)\n",
    "ax = sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]) -\n",
    "    ak.to_numpy(sel_array[\"match_zmag\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax.set_ylim(-150, 175)\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    ")\n",
    "ax.set_xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "ax.set_ylabel(\n",
    "    r\"$\\left\\langle z_{\\mathrm{Mag}}-z_{\\mathrm{Mag}}^{\\mathrm{old}}\\right\\rangle$ [mm]\"\n",
    ")\n",
    "ax2.set_ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_old_regression_plot.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    label=\"True\",\n",
    ")\n",
    "plt.hist(\n",
    "    ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    bins=50,\n",
    "    range=[5150, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    label=\"True\",\n",
    ")\n",
    "plt.xlabel(r\"$z_{\\mathrm{Mag}}$ [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "# plt.show()\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_old_new.pdf\",\n",
    "    format=\"PDF\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 330,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 78888 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 2 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "xlim = 1\n",
    "bins = np.linspace(-xlim, xlim, 50)\n",
    "ax = sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_xEndT\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(\n",
    "    ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    bins=bins,\n",
    "    color=\"#FB8F67\",\n",
    "    alpha=0.5,\n",
    "    align=\"mid\",\n",
    "    density=True,\n",
    "    range=[-xlim, xlim],\n",
    ")\n",
    "ax.set_xlabel(r\"$\\Delta t_x$\")\n",
    "ax.set_ylabel(r\"$\\left\\langle z_{\\mathrm{Mag}}\\right\\rangle$ [mm]\")\n",
    "ax2.set_ylabel(f\"Number of Tracks (normalised)\", labelpad=15.0)\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "# plt.show()\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/parameterisations/plots/magnet_kink_zmag_deltatx_dist.pdf\",\n",
    "    format=\"PDF\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "plt.hist(\n",
    "    sel_array[\"z_mag_corr\"],\n",
    "    bins=100,\n",
    "    # range=[5100, 5700],\n",
    "    color=\"#2A9D8F\",\n",
    "    density=True,\n",
    ")\n",
    "plt.xlabel(r\"z$_{Mag}$ correction [mm]\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.4, 0.4, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_tx\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(5100, 5700)\n",
    "plt.xlabel(\"dx/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-0.25, 0.25, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_9410_ty\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "# plt.ylim(4500, 5700)\n",
    "\n",
    "plt.xlabel(\"dy/dz(T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-300, 300, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"ideal_state_770_x\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"x_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"x (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"y_EndVelo_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (VELO)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(0, 500, 50)\n",
    "sns.regplot(\n",
    "    # x=ak.to_numpy(sel_array[\"ideal_state_770_y\"]),\n",
    "    x=ak.to_numpy(sel_array[\"y_EndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"y (T)\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_xEndT\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_x$\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-1.0, 1.0, 50)\n",
    "sns.regplot(\n",
    "    x=ak.to_numpy(sel_array[\"dSlope_yEndT_abs\"]),\n",
    "    y=ak.to_numpy(sel_array[\"z_mag_corr\"]),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    ")\n",
    "plt.xlabel(\"$\\Delta t_y$\")\n",
    "plt.ylabel(\"$z_{Mag}$ correction [mm]\")\n",
    "mplhep.lhcb.text(\"Simulation\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "features = [\n",
    "    # # \"ideal_state_770_tx\",\n",
    "    # \"ideal_state_770_ty\",\n",
    "    # \"ideal_state_9410_ty\",\n",
    "    # # \"dSlope_yEndT\",\n",
    "    # # \"dSlope_yEndT_abs\",\n",
    "    # # \"ideal_state_770_x\",\n",
    "    # # \"ideal_state_9410_x\",\n",
    "    # # \"x_EndVelo_abs\",\n",
    "    # \"y_EndVelo_abs\",\n",
    "    # # \"y_EndT_abs\",\n",
    "    # \"dSlope_xEndT\",\n",
    "    # \"dSlope_xEndT_abs\",\n",
    "    # \"x_EndT_abs\",\n",
    "    \"ideal_state_770_tx\",\n",
    "    \"ideal_state_770_ty\",\n",
    "    \"ideal_state_9410_ty\",\n",
    "    \"ideal_state_9410_tx\",\n",
    "    \"dSlope_xEndT\",\n",
    "    \"dSlope_xEndT_abs\",\n",
    "    \"x_EndT_abs\",\n",
    "    \"ideal_state_770_x\",\n",
    "    \"x_EndVelo_abs\",\n",
    "]\n",
    "\n",
    "target_feat = \"z_mag_corr\"\n",
    "order = 4\n",
    "\n",
    "data = np.column_stack([ak.to_numpy(sel_array[feat]) for feat in features])\n",
    "target = ak.to_numpy(sel_array[target_feat])\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    data, target, test_size=0.1, random_state=42\n",
    ")\n",
    "\n",
    "poly = PolynomialFeatures(degree=order, include_bias=True)\n",
    "X_train_model = poly.fit_transform(X_train)\n",
    "X_test_model = poly.fit_transform(X_test)\n",
    "\n",
    "poly_features = poly.get_feature_names_out(input_features=features)\n",
    "\n",
    "# keep = [\n",
    "#     \"ideal_state_770_tx^2\",\n",
    "#     \"ideal_state_770_x\",\n",
    "#     # \"x_EndVelo_abs\",\n",
    "#     \"dSlope_xEndT^2\",\n",
    "#     \"dSlope_xEndT_abs\",\n",
    "#     \"x_EndT_abs\",\n",
    "# ]\n",
    "# remove = [i for i, f in enumerate(poly_features) if f not in keep]\n",
    "# remove_term = [\n",
    "#     \"ideal_state_770_x^2 ideal_state_9410_x\",\n",
    "# ]\n",
    "# remove = [i for i, f in enumerate(poly_features) if f in remove_term]\n",
    "\n",
    "# X_train_model = np.delete(X_train_model, remove, axis=1)\n",
    "# X_test_model = np.delete(X_test_model, remove, axis=1)\n",
    "# poly_features = np.delete(poly_features, remove)\n",
    "# print(poly_features)\n",
    "\n",
    "# lin_reg = LinearRegression()\n",
    "lin_reg = Lasso(alpha=0.1, max_iter=2000)\n",
    "# lin_reg = LassoCV()\n",
    "# lin_reg = ElasticNet(alpha=1)\n",
    "# lin_reg = Ridge(alpha=1)\n",
    "lin_reg.fit(X_train_model, y_train)\n",
    "y_pred_test = lin_reg.predict(X_test_model)\n",
    "print(\"intercept=\", lin_reg.intercept_)\n",
    "print(\"coef=\", dict(zip(poly_features, lin_reg.coef_)))\n",
    "print(\"r2 score=\", lin_reg.score(X_test_model, y_test))\n",
    "print(\"RMSE =\", mean_squared_error(y_test, y_pred_test, squared=False))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bins = np.linspace(-200, 200, 50)\n",
    "ax = sns.regplot(\n",
    "    x=y_test,\n",
    "    y=abs(y_test - y_pred_test),\n",
    "    x_bins=bins,\n",
    "    fit_reg=None,\n",
    "    x_estimator=np.mean,\n",
    "    label=\"bla\",\n",
    ")\n",
    "ax2 = ax.twinx()\n",
    "ax2.hist(y_test, bins=30, range=[-200, 200], color=\"#2A9D8F\", alpha=0.8, align=\"left\")\n",
    "ax.set_xlabel(r\"z$_{Mag}$ correction [mm]\")\n",
    "ax.set_ylabel(\"Mean Deviation [mm]\")\n",
    "ax2.set_ylabel(\"Number of Tracks\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
