{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import mplhep\n",
    "\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "plt.rcParams[\"savefig.dpi\"] = 600\n",
    "# %matplotlib inline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_tree = uproot.open(\n",
    "    {\n",
    "        \"/work/cetin/LHCb/reco_tuner/nn_electron_training/result_NewParams_EFilter/matching_ghost_mlp_training.root\": \"MatchNNDataSet/TrainTree\"\n",
    "    }\n",
    ")\n",
    "test_tree = uproot.open(\n",
    "    {\n",
    "        \"/work/cetin/LHCb/reco_tuner/nn_electron_training/result_NewParams_EFilter/matching_ghost_mlp_training.root\": \"MatchNNDataSet/TestTree\"\n",
    "    }\n",
    ")\n",
    "train_array = train_tree.arrays()\n",
    "test_array = test_tree.arrays()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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/newparams_filtered_NN_elec_response.pdf\",\n",
    "    format=\"PDF\",\n",
    ")\n",
    "# plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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, 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.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[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, 100),\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, 100),\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.02),\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.02),\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/filtered_NN_elec_variables.pdf\",\n",
    "    format=\"PDF\")\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
}