feat: Add ratio plot support

examples/dev-example.ipynb

@@ -7,6 +7,13 @@
     "# Development examples"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Style setup"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -28,6 +35,13 @@
     ")  # Default values"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Stacked histogram plots"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -59,13 +73,15 @@
    "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);"
    ]
@@ -76,7 +92,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "ax.figure.savefig(\"example_stack.png\")"
+    "fig.savefig(\"example_stack.png\")"
    ]
   },
   {
@@ -141,6 +157,13 @@
     ")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "## Shape plots"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -212,6 +235,68 @@
     ")"
    ]
   },
+  {
+   "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, 12],\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,

setup.cfg

@@ -39,7 +39,7 @@ install_requires =
     click>=6.0
     awkward>=1.0
     uproot>=4.0
-    hist[plot]>=2.2.0
+    hist[plot]>=2.3.0
     mplhep>=0.2.18
     uproot3>=3.14 # Needed until writing added in uproot4
 

src/heputils/plot.py

@@ -495,3 +495,69 @@ def stack_hist(hists, ax=None, **kwargs):
     ax = draw_experiment_label(ax, max_height=max_hist, **kwargs)
 
     return _plot_ax_kwargs(ax, **kwargs)
+
+
+def stack_ratio_plot(hists, **kwargs):
+    """
+    Stack plot on top, ratio plot on bottom
+    """
+    fig = kwargs.pop("fig", plt.gcf())
+
+    # Scale figure height to deal with ratio subplot being added
+    fig_height_scale = kwargs.pop("fig_height_scale", 1.25)
+    _fig_width, _fig_height = get_style()["figure.figsize"]
+    fig.set_size_inches(_fig_width, _fig_height * fig_height_scale, forward=True)
+
+    # Setup figure subplot grid and axis dict for hist.plot_ratio
+    grid = fig.add_gridspec(2, 1, hspace=0, height_ratios=[3, 1])
+    main_ax = fig.add_subplot(grid[0])
+    subplot_ax = fig.add_subplot(grid[1], sharex=main_ax)
+    ax_dict = {"main_ax": main_ax, "ratio_ax": subplot_ax}
+
+    labels = kwargs.pop("labels", None)
+    color = kwargs.pop("color", None)
+    if color is not None and len(color) != len(hists):
+        color = color[: len(hists)]
+    alpha = kwargs.pop("alpha", None)
+    semilogy = kwargs.pop("logy", True)
+    _data_hist = kwargs.pop("data_hist", None)
+
+    # Setup and plot the ratio plot
+    ratio_plot_numerator = kwargs.pop("ratio_numerator", "data")
+    ratio_plot_kwargs = {
+        "ax_dict": ax_dict,
+        "rp_ylim": kwargs.pop("rp_ylim", None),
+        "rp_uncertainty_type": kwargs.pop("rp_uncertainty_type", "poisson"),
+        "rp_uncert_draw_type": kwargs.pop("rp_uncert_draw_type", "line"),
+    }
+
+    num_hists = utils.sum_hists(hists)
+    if ratio_plot_numerator.lower() in ["simulation", "sim", "mc"]:
+        ratio_plot_kwargs["rp_ylabel"] = kwargs.pop("rp_ylabel", "MC/Data")
+        num_hists.plot_ratio(_data_hist, **ratio_plot_kwargs)
+    else:
+        ratio_plot_kwargs["rp_ylabel"] = kwargs.pop("rp_ylabel", "Data/MC")
+        _data_hist.plot_ratio(num_hists, **ratio_plot_kwargs)
+    subplot_ax.set_xlabel(kwargs.get("xlabel", None))
+
+    # FIXME: Ugly hack to overwrite ratio plot main axis
+    main_ax.clear()
+    main_ax = stack_hist(
+        hists,
+        labels=labels,
+        color=color,
+        alpha=alpha,
+        logy=semilogy,
+        data_hist=_data_hist,
+        ax=main_ax,
+        **kwargs,
+    )
+
+    # Hide tick marks of main_ax
+    plt.setp(main_ax.get_xticklabels(), visible=False)
+    # Trying to get things looking okay
+    if semilogy and fig_height_scale < 1.25:
+        # Ensure enough space for legend
+        main_ax.set_ylim(top=main_ax.get_ylim()[-1] * 100)
+
+    return main_ax, subplot_ax