Plot the shapes of a datacard bin from where the maker reads them (input_categories)

.github/copilot-instructions.md

@@ -102,6 +102,6 @@ Verified 2026-08-27; re-check before relying on any of it.
 | Layout | `dc_make/` (datacard construction: `maker.py`, `create_datacards.py`, `process.py`, `model.py`, `uncertainty.py`, `binner.py`), `bin_opt/` (binning optimisation and limits), `common/` (`tools.py`, `param_parse.py`), `law/tasks.py`, `fit_val/`, `config/`, `legacy/`, `test/` (unit tests and their scratch analysis) |
 | Consumers | HH_bbtautau and HH_bbWW use this as a submodule; H_mumu does not |
 | Cards | analysis-side, e.g. `HH_bbWW/config/Datacards/CI_card.yaml` |
-| Tests | `test/test_law_chain.py` (the law tasks on the real FLAF and law classes over a scratch analysis, `test/statinference_fixture.py`: per card the outputs, inputs and payload commands; no ROOT, Combine or dhi), run by `unit-tests` against a pinned FLAF commit |
+| Tests | `test/test_law_chain.py` (the law tasks on the real FLAF and law classes over a scratch analysis, `test/statinference_fixture.py`: per card the outputs, inputs and payload commands) and `test/test_shape_plots.py` (which histograms a card's shape plots show, through `input_categories`); no ROOT, Combine or dhi; run by `unit-tests` against a pinned FLAF commit |
 | Workflows | `formatting-check`, `unit-tests`, `repo-sanity-checks` (FLAF's sanity checks and FLAF's `run_tools/check_era_rename.py` for leftovers of the `period` → `era` rename), `trigger-flaf-integration`. Formatting is checked automatically — do not comment on it. A change to an assertion of `test/test_law_chain.py` (not to its helpers) changes what the chain is guaranteed to read, write or run — comment on it |
 | Integration test | Triggered by `@cms-flaf-bot please test`; the pipeline configuration lives in `cms-flaf/FLAF_ci` |

.github/workflows/unit-tests.yaml

@@ -46,3 +46,4 @@ jobs:
         run: >-
           python -m pytest -q
           StatInference/test/test_law_chain.py
+          StatInference/test/test_shape_plots.py

README.md

@@ -101,12 +101,13 @@ rm -r output/binning_v2/tauTau_2018/workers
 
 `test/test_law_chain.py` checks the law tasks on the real FLAF and law classes over a scratch
 analysis (`test/statinference_fixture.py`): for each card, what it writes, which merged
-histograms it reads and the commands it runs. It needs no ROOT, Combine or dhi. The
-`unit-tests` workflow runs it against a pinned FLAF commit; to run it locally, put a FLAF
-checkout at that commit (`ref:` in `.github/workflows/unit-tests.yaml`) next to this one and
-run from the directory holding both:
+histograms it reads and the commands it runs. `test/test_shape_plots.py` checks which
+histograms the shape plots of a card show, for categories read through `input_categories`.
+They need no ROOT, Combine or dhi. The `unit-tests` workflow runs them against a pinned FLAF
+commit; to run them locally, put a FLAF checkout at that commit (`ref:` in
+`.github/workflows/unit-tests.yaml`) next to this one and run from the directory holding both:
 
 ```bash
 pip install pyyaml numpy luigi==3.8.1 law==0.1.21 pytest
-PYTHONPATH=$PWD python -m pytest -q StatInference/test/test_law_chain.py
+PYTHONPATH=$PWD python -m pytest -q StatInference/test/test_law_chain.py StatInference/test/test_shape_plots.py
 ```

law/CreateDatacardsTask.py

@@ -8,7 +8,7 @@
 from FLAF.RunKit.run_tools import ps_call
 from FLAF.run_tools.law_customizations import HTCondorWorkflow, copy_param
 
-from StatInference.common.tools import CategoryNaming, importROOT
+from StatInference.common.tools import CategoryNaming, InputCategories, importROOT
 
 from .PreprocessShapesTask import PreprocessShapesTask
 from .StatInferenceTask import StatInferenceTask
@@ -223,23 +223,39 @@ def _stitch_grid(panels, out_path):
         return True
 
     @staticmethod
-    def _present_keys(shape_file, cfg):
-        """The (channel, region, category) triples the summed shape file actually holds."""
+    def _present_keys(shape_file, cfg, points):
+        """({(channel, region, category): input directory} for the datacard bins the summed
+        shape file holds, [every path looked up], or None when the file cannot be read).
+
+        Each category is looked up where the maker reads it for these model points: its
+        `input_categories` directory, else its own name. A category that the points read
+        from different directories is not plotted, since one plot cannot show them.
+        """
         ROOT = importROOT()
-        present = set()
+        input_categories = InputCategories.fromConfig(cfg)
+        present, looked_up = {}, []
         f = ROOT.TFile.Open(shape_file)
         if not f or f.IsZombie():
-            return present
+            return present, None
         try:
             for category in cfg["categories"]:
                 region, _, cat = category.rpartition("/")
+                directories = {input_categories.path(category, p) for p in points}
+                if len(directories) != 1:
+                    print(
+                        f"WARNING: {category} is read from {', '.join(sorted(directories))} "
+                        "for the model points of one variable; it is not plotted"
+                    )
+                    continue
+                (directory,) = directories
                 for channel in cfg["channels"]:
-                    path = "/".join(x for x in (channel, region, cat) if x)
+                    path = "/".join(x for x in (channel, directory) if x)
+                    looked_up.append(path)
                     if f.Get(path):
-                        present.add((channel, region, cat))
+                        present[(channel, region, cat)] = directory
         finally:
             f.Close()
-        return present
+        return present, looked_up
 
     def _stitch_grids(self, cfg, plots_dir, variable, panel_of_key):
         """One grid per base category: its slices across, the channels down."""
@@ -289,6 +305,7 @@ def plot_rebinned_shapes(self, base_dir, config, output_dir):
         # The datacard bin is the sum over the sub-eras (getCombinedShape does the same at
         # card-build time), so the sub-era files are hadd'ed into the one file the plotter
         # reads. Plotting them separately would show something the fit never sees.
+        variable_points = self.get_variable_points()
         for variable in self.get_required_variables():
             inputs = [
                 p
@@ -315,23 +332,37 @@ def plot_rebinned_shapes(self, base_dir, config, output_dir):
             # from the shapes. HistPlotter exits non-zero on the first key it cannot find,
             # so passing the full list costs every panel and grid that would have come
             # after the gap, not just the missing one.
-            present = self._present_keys(summed, cfg)
-            keys, outputs = [], []
+            present, looked_up = self._present_keys(
+                summed, cfg, variable_points[variable]
+            )
+            keys, outputs, panel_of_key = [], [], {}
             for category in cfg["categories"]:
                 # "SR/res2b_dnn0" -> region "SR", category "res2b_dnn0": HistPlotter
                 # navigates channel -> region -> category.
                 region, _, cat = category.rpartition("/")
                 for channel in cfg["channels"]:
-                    if (channel, region, cat) not in present:
+                    directory = present.get((channel, region, cat))
+                    if directory is None:
                         continue
-                    keys.append(f"{channel}:{cat}:{region}")
-                    outputs.append(
-                        os.path.join(
-                            plots_dir, f"{variable}_{channel}_{region}_{cat}.pdf"
-                        )
+                    # The key names the directory the shapes are read from; the file and
+                    # the grid keep the datacard bin's name.
+                    in_region, _, in_cat = directory.rpartition("/")
+                    keys.append(f"{channel}:{in_cat}:{in_region}")
+                    output = os.path.join(
+                        plots_dir, f"{variable}_{channel}_{region}_{cat}.pdf"
                     )
+                    outputs.append(output)
+                    panel_of_key[f"{channel}:{cat}:{region}"] = output
             if not keys:
-                print(f"WARNING: no shapes to plot for {variable}")
+                if looked_up is None:
+                    reason = f"{summed} cannot be read"
+                elif looked_up:
+                    reason = f"none of {', '.join(looked_up)} is in the summed shapes"
+                else:
+                    reason = "no category is read from one directory"
+                print(f"WARNING: no shapes to plot for {variable}: {reason}")
+                if os.path.exists(summed):
+                    os.remove(summed)
                 continue
             cmd = [
                 "python3",
@@ -372,9 +403,11 @@ def plot_rebinned_shapes(self, base_dir, config, output_dir):
                 "--wantLogScale",
                 "y",
             ]
+            if self.user_custom:
+                cmd += ["--user-custom", self.user_custom]
             try:
                 ps_call(cmd, verbose=1)
-                self._stitch_grids(cfg, plots_dir, variable, dict(zip(keys, outputs)))
+                self._stitch_grids(cfg, plots_dir, variable, panel_of_key)
             except Exception as e:
                 # The datacards are the product that matters; a plotting failure should be
                 # loud but must not leave the task looking broken with valid cards on disk.

law/StatInferenceTask.py

@@ -159,27 +159,34 @@ def get_masses(self):
             )
         return sorted(masses)
 
-    def get_required_variables(self):
-        """The Hists_merged variables this chain reads: one per mass point (the 2D
-        shape the datacards are built from), derived from the datacard config's
-        input_file_pattern.
-
-        This is the *complete* input list -- there is no intersection with global.yaml's
-        histTuple_flavor variable list anywhere downstream, so a mass point present here
-        is read whether or not the active flavour happens to mention it.
-        """
+    def get_variable_points(self):
+        """{variable: [model point, ...]}: the mass points whose input is that
+        Hists_merged variable, each as {parameter: value}, from the datacard config's
+        input_file_pattern."""
         data = self.get_config_data()
         model = data["model"]
         pattern = model["input_file_pattern"]
         param_name = model["parameters"][0]
 
-        variables = set()
+        points = {}
         for mass in self.get_masses():
             rel = Template(pattern).safe_substitute({"ERA": "ERA", param_name: mass})
             # "<era>/<variable>/<variable>.root": the variable is the directory holding
             # the file, not the file itself.
-            variables.add(os.path.basename(os.path.dirname(rel)))
-        return sorted(variables)
+            variable = os.path.basename(os.path.dirname(rel))
+            points.setdefault(variable, []).append({param_name: mass})
+        return points
+
+    def get_required_variables(self):
+        """The Hists_merged variables this chain reads: one per mass point (the 2D
+        shape the datacards are built from), derived from the datacard config's
+        input_file_pattern.
+
+        This is the *complete* input list -- there is no intersection with global.yaml's
+        histTuple_flavor variable list anywhere downstream, so a mass point present here
+        is read whether or not the active flavour happens to mention it.
+        """
+        return sorted(self.get_variable_points())
 
     def merged_hist_reqs(self, eras):
         """{(era, variable): MergedHists} -- every merged input file for those eras.