{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import mplhep\n",
    "import awkward as ak\n",
    "import pandas as pd\n",
    "import seaborn as sns\n",
    "\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "plt.rcParams[\"savefig.dpi\"] = 600\n",
    "# %matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "tree = uproot.open({\n",
    "    \"/work/cetin/LHCb/reco_tuner/data_matching/parameterisations/sample4/param_data_B_default_thesis_selected.root\":\n",
    "    \"Selected\"\n",
    "})\n",
    "allcols = tree.arrays()\n",
    "allcols = allcols[allcols.isElectron == 1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "zMagnetParamsEndT = [\n",
    "    5283.043384741342,\n",
    "    28.45437483916609,\n",
    "    0.016308934185688337,\n",
    "    -1668.3602223088105,\n",
    "    301.07580662497,\n",
    "]\n",
    "dSlope_xEndT = allcols.tx_l11 - allcols.tx\n",
    "dSlope_xEndT_abs = abs(dSlope_xEndT)\n",
    "x_EndT_abs = abs(allcols.x_l11 + allcols.tx_l11 * (9410.0 - allcols.z_l11))\n",
    "\n",
    "zMagTerm = zMagnetParamsEndT[0] + zMagnetParamsEndT[3] * allcols.tx * allcols.tx\n",
    "\n",
    "z_mag_xEndT = (\n",
    "    zMagTerm + dSlope_xEndT_abs *\n",
    "    (zMagnetParamsEndT[1] + zMagnetParamsEndT[4] * dSlope_xEndT_abs) +\n",
    "    zMagnetParamsEndT[2] * x_EndT_abs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "zmag = (allcols[\"x\"] - allcols[\"x_l11\"] - allcols[\"tx\"] * allcols[\"z\"] +\n",
    "        allcols[\"tx_l11\"] * allcols[\"z_l11\"]) / dSlope_xEndT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 61660 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.hist(\n",
    "    zmag,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"True\",\n",
    ")\n",
    "plt.hist(\n",
    "    z_mag_xEndT,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    color=\"#F05342\",\n",
    "    label=\"Predicted\",\n",
    ")\n",
    "plt.xlabel(r\"$z_{mag}$\")\n",
    "plt.ylabel(\"Number of tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.legend(loc=\"upper right\")\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/thesis/zmag_new_parameterisation.pdf\", format=\"PDF\"\n",
    "# )\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 35,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "True"
      ]
     },
     "execution_count": 35,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "weird = zmag < 5200\n",
    "True in weird"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 40712 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "array = {\"z_mag\": zmag, \"z_mag_pred\": z_mag_xEndT, \"diff\": zmag - z_mag_xEndT}\n",
    "df = pd.DataFrame(array)\n",
    "sns.regplot(data=df, x=\"z_mag\", y=\"diff\", x_bins=30, fit_reg=False)\n",
    "plt.show()"
   ]
  },
  {
   "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": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tuner",
   "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"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
