diff --git a/bin/rabbit_fit.py b/bin/rabbit_fit.py index 44893f1..847e43c 100755 --- a/bin/rabbit_fit.py +++ b/bin/rabbit_fit.py @@ -487,6 +487,10 @@ def main(): "procs": ifitter.indata.procs, "pois": ifitter.poi_model.pois, "nois": ifitter.parms[ifitter.poi_model.npoi :][indata.noiidxs], + "systs": ifitter.indata.systs, + "systgroups": ifitter.indata.systgroups, + "systgroupidxs": ifitter.indata.systgroupidxs, + "systsnoconstraint": ifitter.indata.systsnoconstraint, } with workspace.Workspace( diff --git a/bin/rabbit_plot_hists.py b/bin/rabbit_plot_hists.py index ce4e352..a3ae402 100755 --- a/bin/rabbit_plot_hists.py +++ b/bin/rabbit_plot_hists.py @@ -331,6 +331,30 @@ def parseArgs(): Additional varNames can be specified to add variations from the nominal input. """, ) + parser.add_argument( + "--varGroupNames", + type=str, + nargs="*", + default=None, + help=( + "Variation group names to build on-the-fly from nuisance variations " + "using fitresult meta mapping (systs/systgroups/systgroupidxs)." + ), + ) + parser.add_argument( + "--varGroupLabels", + type=str, + nargs="*", + default=None, + help="Label(s) for --varGroupNames.", + ) + parser.add_argument( + "--varGroupColors", + type=str, + nargs="*", + default=None, + help="Color(s) for --varGroupNames.", + ) parser.add_argument( "--varLabels", type=str, @@ -1174,6 +1198,9 @@ def make_plots( varFilesFitTypes=None, varMarkers=None, varNames=None, + varGroupNames=None, + varGroupLabels=None, + varGroupColors=None, varLabels=None, varColors=None, binwnorm=None, @@ -1214,13 +1241,16 @@ def make_plots( l if p not in args.suppressProcsLabel else None for l, p in zip(labels, procs) ] - if varNames is not None: + if varNames is not None or varGroupNames is not None: # take the first variations from the varFiles, empty if no varFiles are provided if len(varFilesFitTypes) == 1: varFilesFitTypes = varFilesFitTypes * len(varResults) hists_down = [] hists_up = [] + plot_var_names = [] if varNames is None else list(varNames) + plot_var_labels = [] if varLabels is None else list(varLabels) + plot_var_colors = [] if varColors is None else list(varColors) for r, t in zip(varResults, varFilesFitTypes): h = r[f"hist_{t}_inclusive"].get() @@ -1237,12 +1267,11 @@ def make_plots( hists_up.append(hist_up) # take the next variations from the nominal input file - if len(varNames) > len(varResults): + if varNames is not None and len(varNames) > len(varResults): # variations from the nominal input file hist_var = result[ f"hist_{fittype}_inclusive_variations{'_correlated' if args.correlatedVariations else ''}" ].get() - hists_down.extend( [ hist_var[{"downUpVar": 0, "vars": n}].project( @@ -1259,6 +1288,65 @@ def make_plots( for n in varNames[len(varResults) :] ] ) + + # grouped variations built on-the-fly from grouped global impacts + # (covariance-based decomposition). + if varGroupNames is not None and len(varGroupNames): + axis_names = [a.name for a in axes] + h_nom = result[f"hist_{fittype}_inclusive"].get().project(*axis_names) + name_impacts_grouped = f"hist_{fittype}_inclusive_global_impacts_grouped" + + if name_impacts_grouped not in result.keys(): + raise KeyError( + f"--varGroupNames requires '{name_impacts_grouped}' in the selected result/channel. " + "Run rabbit_fit with --computeHistImpacts to produce grouped impact histograms." + ) + h_impacts_grouped = result[name_impacts_grouped].get() + impact_entries = {x for x in h_impacts_grouped.axes["impacts"]} + + for ig, gname in enumerate(varGroupNames): + if gname not in impact_entries: + logger.warning( + f"Group '{gname}' has no matching entry in grouped impacts axis; skipping." + ) + continue + sigma = ( + h_impacts_grouped[{"impacts": gname}].project(*axis_names).values() + ) + + h_up = h_nom.copy() + h_down = h_nom.copy() + h_up.values()[...] = h_nom.values() + sigma + h_down.values()[...] = h_nom.values() - sigma + hists_up.append(h_up) + hists_down.append(h_down) + + plot_var_names.append(gname) + if varGroupLabels is not None and ig < len(varGroupLabels): + plot_var_labels.append(varGroupLabels[ig]) + else: + plot_var_labels.append(gname) + if varGroupColors is not None and ig < len(varGroupColors): + plot_var_colors.append(varGroupColors[ig]) + + # Ensure metadata arrays align with the effective number of plotted variations. + n_total = len(hists_up) + if len(plot_var_names) < n_total: + plot_var_names.extend( + [f"var{i}" for i in range(len(plot_var_names), n_total)] + ) + if len(plot_var_labels) < n_total: + plot_var_labels.extend(plot_var_names[len(plot_var_labels) : n_total]) + if len(plot_var_colors) < n_total: + default_cols = [ + colormaps["tab10" if n_total < 10 else "tab20"](i) + for i in range(n_total) + ] + plot_var_colors.extend(default_cols[len(plot_var_colors) : n_total]) + + varNames = plot_var_names + varLabels = plot_var_labels + varColors = plot_var_colors else: hists_down = None hists_up = None @@ -1418,6 +1506,7 @@ def main(): varFiles = args.varFiles varNames = args.varNames + varGroupNames = args.varGroupNames varLabels = args.varLabels varColors = args.varColors if varNames is not None: @@ -1434,6 +1523,17 @@ def main(): colormaps["tab10" if len(varNames) < 10 else "tab20"](i) for i in range(len(varNames)) ] + varGroupLabels = args.varGroupLabels + varGroupColors = args.varGroupColors + if varGroupNames is not None: + if varGroupLabels is not None and len(varGroupLabels) != len(varGroupNames): + raise ValueError( + "Must specify the same number of args for --varGroupNames and --varGroupLabels" + ) + if varGroupColors is not None and len(varGroupColors) != len(varGroupNames): + raise ValueError( + "Must specify the same number of args for --varGroupNames and --varGroupColors" + ) fittype = "prefit" if args.prefit else "postfit" @@ -1470,6 +1570,9 @@ def main(): meta=meta, fittype=fittype, varNames=varNames, + varGroupNames=varGroupNames, + varGroupLabels=varGroupLabels, + varGroupColors=varGroupColors, varLabels=varLabels, varColors=varColors, varMarkers=args.varMarkers,