{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Development examples"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Style setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "import uproot\n",
    "\n",
    "import heputils\n",
    "\n",
    "heputils.plot.set_style(\"ATLAS\")\n",
    "heputils.plot.set_experiment_info(status=\"Internal\")  # Internal, Preliminary, Public\n",
    "heputils.plot.set_experiment_info(\n",
    "    center_of_mass_energy=13, center_of_mass_energy_units=\"TeV\"\n",
    ")  # Default values\n",
    "heputils.plot.set_experiment_info(\n",
    "    luminosity=132, luminosity_units=\"fb\"\n",
    ")  # Default values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Stacked histogram plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root_file = uproot.open(\"example.root\")\n",
    "root_hists = [root_file[key].to_hist() for key in root_file.keys()]\n",
    "data_hist = root_hists[-1]\n",
    "simulation_hists = root_hists[:-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "labels = [r\"$t\\bar{t}$\", r\"$W+$jets\", \"Other\", \"Signal\"]\n",
    "colormap = [\"darkorange\", \"steelblue\", \"firebrick\", \"purple\"]\n",
    "\n",
    "scale_factors = np.ones(len(simulation_hists))\n",
    "scale_factors[-1] = 5"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "ax = heputils.plot.stack_hist(\n",
    "    simulation_hists,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    scale_factors=scale_factors,\n",
    "    ax=ax,\n",
    ")\n",
    "ax = heputils.plot.data_hist(data_hist, ax=ax);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig.savefig(\"example_stack.png\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can also give the data histogram directly to `heputils.plot.stack_hist`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ax, artists = heputils.plot.stack_hist(\n",
    "    simulation_hists,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    data_hist=data_hist,\n",
    "    logy=False,\n",
    "    return_artists=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Can easily plot on a `subplots` grid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_fig_size = heputils.plot.get_style()[\"figure.figsize\"]\n",
    "fig, axs = plt.subplots(1, 2, figsize=(2.1 * _fig_size[0], _fig_size[1]))\n",
    "\n",
    "axs[0] = heputils.plot.stack_hist(\n",
    "    simulation_hists,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    data_hist=data_hist,\n",
    "    ax=axs[0],\n",
    ")\n",
    "axs[1] = heputils.plot.stack_hist(\n",
    "    simulation_hists,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    data_hist=data_hist,\n",
    "    logy=False,\n",
    "    ax=axs[1],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Shape plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_fig_size = heputils.plot.get_style()[\"figure.figsize\"]\n",
    "fig, axs = plt.subplots(1, 2, figsize=(2.1 * _fig_size[0], _fig_size[1]))\n",
    "\n",
    "axs[0] = heputils.plot.shape_hist(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    ax=axs[0],\n",
    ")\n",
    "axs[1] = heputils.plot.shape_hist(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    logy=True,\n",
    "    density=False,\n",
    "    ax=axs[1],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "For a shape plot you can also easily switch between a `fill` style (shaded) and a `step` style (lines) using the `histtype` kwarg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_fig_size = heputils.plot.get_style()[\"figure.figsize\"]\n",
    "fig, axs = plt.subplots(1, 2, figsize=(2.1 * _fig_size[0], _fig_size[1]))\n",
    "\n",
    "axs[0] = heputils.plot.shape_hist(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    histtype=\"step\",\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    ax=axs[0],\n",
    ")\n",
    "axs[1] = heputils.plot.shape_hist(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    histtype=\"step\",\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    logy=True,\n",
    "    density=False,\n",
    "    ax=axs[1],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Ratio plots"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "axs = heputils.plot.stack_ratio_plot(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    ratio_numerator=\"data\",  # \"data\", \"simulation\", or \"mc\"\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    rp_uncert_draw_type=\"line\",  # \"line\" or \"bar\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig.savefig(\"example_stack_ratio.png\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "axs = heputils.plot.stack_ratio_plot(\n",
    "    simulation_hists,\n",
    "    data_hist=data_hist,\n",
    "    labels=labels,\n",
    "    color=colormap,\n",
    "    rp_ylim=[-1, 8],\n",
    "    xlabel=r\"$X$ Mass [GeV]\",\n",
    "    ylabel=\"Count\",\n",
    "    logy=False,\n",
    "    rp_uncert_draw_type=\"bar\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Jupyter repr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked_hist = heputils.utils.sum_hists(simulation_hists)\n",
    "stacked_hist.plot();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "stacked_hist"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}