Patch setupRabbit.py operations

scripts/rabbit/setupRabbit.py

@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 import argparse
 import math
+import sys
 
 import hist
 import numpy as np
@@ -26,7 +27,82 @@
 from wums import logging, output_tools
 
 
-def make_subparsers(parser):
+def _parse_axis_range_specs(specs):
+    parsed_specs = []
+    for axis, low, high in specs:
+        parsed_specs.append(
+            (
+                axis,
+                parsing.str_to_complex_or_int(low),
+                parsing.str_to_complex_or_int(high),
+            )
+        )
+    return parsed_specs
+
+
+def _build_fitvar_axlim(axlim_specs, fitvar):
+    if not axlim_specs:
+        return []
+
+    parsed_specs = _parse_axis_range_specs(axlim_specs)
+    axlim = [None] * (2 * len(fitvar))
+    seen_axes = set()
+    for axis, low, high in parsed_specs:
+        if axis not in fitvar:
+            raise ValueError(
+                f"--axlim only accepts fit variables. Axis '{axis}' is not one of {fitvar}"
+            )
+        if axis in seen_axes:
+            raise ValueError(f"Duplicate axis '{axis}' passed to --axlim")
+        seen_axes.add(axis)
+        idx = fitvar.index(axis)
+        axlim[2 * idx] = low
+        axlim[2 * idx + 1] = high
+
+    return axlim
+
+
+def _build_preselection_specs(selection_specs, fitvar):
+    if not selection_specs:
+        return []
+
+    parsed_specs = _parse_axis_range_specs(selection_specs)
+    seen_axes = set()
+    for axis, _, _ in parsed_specs:
+        if axis in fitvar:
+            raise ValueError(
+                f"--preselect only accepts non-fit axes. Axis '{axis}' is one of the fit variables {fitvar}"
+            )
+        if axis in seen_axes:
+            raise ValueError(f"Duplicate axis '{axis}' passed to --preselect")
+        seen_axes.add(axis)
+
+    return parsed_specs
+
+
+def _normalize_negative_imaginary_bounds(argv):
+    normalized_argv = []
+    i = 0
+    while i < len(argv):
+        token = argv[i]
+        normalized_argv.append(token)
+
+        if token in {"--axlim", "--preselect"} and i + 3 < len(argv):
+            normalized_argv.append(argv[i + 1])
+            for value in (argv[i + 2], argv[i + 3]):
+                if value.startswith("-") and value.endswith("j"):
+                    normalized_argv.append(f" {value}")
+                else:
+                    normalized_argv.append(value)
+            i += 4
+            continue
+
+        i += 1
+
+    return normalized_argv
+
+
+def make_subparsers(parser, argv=None):
 
     parser.add_argument(
         "--analysisMode",
@@ -36,7 +112,7 @@ def make_subparsers(parser):
         help="Select analysis mode to run. Default is the traditional analysis",
     )
 
-    tmpKnownArgs, _ = parser.parse_known_args()
+    tmpKnownArgs, _ = parser.parse_known_args(argv)
     subparserName = tmpKnownArgs.analysisMode
     if subparserName is None:
         return parser
@@ -133,7 +209,7 @@ def make_subparsers(parser):
     return parser
 
 
-def make_parser(parser=None):
+def make_parser(parser=None, argv=None):
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "-o",
@@ -265,12 +341,14 @@ def make_parser(parser=None):
     )
     parser.add_argument(
         "--axlim",
-        type=parsing.str_to_complex_or_int,
         default=[],
-        nargs="*",
+        nargs=3,
+        action="append",
+        metavar=("AXIS", "LOW", "HIGH"),
         help="""
-        Restrict axis to this range or these bins (assumes pairs of values by axis, with trailing axes optional).
-        Arguments must be pure real or pure imaginary numbers to select bin indices or values, respectively.
+        Restrict a fit axis to this range or these bins.
+        Repeat as '--axlim AXIS LOW HIGH'. LOW and HIGH must be pure real integers for bin indices
+        or pure imaginary numbers for axis values.
         """,
     )
     parser.add_argument(
@@ -871,16 +949,15 @@ def make_parser(parser=None):
         help="probability density for systematic variations",
     )
     parser.add_argument(
-        "--select",
-        nargs="+",
-        dest="selection",
-        type=str,
-        default=None,
-        help="Apply a selection to the histograms, if the axis exists."
-        "This option can be applied to any of the axis, not necessarily one of the fitaxes, unlike --axlim."
-        "Use complex numbers for axis value, integers for bin number."
-        "e.g. --select 'ptll 0 10"
-        "e.g. --select 'ptll 0j 10j",
+        "--presel",
+        nargs=3,
+        action="append",
+        default=[],
+        metavar=("AXIS", "LOW", "HIGH"),
+        help="Apply a strict preselection on a non-fit axis before downstream projections."
+        " Repeat as '--presel AXIS LOW HIGH'."
+        " LOW and HIGH must be pure real integers for bin indices or pure imaginary numbers for axis values."
+        " The command fails if a requested axis is missing from any loaded histogram.",
     )
     parser.add_argument(
         "--noTheoryCorrsViaHelicities",
@@ -892,7 +969,7 @@ def make_parser(parser=None):
         action="store_true",
         help="Use the Breit-Wigner mass wights for mW.",
     )
