{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import uproot\n",
    "input_tree_md = uproot.open({\"/work/guenther/reco_tuner/data/param_data_MD_selected.root\": \"Selected\"})\n",
    "input_tree_mu = uproot.open({\"/work/guenther/reco_tuner/data/param_data_MU_selected.root\": \"Selected\"})\n",
    "    # this is an event list of dictionaries containing awkward arrays\n",
    "array_md = input_tree_md.arrays()\n",
    "array_mu = input_tree_mu.arrays()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "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": "[encoded data, 55100 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57964 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 55708 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 59188 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 57364 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 64032 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 60276 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 62208 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "image/png": "[encoded data, 52672 characters]",
      "text/plain": [
       "<Figure size 1200x900 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import mplhep\n",
    "mplhep.style.use([\"LHCbTex2\"])\n",
    "default_args = dict(\n",
    "    alpha=0.5\n",
    ")\n",
    "\n",
    "variables = {\n",
    "    \"z_mag_x\": dict(\n",
    "        range=(5100, 5600),\n",
    "        label=\"z_mag_x\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"z_mag_y\": dict(\n",
    "        range=(2000, 10000),\n",
    "        label=\"z_mag_y\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"AX\": dict(\n",
    "        range=(-3000, 3000),\n",
    "        label=\"AX\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"BX\": dict(\n",
    "        range=(-1.2, 1.2),\n",
    "        label=\"BX\",\n",
    "        unit=\"1/mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"CX\": dict(\n",
    "        range=(-7e-6, 7e-6),\n",
    "        label=\"CX\",\n",
    "        unit=\"1/mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"DX\": dict(\n",
    "        range=(-7e-9, 7e-9),\n",
    "        label=\"DX\",\n",
    "        unit=\"1/mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"AY\": dict(\n",
    "        range=(-2500, 2500),\n",
    "        label=\"AY\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"BY\": dict(\n",
    "        range=(-0.4, 0.4),\n",
    "        label=\"BY\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "    \"CY\": dict(\n",
    "        range=(-4e-6, 4e-6),\n",
    "        label=\"CY\",\n",
    "        unit=\"mm\",\n",
    "        bins=50,\n",
    "    ),\n",
    "}\n",
    "for name in variables.keys():\n",
    "    plt.figure()\n",
    "    n, bins, patches = plt.hist(\n",
    "        array_md[name],\n",
    "        bins=variables[name][\"bins\"],\n",
    "        **default_args,\n",
    "        range=variables[name][\"range\"],\n",
    "        label=\"MagDown\"\n",
    "    )\n",
    "    plt.hist(array_mu[name], bins=variables[name][\"bins\"],\n",
    "        **default_args, \n",
    "        range=variables[name][\"range\"], label=\"MagUp\")\n",
    "\n",
    "    label = variables[name][\"label\"]\n",
    "    unit = variables[name][\"unit\"]\n",
    "    bin_width = \"%.1f\" % round(bins[1] - bins[0], 1)\n",
    "    plt.xlabel(f\"{label} [{unit}]\")\n",
    "    plt.ylabel(f\"Particles/({bin_width} {unit})\")\n",
    "    mplhep.lhcb.text(\"Simulation\")\n",
    "    plt.legend()\n",
    "\n",
    "\n"
   ]
  },
  {
   "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.6"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "a2eff8b4da8b8eebf5ee2e5f811f31a557e0a202b4d2f04f849b065340a6eda6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
