{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "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",
    "\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "plt.rcParams[\"savefig.dpi\"] = 600\n",
    "# %matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tree = uproot.open({\n",
    "    \"/work/cetin/LHCb/reco_tuner/nn_electron_training/result_bestprecuts_NozMagCut/matching_ghost_mlp_training.root\":\n",
    "    \"MatchNNDataSet/TrainTree\"\n",
    "})\n",
    "test_tree = uproot.open({\n",
    "    \"/work/cetin/LHCb/reco_tuner/nn_electron_training/result_bestprecuts_NozMagCut/matching_ghost_mlp_training.root\":\n",
    "    \"MatchNNDataSet/TestTree\"\n",
    "})\n",
    "train_array = train_tree.arrays()\n",
    "test_array = test_tree.arrays()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 87696 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "train_bkg = train_array[train_array.classID == 1]\n",
    "train_sig = train_array[train_array.classID == 0]\n",
    "test_bkg = test_array[test_array.classID == 1]\n",
    "test_sig = test_array[test_array.classID == 0]\n",
    "plt.hist(\n",
    "    train_sig.matching_mlp,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "plt.hist(\n",
    "    train_bkg.matching_mlp,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "mplhep.histplot(\n",
    "    np.histogram(np.array(test_sig.matching_mlp), 50),\n",
    "    histtype=\"errorbar\",\n",
    "    density=True,\n",
    "    yerr=True,\n",
    "    color=\"#107E7D\",\n",
    "    marker=\"^\",\n",
    "    markersize=7,\n",
    "    label=\"test sample, true pairs\",\n",
    ")\n",
    "mplhep.histplot(\n",
    "    np.histogram(np.array(test_bkg.matching_mlp), 50),\n",
    "    histtype=\"errorbar\",\n",
    "    density=True,\n",
    "    yerr=True,\n",
    "    color=\"#F05342\",\n",
    "    label=\"test sample, wrong pairs\",\n",
    ")\n",
    "plt.xlabel(\"neural network response\")\n",
    "plt.ylabel(\"Number of Tracks (normalised)\")\n",
    "mplhep.lhcb.text(\"Simulation\", loc=0)\n",
    "plt.legend(loc=\"upper center\")\n",
    "# plt.savefig(\n",
    "#     \"/work/cetin/LHCb/reco_tuner/thesis/new_electron_NN_response.pdf\",\n",
    "#     format=\"PDF\",\n",
    "# )\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The cut is set to 0.5, thereby rejecting 0.901 of fake tracks, while keeping 0.935 of the true matches.\n"
     ]
    }
   ],
   "source": [
    "response = 0.5\n",
    "\n",
    "n_sig_keep = ak.num(train_sig[train_sig.matching_mlp >= response], axis=0)\n",
    "n_bkg_rej = ak.num(train_bkg[train_bkg.matching_mlp < response], axis=0)\n",
    "nevents = 115e3\n",
    "\n",
    "print(\n",
    "    f\"The cut is set to {response}, thereby rejecting {np.round(n_bkg_rej/nevents,3)} of fake tracks, while keeping {np.round(n_sig_keep/nevents,3)} of the true matches.\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 105024 characters]",
      "text/plain": [
       "<Figure size 1400x1500 with 6 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "fig, axes = plt.subplots(3, 2, figsize=(14, 15), sharey=False)\n",
    "# 0,0\n",
    "axes[0, 0].hist(\n",
    "    train_sig.chi2,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[0, 0].hist(\n",
    "    train_bkg.chi2,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[0, 0].set_xlabel(r\"$\\chi^{2}_{\\mathrm{match}}$\")\n",
    "axes[0, 0].legend(prop={\"size\": 20})\n",
    "# 1,0\n",
    "axes[1, 0].hist(\n",
    "    train_sig.distX,\n",
    "    bins=50,\n",
    "    range=(0, 150),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[1, 0].hist(\n",
    "    train_bkg.distX,\n",
    "    bins=50,\n",
    "    range=(0, 150),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[1, 0].set_xlabel(r\"$D_{x}$ [mm]\")\n",
    "axes[1, 0].set_ylabel(\"Number of Tracks (normalised)\",\n",
    "                      va=\"bottom\",\n",
    "                      ha=\"center\")\n",
    "# 0,1\n",
    "axes[0, 1].hist(\n",
    "    train_sig.teta2,\n",
    "    bins=50,\n",
    "    range=(0.0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[0, 1].hist(\n",
    "    train_bkg.teta2,\n",
    "    bins=50,\n",
    "    range=(0.0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[0, 1].set_xlabel(r\"$t_{x}^{2}+t_{y}^{2}$\")\n",
    "# 1,1\n",
    "axes[1, 1].hist(\n",
    "    train_sig.distY,\n",
    "    bins=50,\n",
    "    range=(0, 150),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[1, 1].hist(\n",
    "    train_bkg.distY,\n",
    "    bins=50,\n",
    "    range=(0, 150),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[1, 1].set_xlabel(r\"$D_{y}$ [mm]\")\n",
    "# 2,0\n",
    "axes[2, 0].hist(\n",
    "    train_sig.dSlope,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[2, 0].hist(\n",
    "    train_bkg.dSlope,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[2, 0].set_xlabel(r\"$|\\Delta t_{x}^{\\mathrm{match}}|$\")\n",
    "# 2,1\n",
    "axes[2, 1].hist(\n",
    "    train_sig.dSlopeY,\n",
    "    bins=50,\n",
    "    range=(0, 0.03),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[2, 1].hist(\n",
    "    train_bkg.dSlopeY,\n",
    "    bins=50,\n",
    "    range=(0, 0.03),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[2, 1].set_xlabel(r\"$|\\Delta t_{y}^{\\mathrm{match}}|$\")\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/thesis/new_electron_NN_variables.pdf\",\n",
    "    format=\"PDF\",\n",
    ")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axes = plt.subplots(2, 3, figsize=(25, 13), sharey=False)\n",
    "# 0,0\n",
    "axes[0, 0].hist(\n",
    "    train_sig.chi2,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[0, 0].hist(\n",
    "    train_bkg.chi2,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[0, 0].set_xlabel(r\"$\\chi^{2}_{\\mathrm{match}}$\")\n",
    "axes[0, 0].legend(prop={\"size\": 20})\n",
    "\n",
    "# 0,1\n",
    "axes[0, 1].hist(\n",
    "    train_sig.teta2,\n",
    "    bins=50,\n",
    "    range=(0.0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[0, 1].hist(\n",
    "    train_bkg.teta2,\n",
    "    bins=50,\n",
    "    range=(0.0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[0, 1].set_xlabel(r\"$t_{x}^{2}+t_{y}^{2}$\")\n",
    "# 0,2\n",
    "axes[0, 2].hist(\n",
    "    train_sig.distX,\n",
    "    bins=50,\n",
    "    range=(0, 100),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[0, 2].hist(\n",
    "    train_bkg.distX,\n",
    "    bins=50,\n",
    "    range=(0, 100),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[0, 2].set_xlabel(r\"$D_{x}$ [mm]\")\n",
    "axes[0, 0].set_ylabel(\"Number of tracks (normalised)\", va=\"bottom\", ha=\"center\")\n",
    "# 1,0\n",
    "axes[1, 0].hist(\n",
    "    train_sig.distY,\n",
    "    bins=50,\n",
    "    range=(0, 100),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[1, 0].hist(\n",
    "    train_bkg.distY,\n",
    "    bins=50,\n",
    "    range=(0, 100),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[1, 0].set_xlabel(r\"$D_{y}$ [mm]\")\n",
    "# 2,0\n",
    "axes[1, 1].hist(\n",
    "    train_sig.dSlope,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[1, 1].hist(\n",
    "    train_bkg.dSlope,\n",
    "    bins=50,\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[1, 1].set_xlabel(r\"$|\\Delta t_{x}^{\\mathrm{match}}|$\")\n",
    "# 2,1\n",
    "axes[1, 2].hist(\n",
    "    train_sig.dSlopeY,\n",
    "    bins=50,\n",
    "    range=(0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#107E7D\",\n",
    "    label=\"training sample, true pairs\",\n",
    ")\n",
    "axes[1, 2].hist(\n",
    "    train_bkg.dSlopeY,\n",
    "    bins=50,\n",
    "    range=(0, 0.02),\n",
    "    alpha=0.5,\n",
    "    density=True,\n",
    "    log=False,\n",
    "    color=\"#F05342\",\n",
    "    label=\"training sample, wrong pairs\",\n",
    ")\n",
    "axes[1, 2].set_xlabel(r\"$|\\Delta t_{y}^{\\mathrm{match}}|$\")\n",
    "plt.savefig(\n",
    "    \"/work/cetin/LHCb/reco_tuner/thesis/newparams_filtered_NN_elec_variables_landscape.pdf\",\n",
    "    format=\"PDF\",\n",
    ")\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": []
  },
  {
   "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
}
