Merge pull request #61 from uhh-cms/feature/s_over_b_plots

S over B plotting function

hbw/config/processes.py

@@ -111,3 +111,12 @@ def add_hbw_processes(config: od.Config, campaign: od.Campaign):
         )
         for proc in signal_processes:
             sig.add_process(proc)
+
+    # add auxiliary information if process is signal
+    for proc_inst in config.processes:
+        is_signal = any([
+            signal_tag in proc_inst.name
+            for signal_tag in ("qqHH", "ggHH", "radion", "gravition")
+        ])
+        if is_signal:
+            proc_inst.add_tag("is_signal")

hbw/plotting/example.py

@@ -0,0 +1,83 @@
+# coding: utf-8
+
+"""
+Examples for custom plot functions.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+from columnflow.util import maybe_import
+from columnflow.plotting.plot_util import (
+    remove_residual_axis,
+    apply_variable_settings,
+    apply_process_settings,
+)
+
+hist = maybe_import("hist")
+np = maybe_import("numpy")
+mpl = maybe_import("matplotlib")
+plt = maybe_import("matplotlib.pyplot")
+mplhep = maybe_import("mplhep")
+od = maybe_import("order")
+
+
+def my_plot1d_func(
+    hists: OrderedDict[od.Process, hist.Hist],
+    config_inst: od.Config,
+    category_inst: od.Category,
+    variable_insts: list[od.Variable],
+    style_config: dict | None = None,
+    yscale: str | None = "",
+    process_settings: dict | None = None,
+    variable_settings: dict | None = None,
+    example_param: str | float | bool | None = None,
+    **kwargs,
+) -> plt.Figure:
+    """
+    This is an exemplary custom plotting function.
+
+    Exemplary task call:
+
+    .. code-block:: bash
+        law run cf.PlotVariables1D --version v1 --processes st,tt --variables jet1_pt \
+            --plot-function ana_short.plotting.example.my_plot1d_func \
+            --general-settings example_param=some_text
+    """
+    # we can add arbitrary parameters via the `general_settings` parameter to access them in the
+    # plotting function. They are automatically parsed either to a bool, float, or string
+    print(f"The example_param has been set to '{example_param}' (type: {type(example_param)})")
+
+    # call helper function to remove shift axis from histogram
+    remove_residual_axis(hists, "shift")
+
+    # call helper functions to apply the variable_settings and process_settings
+    variable_inst = variable_insts[0]
+    hists = apply_variable_settings(hists, variable_insts, variable_settings)
+    hists = apply_process_settings(hists, process_settings)
+
+    # use the mplhep CMS stype
+    plt.style.use(mplhep.style.CMS)
+
+    # create a figure and fill it with content
+    fig, ax = plt.subplots()
+    for proc_inst, h in hists.items():
+        h.plot1d(
+            ax=ax,
+            label=proc_inst.label,
+            color=proc_inst.color1,
+        )
+
+    # styling and parameter implementation (e.g. `yscale`)
+    ax.set(
+        yscale=yscale,
+        ylabel=variable_inst.get_full_y_title(),
+        xlabel=variable_inst.get_full_x_title(),
+        xscale="log" if variable_inst.log_x else "linear",
+    )
+    ax.legend()
+    mplhep.cms.label(ax=ax, fontsize=22, llabel="private work")
+
+    # task expects a figure and a tuple of axes as output
+    return fig, (ax,)

hbw/plotting/s_over_b.py

@@ -0,0 +1,126 @@
+# coding: utf-8
+
+"""
+Examples for custom plot functions.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+
+import law
+
+from columnflow.util import maybe_import
+from columnflow.plotting.plot_all import plot_all
+from columnflow.plotting.plot_util import (
+    # prepare_plot_config,
+    prepare_style_config,
+    remove_residual_axis,
+    apply_variable_settings,
+    apply_process_settings,
+    apply_density_to_hists,
+)
+
+from hbw.util import round_sig
+
+hist = maybe_import("hist")
+np = maybe_import("numpy")
+mpl = maybe_import("matplotlib")
+plt = maybe_import("matplotlib.pyplot")
+mplhep = maybe_import("mplhep")
+od = maybe_import("order")
+
+logger = law.logger.get_logger(__name__)
+
+
+def plot_s_over_b(
+    hists: OrderedDict,
+    config_inst: od.Config,
+    category_inst: od.Category,
+    variable_insts: list[od.Variable],
+    style_config: dict | None = None,
+    density: bool | None = False,
+    shape_norm: bool | None = False,
+    yscale: str | None = "",
+    hide_errors: bool | None = None,
+    sqrt_b: bool | None = None,
+    process_settings: dict | None = None,
+    variable_settings: dict | None = None,
+    **kwargs,
+) -> plt.Figure:
+    """
+    Plotting function to create a single line presenting signal vs background ratio.
+    Some plot parameters might be supported but have not been tested.
+    Exemplary task call:
+
+    .. code-block:: bash
+        law run cf.PlotVariables1D --version prod1 \
+            --processes ggHH_kl_1_kt_1_sl_hbbhww,tt_sl --variables jet1_pt \
+            --plot-function hbw.plotting.s_over_b.plot_s_over_b \
+            --general-settings sqrt_b
+    """
+    remove_residual_axis(hists, "shift")
+
+    variable_inst = variable_insts[0]
+    hists = apply_variable_settings(hists, variable_insts, variable_settings)
+    hists = apply_process_settings(hists, process_settings)
+    hists = apply_density_to_hists(hists, density)
+
+    # separate histograms into signal and background based on process tag 'is_signal'
+    h_sig, h_bkg = None, None
+    for proc_inst, h in hists.items():
+        if proc_inst.has_tag("is_signal"):
+            h_sig = h + h_sig if h_sig else h
+        else:
+            h_bkg = h + h_bkg if h_bkg else h
+
+    if not h_sig:
+        raise Exception(
+            "No signal processes given. Remember to add the 'is_signal' auxiliary to your "
+            "signal processes",
+        )
+    if not h_bkg:
+        raise Exception("No background processes given")
+
+    # TODO: include uncertainties
+    S_over_B = h_sig[::sum].value / h_bkg[::sum].value
+    S_over_sqrtB = h_sig[::sum].value / np.sqrt(h_bkg[::sum].value)
+    logger.info(
+        f"\n    Integrated S over B:     {round_sig(S_over_B)}" +
+        f"\n    Integrated S over sqrtB: {round_sig(S_over_sqrtB)}",
+    )
+
+    # NOTE: this does not take into account the variances on the background stack
+    if sqrt_b:
+        h_out = h_sig / np.sqrt(h_bkg.values())
+        label = "S over sqrt(B)"
+    else:
+        h_out = h_sig / h_bkg.values()
+        label = "S over B"
+
+    # draw lines
+    plot_config = {}
+    line_norm = sum(h_out.values()) if shape_norm else 1
+    plot_config["s_over_b"] = {
+        "method": "draw_hist",
+        "hist": h_out,
+        "kwargs": {
+            "norm": line_norm,
+            "label": label,
+        },
+    }
+
+    default_style_config = prepare_style_config(
+        config_inst, category_inst, variable_inst, density, shape_norm, yscale,
+    )
+    # disable autmatic setting of ylim
+    default_style_config["ax_cfg"]["ylim"] = None
+
+    style_config = law.util.merge_dicts(default_style_config, style_config, deep=True)
+    if shape_norm:
+        style_config["ax_cfg"]["ylabel"] = r"$\Delta N/N$"
+
+    # ratio plot not used here; set `skip_ratio` to True
+    kwargs["skip_ratio"] = True
+
+    return plot_all(plot_config, style_config, **kwargs)