-    parser = make_subparsers(parser)
+    parser = make_subparsers(parser, argv=argv)
 
     return parser
 
@@ -964,18 +1041,19 @@ def setup(
     datagroups.fit_axes = fitvar
     datagroups.channel = channel
 
-    if args.selection:
-        for sel in args.selection:
-            sel_ax, sel_lb, sel_ub = sel.split()
-            sel_lb = parsing.str_to_complex_or_int(sel_lb)
-            sel_ub = parsing.str_to_complex_or_int(sel_ub)
-            datagroups.setGlobalAction(
-                lambda h: (
-                    h[{sel_ax: slice(sel_lb, sel_ub, hist.sum)}]
-                    if sel_ax in h.axes.name
-                    else h
-                ),
-            )
+    preselection_specs = _build_preselection_specs(args.presel, fitvar)
+    if preselection_specs:
+
+        def apply_preselection(h, specs=tuple(preselection_specs)):
+            for axis, low, high in specs:
+                if axis not in h.axes.name:
+                    raise ValueError(
+                        f"--preselect requested axis '{axis}', but histogram axes are {h.axes.name}"
+                    )
+                h = h[{axis: slice(low, hh.get_hist_slice_upper(h, axis, high))}]
+            return h
+
+        datagroups.setGlobalAction(apply_preselection)
 
     if args.angularCoeffs:
         datagroups.setGlobalAction(
@@ -984,10 +1062,11 @@ def setup(
             )
         )
 
-    if args.axlim or args.rebin or args.absval:
+    fitvar_axlim = _build_fitvar_axlim(args.axlim, fitvar)
+    if fitvar_axlim or args.rebin or args.absval:
         datagroups.set_rebin_action(
             fitvar,
-            args.axlim,
+            fitvar_axlim,
             args.rebin,
             args.absval,
             args.rebinBeforeSelection,
@@ -1064,7 +1143,7 @@ def setup(
                         hist.axis.Variable(run_edges + 0.5, name="run"),
                         add_trailing=False,
                     ),
-                    lumis[:, *[np.newaxis for a in h.axes]],
+                    lumis[(slice(None),) + (np.newaxis,) * len(h.axes)],
                 )
             )
         )
@@ -2842,8 +2921,9 @@ def outputFolderName(outfolder, datagroups, doStatOnly, postfix):
 
 
 if __name__ == "__main__":
-    parser = make_parser()
-    args = parser.parse_args()
+    argv = _normalize_negative_imaginary_bounds(sys.argv[1:])
+    parser = make_parser(argv=argv)
+    args = parser.parse_args(argv)
 
     logger = logging.setup_logger(__file__, args.verbose, args.noColorLogger)
 
@@ -3042,7 +3122,6 @@ def outputFolderName(outfolder, datagroups, doStatOnly, postfix):
                     datagroups,
                     "Z",
                     True,
-                    False,
                     poi_axes=poi_axes,
                     prior_norm=args.priorNormXsec,
                     scale_norm=args.scaleNormXsecHistYields,

wremnants/postprocessing/datagroups/datagroups.py

@@ -547,21 +547,21 @@ def loadHistsForDatagroups(
                     )
                     h = preOpMap[member.name](h, **preOpArgs)
 
-                sum_axes = [x for x in self.sum_gen_axes if x in h.axes.name]
-                if len(sum_axes) > 0:
+                if self.globalAction:
+                    logger.debug("Applying global action")
+                    h = self.globalAction(h)
+
+                sum_gen_axes = [x for x in self.sum_gen_axes if x in h.axes.name]
+                if len(sum_gen_axes) > 0:
                     # sum over remaining axes (avoid integrating over fit axes & fakerate axes)
-                    logger.debug(f"Sum over axes {sum_axes}")
-                    h = h.project(*[x for x in h.axes.name if x not in sum_axes])
+                    logger.debug(f"Sum over axes {sum_gen_axes}")
+                    h = h.project(*[x for x in h.axes.name if x not in sum_gen_axes])
                     logger.debug(f"Hist axes are now {h.axes.name}")
 
                 if h_id == id(h):
                     logger.debug(f"Make explicit copy")
                     h = h.copy()
 
-                if self.globalAction:
-                    logger.debug("Applying global action")
-                    h = self.globalAction(h)
-
                 if forceNonzero:
                     logger.debug("force non zero")
                     h = hh.clipNegativeVals(h, createNew=False)
@@ -1038,11 +1038,12 @@ def set_rebin_action(
         rename=True,
     ):
         if len(ax_lim):
-            if not all(x.real == 0 or x.imag == 0 for x in ax_lim):
+            specified_ax_lim = [x for x in ax_lim if x is not None]
+            if not all(x.real == 0 or x.imag == 0 for x in specified_ax_lim):
                 raise ValueError(
                     "In set_rebin_action(): ax_lim only accepts pure real or imaginary numbers"
                 )
-            if any(x.imag == 0 and (x.real % 1) != 0.0 for x in ax_lim):
+            if any(x.imag == 0 and (x.real % 1) != 0.0 for x in specified_ax_lim):
                 raise ValueError(
                     "In set_rebin_action(): ax_lim requires real numbers to be of integer type"
                 )