Refactor eft and resonant limit tasks.

Author: riga

dhi/tasks/combine.py

@@ -779,48 +779,6 @@ def store_parts(self):
         return parts
 
 
-class MultiDatacardTransposedTask(MultiDatacardTask):
-
-    exclude_params_index = {"datacard_names", "datacard_order"}
-
-    datacard_names = None
-    datacard_order = None
-    group_duplicate_cards = False
-
-    @classmethod
-    def extract_info_from_datacard_path(cls, datacard):
-        return os.path.splitext(os.path.basename(datacard).rsplit("_", 1)[-1])[0]
-
-    def __init__(self, *args, **kwargs):
-        super(MultiDatacardTransposedTask, self).__init__(*args, **kwargs)
-
-        # create a map of datacard info strings to lists of cards that contain it
-        self.multi_datacards_transposed = OrderedDict()
-        seen = set()
-        for datacards in self.multi_datacards:
-            contains_duplicate = any(datacard in seen for datacard in datacards)
-            groups = OrderedDict()
-            for datacard in datacards:
-                # extract the info string from the basename
-                info = self.extract_info_from_datacard_path(datacard)
-
-                if not self.group_duplicate_cards:
-                    # when not grouping, just add the card
-                    self.multi_datacards_transposed.setdefault(info, [[]])[0].append(datacard)
-                elif not contains_duplicate:
-                    # when the sequence contains only unseen cards, just add the card
-                    self.multi_datacards_transposed.setdefault(info, []).append([datacard])
-                else:
-                    # add it to groups for the current sequence
-                    groups.setdefault(info, []).append(datacard)
-                seen.add(datacard)
-
-            # add groups if any
-            if groups:
-                for info, group in groups.items():
-                    self.multi_datacards_transposed.setdefault(info, []).append(group)
-
-
 class ParameterValuesTask(AnalysisTask):
 
     parameter_values = ModelParameters(

dhi/tasks/eft.py

@@ -4,7 +4,8 @@
 Tasks related to EFT benchmarks and scans.
 """
 
-from collections import OrderedDict
+import re
+from collections import OrderedDict, defaultdict
 
 import law
 import luigi
@@ -13,7 +14,6 @@
 from dhi.tasks.remote import HTCondorWorkflow
 from dhi.tasks.combine import (
     MultiDatacardTask,
-    MultiDatacardTransposedTask,
     POITask,
     POIPlotTask,
     CombineCommandTask,
@@ -24,28 +24,44 @@
 from dhi.config import br_hh
 
 
-class EFTBase(POITask, MultiDatacardTransposedTask):
+class EFTBase(POITask):
 
+    datacard_pattern = luigi.Parameter(
+        default=r"^.*_([^_]+)\.txt$",
+        description="a regular expression with a single match group that is supposed to point to "
+        "the benchmark name in the datacard path; default: ^.*_([^_]+)\\.txt$",
+    )
     hh_model = law.NO_STR
     allow_empty_hh_model = True
 
     poi = "r_gghh"
 
     @classmethod
-    def modify_param_values(cls, params):
-        params = POITask.modify_param_values.__func__.__get__(cls)(params)
-        params = MultiDatacardTransposedTask.modify_param_values.__func__.__get__(cls)(params)
-        return params
+    def _group_datacards(cls, datacards, cre):
+        groups = defaultdict(list)
+        for datacard in datacards:
+            m = cre.match(datacard)
+            if not m:
+                raise Exception(
+                    f"no benchmark value could be extracted from datacard '{datacard}' "
+                    f"with pattern '{cre.pattern}'",
+                )
+            groups[m.group(1)].append(datacard)
+
+        return OrderedDict([
+            (bm, sorted(groups[bm]))
+            for bm in sort_eft_benchmark_names(groups.keys())
+        ])
 
     def __init__(self, *args, **kwargs):
         super(EFTBase, self).__init__(*args, **kwargs)
 
-        # sort EFT datacards according to benchmark names
-        names = sort_eft_benchmark_names(self.multi_datacards_transposed.keys())
-        self.benchmark_datacards = OrderedDict(
-            (name, self.multi_datacards_transposed[name])
-            for name in names
-        )
+        # group datacards into a dictionary benchmark -> [cards]
+        self.benchmark_datacards = self.group_datacards()
+
+    def group_datacards(self):
+        cre = re.compile(self.datacard_pattern)
+        return self._group_datacards(self.datacards, cre)
 
     @property
     def other_pois(self):
@@ -62,11 +78,10 @@ class EFTBenchmarkLimits(EFTBase, CombineCommandTask, law.LocalWorkflow, HTCondo
     run_command_in_tmp = True
 
     def create_branch_map(self):
-        branch_map = []
-        for name, cards in self.benchmark_datacards.items():
-            for _cards in cards:
-                branch_map.append({"benchmark": name, "cards": _cards})
-        return branch_map
+        return [
+            {"benchmark": benchmark, "cards": cards}
+            for benchmark, cards in self.benchmark_datacards.items()
+        ]
 
     def workflow_requires(self):
         reqs = super(EFTBenchmarkLimits, self).workflow_requires()
@@ -273,45 +288,21 @@ def run(self):
         )
 
 
-class PlotMultipleEFTBenchmarkLimits(PlotEFTBenchmarkLimits):
-
-    datacard_names = MultiDatacardTask.datacard_names
-    datacard_order = MultiDatacardTask.datacard_order
-    group_duplicate_cards = True
+class PlotMultipleEFTBenchmarkLimits(PlotEFTBenchmarkLimits, MultiDatacardTask):
 
     default_plot_function = "dhi.plots.eft.plot_multi_benchmark_limits"
 
-    def __init__(self, *args, **kwargs):
-        super(PlotMultipleEFTBenchmarkLimits, self).__init__(*args, **kwargs)
-
-        # check that each mass point has the same amount of cards
-        n_entries = {len(cards) for cards in self.benchmark_datacards.values()}
-        if len(n_entries) != 1:
-            raise Exception("founds different amount of entries in input datacards: {}".format(
-                ",".join(map(str, n_entries)),
-            ))
-        self.n_entries = list(n_entries)[0]
-
-        # the lengths of names and order indices must match multi_datacards when given
-        if self.datacard_names and len(self.datacard_names) != self.n_entries:
-            raise Exception("found {} entries in datacard_names whereas {} are expected".format(
-                len(self.datacard_names), self.n_entries,
-            ))
-        if self.datacard_order and len(self.datacard_order) != self.n_entries:
-            raise Exception("found {} entries in datacard_order whereas {} are expected".format(
-                len(self.datacard_order), self.n_entries,
-            ))
+    def group_datacards(self):
+        cre = re.compile(self.datacard_pattern)
+        return [
+            self._group_datacards(datacards, cre)
+            for datacards in self.multi_datacards
+        ]
 
     def requires(self):
         return [
-            MergeEFTBenchmarkLimits.req(
-                self,
-                multi_datacards=tuple(
-                    tuple(cards[i])
-                    for cards in self.benchmark_datacards.values()
-                ),
-            )
-            for i in range(self.n_entries)
+            MergeEFTBenchmarkLimits.req(self, datacards=tuple(sum(groups.values(), [])))
+            for groups in self.benchmark_datacards
         ]
 
     def output(self):

dhi/tasks/resonant.py

@@ -4,7 +4,8 @@
 Tasks related to upper limits on resonant scenarios.
 """
 
-from collections import OrderedDict
+import re
+from collections import OrderedDict, defaultdict
 
 import law
 import luigi
@@ -13,7 +14,6 @@
 from dhi.tasks.remote import HTCondorWorkflow
 from dhi.tasks.combine import (
     MultiDatacardTask,
-    MultiDatacardTransposedTask,
     POITask,
     POIPlotTask,
     CombineCommandTask,
@@ -23,35 +23,45 @@
 from dhi.config import br_hh
 
 
-class ResonantBase(POITask, MultiDatacardTransposedTask):
+class ResonantBase(POITask):
 
+    datacard_pattern = luigi.Parameter(
+        default=r"^.*_(\d+)\.txt$",
+        description="a regular expression with a single match group that is supposed to point to "
+        "the resonance mass value in the datacard path; default: ^.*_(\\d+)\\.txt$",
+    )
     hh_model = law.NO_STR
     allow_empty_hh_model = True
 
     poi = "r_xhh"
     scan_parameter = "mhh"
 
     @classmethod
-    def modify_param_values(cls, params):
-        params = POITask.modify_param_values.__func__.__get__(cls)(params)
-        params = MultiDatacardTransposedTask.modify_param_values.__func__.__get__(cls)(params)
-        return params
+    def _group_datacards(cls, datacards, cre):
+        groups = defaultdict(list)
+        for datacard in datacards:
+            m = cre.match(datacard)
+            if not m:
+                raise Exception(
+                    f"no resonance mass could be extracted from datacard '{datacard}' "
+                    f"with pattern '{cre.pattern}'",
+                )
+            groups[int(m.group(1))].append(datacard)
+
+        return OrderedDict([
+            (mass, sorted(groups[mass]))
+            for mass in sorted(groups)
+        ])
 
     def __init__(self, *args, **kwargs):
         super(ResonantBase, self).__init__(*args, **kwargs)
 
-        # convert keys in multi_datacards_transposed to integers and store them as resonant cards
-        pairs = []
-        for info, datacards in self.multi_datacards_transposed.items():
-            try:
-                mass = int(info)
-            except:
-                raise Exception(
-                    "datacards contain a mass point '{}' which cannot be interpreted as an "
-                    "integer".format(info),
-                )
-            pairs.append((mass, datacards))
-        self.resonant_datacards = OrderedDict(sorted(pairs, key=lambda pair: pair[0]))
+        # group datacards into a dictionary mass -> [cards]
+        self.resonant_datacards = self.group_datacards()
+
+    def group_datacards(self):
+        cre = re.compile(self.datacard_pattern)
+        return self._group_datacards(self.datacards, cre)
 
     @property
     def other_pois(self):
@@ -68,11 +78,10 @@ class ResonantLimits(ResonantBase, CombineCommandTask, law.LocalWorkflow, HTCond
     run_command_in_tmp = True
 
     def create_branch_map(self):
-        branch_map = []
-        for mass, cards in self.resonant_datacards.items():
-            for _cards in cards:
-                branch_map.append({"mass": mass, "cards": _cards})
-        return branch_map
+        return [
+            {"mass": mass, "cards": cards}
+            for mass, cards in self.resonant_datacards.items()
+        ]
 
     def workflow_requires(self):
         reqs = super(ResonantLimits, self).workflow_requires()
@@ -288,45 +297,21 @@ def run(self):
         )
 
 
-class PlotMultipleResonantLimits(PlotResonantLimits):
-
-    datacard_names = MultiDatacardTask.datacard_names
-    datacard_order = MultiDatacardTask.datacard_order
-    group_duplicate_cards = True
+class PlotMultipleResonantLimits(PlotResonantLimits, MultiDatacardTask):
 
     default_plot_function = "dhi.plots.limits.plot_limit_scans"
 
-    def __init__(self, *args, **kwargs):
-        super(PlotMultipleResonantLimits, self).__init__(*args, **kwargs)
-
-        # check that each mass point has the same amount of cards
-        n_entries = {len(cards) for cards in self.resonant_datacards.values()}
-        if len(n_entries) != 1:
-            raise Exception("founds different amount of entries in input datacards: {}".format(
-                ",".join(map(str, n_entries)),
-            ))
-        self.n_entries = list(n_entries)[0]
-
-        # the lengths of names and order indices must match multi_datacards when given
-        if self.datacard_names and len(self.datacard_names) != self.n_entries:
-            raise Exception("found {} entries in datacard_names whereas {} are expected".format(
-                len(self.datacard_names), self.n_entries,
-            ))
-        if self.datacard_order and len(self.datacard_order) != self.n_entries:
-            raise Exception("found {} entries in datacard_order whereas {} are expected".format(
-                len(self.datacard_order), self.n_entries,
-            ))
+    def group_datacards(self):
+        cre = re.compile(self.datacard_pattern)
+        return [
+            self._group_datacards(datacards, cre)
+            for datacards in self.multi_datacards
+        ]
 
     def requires(self):
         return [
-            MergeResonantLimits.req(
-                self,
-                multi_datacards=tuple(
-                    tuple(cards[i])
-                    for cards in self.resonant_datacards.values()
-                ),
-            )
-            for i in range(self.n_entries)
+            MergeResonantLimits.req(self, datacards=tuple(sum(groups.values(), [])))
+            for groups in self.resonant_datacards
         ]
 
     def output(self):

dhi/tasks/snapshot.py

@@ -10,6 +10,7 @@
 import luigi
 
 from dhi.tasks.remote import HTCondorWorkflow
+from dhi.tasks.base import AnalysisTask
 from dhi.tasks.combine import (
     CombineCommandTask,
     POITask,
@@ -82,7 +83,7 @@ def build_command(self, fallback_level):
         return cmd
 
 
-class SnapshotUser(object):
+class SnapshotUser(AnalysisTask):
 
     use_snapshot = luigi.BoolParameter(
         default=False,

dhi/tasks/test.py

@@ -413,7 +413,7 @@ def requires(self):
         if self.check_enabled("eft_benchmark_limits"):
             reqs["eft_benchmark_limits"] = PlotEFTBenchmarkLimits.req(
                 self,
-                multi_datacards=(eft_bm_cards,),
+                datacards=eft_bm_cards,
                 unblinded=True,
                 xsec="fb",
                 y_log=True,
@@ -459,7 +459,7 @@ def requires(self):
         if self.check_enabled("resonant_limits"):
             reqs["resonant_limits"] = PlotResonantLimits.req(
                 self,
-                multi_datacards=(res_cards,),
+                datacards=res_cards,
                 unblinded=False,
                 xsec="fb",
                 y_log=True,

docs/content/snippets/createworkspace_param_tab.md

@@ -2,6 +2,6 @@ The `CreateWorkspace` task takes the combined datacard and the PhysicsModel as i
 
 <div class="dhi_parameter_table">
 
---8<-- "content/snippets/parameters.md@-2,20,19,34"
+--8<-- "content/snippets/parameters.md@-2,20,19,34,98,99"
 
 </div>

docs/content/snippets/eftbenchmarklimits_param_tab.md

@@ -2,6 +2,6 @@ The `EFTBenchmarkLimits` task computes the limits of each benchmark datacard.
 
 <div class="dhi_parameter_table">
 
---8<-- "content/snippets/parameters.md@-2,21,48,34,17,18,47"
+--8<-- "content/snippets/parameters.md@-2,20,101,48,34,17,18,47"
 
 </div>