Repository: scikit-hep/cabinetry

Issue:
Histogram colors in stacks - user interface creation
Hi!
It would be great to set colors of histograms in data / MC stack plots as soon as those are [created](https://github.com/scikit-hep/cabinetry/blob/master/src/cabinetry/visualize/plot_model.py#L24) .
Maybe a good way to do it is to add an argument to the function call, namely a dictionary like {sample name: color} ?

Comments:

Patch:
diff --git a/src/cabinetry/visualize/__init__.py b/src/cabinetry/visualize/__init__.py
index 3f53b86..91e40fb 100644
--- a/src/cabinetry/visualize/__init__.py
+++ b/src/cabinetry/visualize/__init__.py
@@ -56,6 +56,8 @@ def data_mc_from_histograms(
     figure_folder: Union[str, pathlib.Path] = "figures",
     log_scale: Optional[bool] = None,
     log_scale_x: bool = False,
+    channels: Optional[Union[str, List[str]]] = None,
+    colors: Optional[Dict[str, str]] = None,
     close_figure: bool = False,
     save_figure: bool = True,
 ) -> List[Dict[str, Any]]:
@@ -71,11 +73,18 @@ def data_mc_from_histograms(
             defaults to None (automatically determine whether to use linear/log scale)
         log_scale_x (bool, optional): whether to use logarithmic horizontal axis,
             defaults to False
+        channels (Optional[Union[str, List[str]]], optional): name of channel to show,
+            or list of names to include, defaults to None (uses all channels)
+        colors (Optional[Dict[str, str]], optional): map of sample names and colors to
+            use in plot, defaults to None (uses default colors)
         close_figure (bool, optional): whether to close each figure, defaults to False
             (enable when producing many figures to avoid memory issues, prevents
             automatic rendering in notebooks)
         save_figure (bool, optional): whether to save figures, defaults to True
 
+    Raises:
+        ValueError: if color specification is incomplete
+
     Returns:
         List[Dict[str, Any]]: list of dictionaries, where each dictionary contains a
             figure and the associated region name
@@ -83,7 +92,24 @@ def data_mc_from_histograms(
     log.info("visualizing histogram")
     histogram_folder = pathlib.Path(config["General"]["HistogramFolder"])
     figure_dict_list = []
+
+    # if custom colors are provided, ensure that they cover all samples
+    if colors is not None:
+        c_missing = {
+            sample["Name"] for sample in config["Samples"] if not sample.get("Data")
+        }.difference(colors.keys())
+        if c_missing:
+            raise ValueError(
+                f"colors need to be provided for all samples, missing for {c_missing}"
+            )
+
+    # create a list of channels to process
+    if channels is not None and isinstance(channels, str):
+        channels = [channels]
+
     for region in config["Regions"]:
+        if channels is not None and region["Name"] not in channels:
+            continue  # skip region
         histogram_dict_list = []
         model_stdevs = []
         # loop over samples in reverse order, such that samples that appear first in the
@@ -119,6 +145,7 @@ def data_mc_from_histograms(
             log_scale=log_scale,
             log_scale_x=log_scale_x,
             label=label,
+            colors=colors,
             close_figure=close_figure,
         )
         figure_dict_list.append({"figure": fig, "region": region["Name"]})
@@ -134,6 +161,7 @@ def data_mc(
     log_scale: Optional[bool] = None,
     log_scale_x: bool = False,
     channels: Optional[Union[str, List[str]]] = None,
+    colors: Optional[Dict[str, str]] = None,
     close_figure: bool = False,
     save_figure: bool = True,
 ) -> Optional[List[Dict[str, Any]]]:
@@ -159,11 +187,16 @@ def data_mc(
             defaults to False
         channels (Optional[Union[str, List[str]]], optional): name of channel to show,
             or list of names to include, defaults to None (uses all channels)
+        colors (Optional[Dict[str, str]], optional): map of sample names and colors to
+            use in plot, defaults to None (uses default colors)
         close_figure (bool, optional): whether to close each figure, defaults to False
             (enable when producing many figures to avoid memory issues, prevents
             automatic rendering in notebooks)
         save_figure (bool, optional): whether to save figures, defaults to True
 
+    Raises:
+        ValueError: if color specification is incomplete
+
     Returns:
         Optional[List[Dict[str, Any]]]: list of dictionaries, where each dictionary
             contains a figure and the associated region name, or None if no figure was
@@ -172,11 +205,19 @@ def data_mc(
     # strip off auxdata (if needed) and obtain data indexed by channel (and bin)
     data_yields = model_utils._data_per_channel(model_prediction.model, data)
 
-    # channels to include in table, with optional filtering applied
+    # if custom colors are provided, ensure that they cover all samples
+    if colors is not None:
+        c_missing = set(model_prediction.model.config.samples).difference(colors.keys())
+        if c_missing:
+            raise ValueError(
+                f"colors need to be provided for all samples, missing for {c_missing}"
+            )
+
+    # channels to include in plot, with optional filtering applied
     filtered_channels = model_utils._filter_channels(model_prediction.model, channels)
 
     if filtered_channels == []:
-        # nothing to include in tables, warning already raised via _filter_channels
+        # nothing to include in plots, warning already raised via _filter_channels
         return None
 
     # indices of included channels
@@ -241,6 +282,7 @@ def data_mc(
             log_scale=log_scale,
             log_scale_x=log_scale_x,
             label=label,
+            colors=colors,
             close_figure=close_figure,
         )
         figure_dict_list.append({"figure": fig, "region": channel_name})
diff --git a/src/cabinetry/visualize/plot_model.py b/src/cabinetry/visualize/plot_model.py
index 4ae8357..dda23ce 100644
--- a/src/cabinetry/visualize/plot_model.py
+++ b/src/cabinetry/visualize/plot_model.py
@@ -30,6 +30,7 @@ def data_mc(
     log_scale: Optional[bool] = None,
     log_scale_x: bool = False,
     label: str = "",
+    colors: Optional[Dict[str, str]] = None,
     close_figure: bool = False,
 ) -> mpl.figure.Figure:
     """Draws a data/MC histogram with uncertainty bands and ratio panel.
@@ -48,6 +49,8 @@ def data_mc(
         log_scale_x (bool, optional): whether to use logarithmic horizontal axis,
             defaults to False
         label (str, optional): label written on the figure, defaults to ""
+        colors (Optional[Dict[str, str]], optional): map of sample names and colors to
+            use in plot, defaults to None (uses default colors)
         close_figure (bool, optional): whether to close each figure immediately after
             saving it, defaults to False (enable when producing many figures to avoid
             memory issues, prevents rendering in notebooks)
@@ -102,9 +105,13 @@ def data_mc(
         else bin_centers
     )
     mc_containers = []
-    for mc_sample_yield in mc_histograms_yields:
+    for mc_sample_yield, sample_label in zip(mc_histograms_yields, mc_labels):
         mc_container = ax1.bar(
-            bin_centers, mc_sample_yield, width=bin_width, bottom=total_yield
+            bin_centers,
+            mc_sample_yield,
+            width=bin_width,
+            bottom=total_yield,
+            color=colors[sample_label] if colors else None,
         )
         mc_containers.append(mc_container)
 


Test patch:
diff --git a/tests/visualize/test_visualize.py b/tests/visualize/test_visualize.py
index 7ece48b..10c2394 100644
--- a/tests/visualize/test_visualize.py
+++ b/tests/visualize/test_visualize.py
@@ -105,17 +105,21 @@ def test_data_mc_from_histograms(mock_load, mock_draw, mock_stdev):
                 "log_scale": None,
                 "log_scale_x": False,
                 "label": "reg_1\npre-fit",
+                "colors": None,
                 "close_figure": False,
             },
         )
     ]
 
-    # custom log scale settings, close figure, do not save figure
+    # custom log scale settings, close figure, do not save figure, channel specified,
+    # custom colors
     _ = visualize.data_mc_from_histograms(
         config,
         figure_folder=figure_folder,
         log_scale=True,
         log_scale_x=True,
+        channels=["reg_1"],
+        colors={"sample_1": "red"},
         close_figure=True,
         save_figure=False,
     )
@@ -124,9 +128,20 @@ def test_data_mc_from_histograms(mock_load, mock_draw, mock_stdev):
         "log_scale": True,
         "log_scale_x": True,
         "label": "reg_1\npre-fit",
+        "colors": {"sample_1": "red"},
         "close_figure": True,
     }
 
+    # no matching channels
+    assert visualize.data_mc_from_histograms(config, channels="abc") == []
+
+    # incomplete color specification
+    with pytest.raises(
+        ValueError,
+        match="colors need to be provided for all samples, missing for {'sample_1'}",
+    ):
+        _ = visualize.data_mc_from_histograms(config, colors={})
+
 
 @mock.patch(
     "cabinetry.visualize.plot_model.data_mc", return_value=matplotlib.figure.Figure()
@@ -191,12 +206,13 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
         "log_scale": None,
         "log_scale_x": False,
         "label": "Signal Region\npre-fit",
+        "colors": None,
         "close_figure": False,
     }
 
     # post-fit plot (different label in model prediction), custom scale, close figure,
-    # do not save figure, histogram input mode: no binning or variable specified (via
-    # side effect)
+    # do not save figure, custom colors, histogram input mode: no binning or variable
+    # specified (via side effect)
     model_pred = model_utils.ModelPrediction(
         model, [[[11.0]]], [[[0.2], [0.2]]], [[0.2, 0.2]], "post-fit"
     )
@@ -206,6 +222,7 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
         config=config,
         figure_folder=figure_folder,
         log_scale=False,
+        colors={"Signal": "red"},
         close_figure=True,
         save_figure=False,
     )
@@ -224,6 +241,7 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
         "log_scale": False,
         "log_scale_x": False,
         "label": "Signal Region\npost-fit",
+        "colors": {"Signal": "red"},
         "close_figure": True,
     }
 
@@ -239,6 +257,13 @@ def test_data_mc(mock_data, mock_filter, mock_dict, mock_bins, mock_draw, exampl
     assert mock_filter.call_args == ((model, "abc"), {})
     assert mock_draw.call_count == 3  # no new call
 
+    # incomplete color specification
+    with pytest.raises(
+        ValueError,
+        match="colors need to be provided for all samples, missing for {'Signal'}",
+    ):
+        _ = visualize.data_mc(model_pred, data, colors={})
+
 
 @mock.patch(
     "cabinetry.histo.Histogram.from_path",
diff --git a/tests/visualize/test_visualize_plot_model.py b/tests/visualize/test_visualize_plot_model.py
index a6e49f5..2de5908 100644
--- a/tests/visualize/test_visualize_plot_model.py
+++ b/tests/visualize/test_visualize_plot_model.py
@@ -1,6 +1,7 @@
 import copy
 from unittest import mock
 
+import matplotlib as mpl
 import matplotlib.pyplot as plt
 from matplotlib.testing.compare import compare_images
 import numpy as np
@@ -103,19 +104,30 @@ def test_data_mc(tmp_path, caplog):
 
     # do not save figure, but close it
     # one bin with zero model prediction
+    # custom histogram colors
     histo_dict_list[0]["yields"] = np.asarray([0, 14])
     histo_dict_list[1]["yields"] = np.asarray([0, 5])
     caplog.clear()
     with mock.patch("cabinetry.visualize.utils._save_and_close") as mock_close_safe:
         with pytest.warns() as warn_record:
-            fig = fig = plot_model.data_mc(
+            fig = plot_model.data_mc(
                 histo_dict_list,
                 total_model_unc_log,
                 bin_edges_log,
                 label="",
+                colors={"Background": "blue", "Signal": "red"},
                 close_figure=True,
             )
             assert mock_close_safe.call_args_list == [((fig, None, True), {})]
+            # custom colors propagated to histogram
+            assert (
+                fig.axes[0].containers[0].patches[0].get_facecolor()
+                == mpl.colors.to_rgba_array("blue")
+            ).all()
+            assert (
+                fig.axes[0].containers[1].patches[0].get_facecolor()
+                == mpl.colors.to_rgba_array("red")
+            ).all()
 
     assert "predicted yield is zero in 1 bin(s), excluded from ratio plot" in [
         rec.message for rec in caplog.records