Comparison update, poisson uncert.

plothist/hep_plotters.py

@@ -12,6 +12,7 @@
     compare_two_hist,
     plot_comparison,
     create_comparison_figure,
+    _hist_ratio_variances,
 )
 
 
@@ -84,7 +85,6 @@ def compare_data_mc(
     comparison_kwargs.setdefault("comparison", "ratio")
     comparison_kwargs.setdefault("ratio_uncertainty", "split")
 
-
     if fig is None and ax_main is None and ax_comparison is None:
         fig, (ax_main, ax_comparison) = create_comparison_figure()
     elif fig is None or ax_main is None or ax_comparison is None:
@@ -121,6 +121,29 @@ def compare_data_mc(
             uncertainties_high**2,
         )
         data_hist[:] = np.stack([data_hist.values(), data_variances], axis=-1)
+    elif comparison_kwargs["comparison"] == "ratio":
+        if comparison_kwargs["ratio_uncertainty"] == "split":
+            np.seterr(divide="ignore", invalid="ignore")
+            # Compute asymmetrical uncertainties to plot_comparison()
+            comparison_kwargs["yerr"] = [
+                uncertainties_low / mc_hist_total.values(),
+                uncertainties_high / mc_hist_total.values(),
+            ]
+            np.seterr(divide="warn", invalid="warn")
+        elif comparison_kwargs["ratio_uncertainty"] == "uncorrelated":
+            data_hist_high = data_hist.copy()
+            data_hist_high[:] = np.stack(
+                [data_hist_high.values(), uncertainties_high**2], axis=-1
+            )
+            data_hist_low = data_hist.copy()
+            data_hist_low[:] = np.stack(
+                [data_hist_low.values(), uncertainties_low**2], axis=-1
+            )
+            # Compute asymmetrical uncertainties to plot_comparison()
+            comparison_kwargs["yerr"] = [
+                np.sqrt(_hist_ratio_variances(data_hist_low, mc_hist_total)),
+                np.sqrt(_hist_ratio_variances(data_hist_high, mc_hist_total)),
+            ]
 
     plot_error_hist(
         data_hist,
@@ -163,6 +186,9 @@ def compare_data_mc(
 
 
 def _get_poisson_uncertainties(data_hist):
+    """
+    Get Poisson asymmetrical uncertainties for a histogram.
+    """
     conf_level = 0.682689492
     alpha = 1.0 - conf_level
     n = data_hist.values()

plothist/plotters.py

@@ -416,6 +416,42 @@ def compare_two_hist(
     return fig, ax_main, ax_comparison
 
 
+def _hist_ratio_variances(hist_1, hist_2):
+    """
+    Calculate the variances of the ratio of two histograms (hist_1/hist_2).
+
+    Parameters
+    ----------
+    hist_1 : boost_histogram.Histogram
+        The first histogram.
+    hist_2 : boost_histogram.Histogram
+        The second histogram.
+
+    Returns
+    -------
+    variances : np.ndarray
+        The variances of the ratio of the two histograms.
+
+    Raises
+    ------
+    ValueError
+        If the bins of the histograms are not equal.
+    """
+    if not np.all(hist_1.axes[0].edges == hist_2.axes[0].edges):
+        raise ValueError("The bins of the histograms must be equal.")
+
+    np.seterr(divide="ignore", invalid="ignore")
+    ratio_variances = np.where(
+        hist_2.values() != 0,
+        hist_1.variances() / hist_2.values() ** 2
+        + hist_2.variances() * hist_1.values() ** 2 / hist_2.values() ** 4,
+        np.nan,
+    )
+    np.seterr(divide="warn", invalid="warn")
+
+    return ratio_variances
+
+
 def plot_comparison(
     hist_1,
     hist_2,
@@ -427,6 +463,7 @@ def plot_comparison(
     comparison_ylabel=None,
     comparison_ylim=None,
     ratio_uncertainty="uncorrelated",
+    **plot_hist_kwargs,
 ):
     """
     Plot the comparison between two histograms.
@@ -453,6 +490,8 @@ def plot_comparison(
         The y-axis limits for the comparison plot. Default is None.
     ratio_uncertainty : str, optional
         How to treat the uncertainties of the histograms when comparison = "ratio" ("uncorrelated" for simple comparison, "split" for scaling and split hist_1 and hist_2 uncertainties). This argument has no effect if comparison != "ratio". Default is "uncorrelated".
+    **plot_hist_kwargs : optional
+        Arguments to be passed to plot_hist() or plot_error_hist(), called in case the comparison is "pull" or "ratio", respectively.
 
     Returns
     -------
@@ -479,12 +518,7 @@ def plot_comparison(
                 np.nan,
             )
         elif ratio_uncertainty == "uncorrelated":
-            ratio_variances = np.where(
-                hist_2.values() != 0,
-                hist_1.variances() / hist_2.values() ** 2
-                + hist_2.variances() * hist_1.values() ** 2 / hist_2.values() ** 4,
-                np.nan,
-            )
+            ratio_variances = _hist_ratio_variances(hist_1, hist_2)
         else:
             raise ValueError("ratio_uncertainty not in ['uncorrelated', 'split'].")
         if ratio_uncertainty == "split":
@@ -512,16 +546,19 @@ def plot_comparison(
     comparison_variances = (
         ratio_variances
         if (ratio_uncertainty == "uncorrelated" or comparison == "pull")
-        else h1_scaled_uncertainties ** 2
+        else h1_scaled_uncertainties**2
     )
 
     hist_comparison = bh.Histogram(hist_2.axes[0], storage=bh.storage.Weight())
     hist_comparison[:] = np.stack([comparison_values, comparison_variances], axis=-1)
 
     if comparison == "pull":
-        plot_hist(hist_comparison, ax=ax, histtype="stepfilled", color="darkgrey")
+        plot_hist_kwargs.setdefault("histtype", "stepfilled")
+        plot_hist_kwargs.setdefault("color", "darkgrey")
+        plot_hist(hist_comparison, ax=ax, **plot_hist_kwargs)
     else:
-        plot_error_hist(hist_comparison, ax=ax, color="black")
+        plot_hist_kwargs.setdefault("color", "black")
+        plot_error_hist(hist_comparison, ax=ax, **plot_hist_kwargs)
 
     if comparison == "ratio":
         if comparison_ylim is None:
@@ -614,7 +651,7 @@ def f(x0, x1):
         def color(lambda_):
             # emphasise either low intensity values (gamma < 1),
             # or high intensity values (γ > 1)
-            lambda_gamma = lambda_ ** gamma
+            lambda_gamma = lambda_**gamma
 
             # Angle and amplitude for the deviation
             # from the black to white diagonal