diff --git a/bin_opt_2d/rebin_2d.py b/bin_opt_2d/rebin_2d.py index 9132a16..8be247e 100644 --- a/bin_opt_2d/rebin_2d.py +++ b/bin_opt_2d/rebin_2d.py @@ -2,7 +2,9 @@ import json import math import os +import re import sys +import numpy as np import yaml if __name__ == "__main__": @@ -14,9 +16,17 @@ sys.path.append(base_dir) __package__ = pkg_dir_name -from StatInference.common.tools import importROOT, CategoryNaming +from StatInference.common.tools import importROOT, CategoryNaming, InputCategories from StatInference.common.param_parse import extractParameters, applyParameters from StatInference.dc_make.model import Model +from StatInference.common.binning_dp import ( + binning_objective, + build_cells, + window_tables, + check_window_mask, + partition_dp, + trim_by_marginal_gain, +) ROOT = importROOT() @@ -46,12 +56,31 @@ "min_slice_bkg_each_neff": 0.0, "min_bin_bkg_each": 0.01, "min_bin_bkg_neff": 4.0, + "min_bin_bkg_each_neff": 0.0, "bkg_per_bin": 5.0, "min_bkg_frac": 0.05, "min_signal": 0.5, "significance_mode": "asimov", + # "greedy": slice x, equal-signal bins along y. "window": one window per category on + # window_axis, the other axis binned inside it by an exact search (_discover_window). + "strategy": "greedy", + # window only: the axis the window is cut on, "y" or "x" + "window_axis": "y", + # window only: a bin is kept only if it raises the category's Z^2 by at least this + # fraction of its total (0 = keep up to max_bins_per_slice) + "dp_min_bin_gain": 0.0, + # window only: window edges are scanned on this stride, then refined within + # +-window_refine of the best + "window_stride": 6, + "window_refine": 5, + # window only: edges are searched between these signal quantiles + "window_signal_quantile": 0.001, } +# checked for the yaml as well as by the command line's choices= +BINNING_STRATEGIES = ("greedy", "window") +WINDOW_AXES = ("y", "x") + def load_binning_config(path, overrides=None): """Merge the binning yaml over the defaults, then command-line overrides over that. @@ -82,6 +111,16 @@ def load_binning_config(path, overrides=None): f"{path or 'binning configuration'}: significance_mode " f"'{knobs['significance_mode']}' is not one of {sorted(SIGNIFICANCE_MODES)}." ) + if knobs["strategy"] not in BINNING_STRATEGIES: + raise RuntimeError( + f"{path or 'binning configuration'}: strategy '{knobs['strategy']}' is not " + f"one of {sorted(BINNING_STRATEGIES)}." + ) + if knobs["window_axis"] not in WINDOW_AXES: + raise RuntimeError( + f"{path or 'binning configuration'}: window_axis '{knobs['window_axis']}' is " + f"not one of {sorted(WINDOW_AXES)}." + ) return knobs @@ -151,6 +190,8 @@ def load_config(config_path): # it against the pattern the shapes are actually written with. None when the # configuration names no pattern, and then there is nothing to check. "category_pattern": cfg.get("category_pattern"), + # where each base category is read from, if the input keys it per model point + "input_categories": InputCategories.fromConfig(cfg), "signal_hist_name_patterns": signal_hist_names, "signal_param_name": extractParameters(signal_hist_names[0])[0], "mass_values": mass_values, @@ -336,15 +377,17 @@ def _slice_passes( if min_neff > 0 and total_error is not None: if effective_entries(total, total_error) < min_neff: return False - if min_each > 0 or min_proc_neff > 0: - for name, value in yields.items(): - if name in exempt: - continue - if value <= min_each: + for name, value in yields.items(): + # exemption covers the magnitude floor, never the sign (see _bin_passes) + if value < 0: + return False + if name in exempt: + continue + if min_each > 0 and value <= min_each: + return False + if min_proc_neff > 0 and errors is not None: + if effective_entries(value, errors.get(name, 0.0)) < min_proc_neff: return False - if min_proc_neff > 0 and errors is not None: - if effective_entries(value, errors.get(name, 0.0)) < min_proc_neff: - return False return True @@ -368,7 +411,15 @@ def minor_backgrounds(bkg_hists_by_name, lo, hi, min_frac): return {name for name, value in yields.items() if value < min_frac * total} -def _bin_passes(yields, min_each, exempt=(), total_error=None, min_neff=0.0): +def _bin_passes( + yields, + min_each, + exempt=(), + total_error=None, + min_neff=0.0, + errors=None, + min_proc_neff=0.0, +): """Mass-bin validity: every *relevant* background must exceed min_each, and (when min_neff > 0) the summed background must be known to at least min_neff effective MC entries. @@ -392,10 +443,18 @@ def _bin_passes(yields, min_each, exempt=(), total_error=None, min_neff=0.0): if effective_entries(total, total_error) < min_neff: return False for name, value in yields.items(): + # Positivity is never exempt: an exempt process that goes negative in a + # signal-relevant bin makes resolveNegativeBins reject the datacard. + if value < 0: + return False if name in exempt: continue if value <= min_each: return False + # each background must be measured, not only the sum + if min_proc_neff > 0 and errors is not None: + if effective_entries(value, errors.get(name, 0.0)) < min_proc_neff: + return False return True @@ -551,7 +610,13 @@ def signal_quantile_ranges(sig_hist, n_bins, first_bin, last_bin): def merge_until_valid( - ranges, sig_hist, bkg_hists_by_name, min_each, exempt=(), min_neff=0.0 + ranges, + sig_hist, + bkg_hists_by_name, + min_each, + exempt=(), + min_neff=0.0, + min_proc_neff=0.0, ): """Merge adjacent ranges until every one satisfies the background gates. @@ -572,6 +637,8 @@ def merge_until_valid( exempt, _total_bkg_error(bkg_hists_by_name, lo, hi), min_neff, + _bkg_errors(bkg_hists_by_name, lo, hi) if min_proc_neff > 0 else None, + min_proc_neff, ): bad = i break @@ -599,6 +666,7 @@ def find_bins( min_each, min_frac=0.0, min_neff=0.0, + min_proc_neff=0.0, ): """Bins inside one slice: signal quantiles for the edges, background gates for the count. @@ -620,7 +688,7 @@ def find_bins( exempt = minor_backgrounds(bkg_hists_by_name, first_bin, last_bin, min_frac) ranges = signal_quantile_ranges(sig_hist, max_bins, first_bin, last_bin) return merge_until_valid( - ranges, sig_hist, bkg_hists_by_name, min_each, exempt, min_neff + ranges, sig_hist, bkg_hists_by_name, min_each, exempt, min_neff, min_proc_neff ) @@ -657,21 +725,7 @@ def bin_budget(bkg_hists_by_name, lo, hi, max_bins_per_slice, bkg_per_bin): return max(1, min(max_bins_per_slice, int(total / bkg_per_bin))) -def discover_binning( - sig2d, - bkg2d_by_name, - n_slices, - max_bins_per_slice, - min_slice_sum, - min_bin_each, - min_slice_bkg_neff=0.0, - min_bkg_frac=0.0, - min_bin_bkg_neff=0.0, - bkg_per_bin=0.0, - sig_mode="sb", - min_slice_bkg_each=0.0, - min_slice_bkg_each_neff=0.0, -): +def _discover_greedy(sig2d, bkg2d_by_name, knobs): """bkg2d_by_name: {background_name: [hist per discovery era, ...]}. The list is usually a single era's own histogram (standalone limit) or all of a meta-era's sub-eras (combined limit) -- see --discovery-eras. @@ -679,6 +733,19 @@ def discover_binning( the same discovery reference's (already-summed) signal histogram: its x projection picks the significance-maximizing slice boundaries, and its y projection within each slice places the mass bin edges by signal quantile.""" + n_slices = knobs["n_slices"] + max_bins_per_slice = knobs["max_bins_per_slice"] + min_slice_sum = knobs["min_slice_bkg_sum"] + min_bin_each = knobs["min_bin_bkg_each"] + min_slice_bkg_neff = knobs["min_slice_bkg_neff"] + min_bkg_frac = knobs["min_bkg_frac"] + min_bin_bkg_neff = knobs["min_bin_bkg_neff"] + bkg_per_bin = knobs["bkg_per_bin"] + sig_mode = knobs["significance_mode"] + min_slice_bkg_each = knobs["min_slice_bkg_each"] + min_slice_bkg_each_neff = knobs["min_slice_bkg_each_neff"] + min_bin_bkg_each_neff = knobs["min_bin_bkg_each_neff"] + any_hist = next(iter(bkg2d_by_name.values()))[0] nx = any_hist.GetNbinsX() ny = any_hist.GetNbinsY() @@ -728,12 +795,138 @@ def discover_binning( min_bin_each, min_bkg_frac, min_bin_bkg_neff, + min_bin_bkg_each_neff, ) bin_ranges = extend_outer_edges(bin_ranges, 0, ny + 1) result.append({"x_range": (xlo, xhi), "y_ranges": bin_ranges}) return result +def _best_count(values, cap): + """The bin count at or below `cap` that scores best, or None if none is feasible. + + Not simply the largest: with the background error in the figure of merit, an extra + split can lower Z^2. + """ + best, best_n = -np.inf, None + for n in range(1, min(cap, len(values) - 1) + 1): + if values[n] > best: + best, best_n = values[n], n + return best_n + + +def _window_value(cells, y0, y1, knobs, check=False): + """(Z^2, bins) of the best binning inside window y0..y1 (in the cells' frame). + + A window is scored by the binning it admits, not by its integrated significance, + which would favour windows that swallow signal and ignore the shape. + """ + mode = knobs["significance_mode"] + # negligible backgrounds are judged once, over the whole window + exempt = cells.exempt(1, cells.nx, y0, y1, knobs["min_bkg_frac"]) + score, valid = window_tables(cells, y0, y1, exempt, knobs, mode) + if check: + + def bin_passes(a, b, ylo, yhi, exempt_set): + return _bin_passes( + cells.yields(a, b, ylo, yhi), + knobs["min_bin_bkg_each"], + exempt_set, + cells.total_bkg_error(a, b, ylo, yhi), + knobs["min_bin_bkg_neff"], + cells.errors(a, b, ylo, yhi), + knobs["min_bin_bkg_each_neff"], + ) + + check_window_mask( + cells, + y0, + y1, + valid, + score, + exempt, + bin_passes, + mode, + ) + values, partitions = partition_dp(score, valid, knobs["max_bins_per_slice"]) + n = _best_count(values, knobs["max_bins_per_slice"]) + if n is None: + return float("-inf"), None + if knobs["dp_min_bin_gain"] > 0 and values[n] > 0: + n = trim_by_marginal_gain([values], [n], knobs["dp_min_bin_gain"] * values[n])[ + 0 + ] + return values[n], [(a + 1, b + 1) for a, b in partitions[n]] + + +def _discover_window(sig2d, bkg2d_by_name, knobs): + """One window per category on window_axis, with the other axis binned inside it. + + Events outside the window are dropped. Suited to a window axis the signal peaks in + (HH->bbWW: HME on y, DNN on x); only one region of it is kept. + + The search is written with the window on y; for window_axis x the cells are built + transposed and the result is mapped back to the input's axes. + """ + transpose = knobs["window_axis"] == "x" + cells = build_cells(sig2d, bkg2d_by_name, significance, transpose=transpose) + ny = cells.ny + # only search where the signal is + cumulative = np.array( + [cells.signal(0, cells.nx + 1, 1, j) for j in range(1, ny + 1)] + ) + total = cumulative[-1] if len(cumulative) else 0.0 + if total <= 0: + return None + q = knobs["window_signal_quantile"] + lo_edge = max(1, int(np.searchsorted(cumulative, q * total))) + hi_edge = min(ny, int(np.searchsorted(cumulative, (1.0 - q) * total)) + 1) + + best_value, best_box = float("-inf"), None + stride = max(1, int(knobs["window_stride"])) + for a in range(lo_edge, hi_edge + 1, stride): + for b in range(a, hi_edge + 1, stride): + value, _ = _window_value(cells, a, b, knobs) + if value > best_value: + best_value, best_box = value, (a, b) + if best_box is None: + return None + refine = max(0, int(knobs["window_refine"])) + a0, b0 = best_box + for a in range(max(lo_edge, a0 - refine), min(hi_edge, a0 + refine) + 1): + for b in range(max(a, b0 - refine), min(hi_edge, b0 + refine) + 1): + value, _ = _window_value(cells, a, b, knobs) + if value > best_value: + best_value, best_box = value, (a, b) + + y0, y1 = best_box + value, bins = _window_value(cells, y0, y1, knobs, check=True) + if bins is None: + return None + # a window at an end of the axis includes that end's under/overflow + y0 = 0 if y0 <= 1 else y0 + y1 = ny + 1 if y1 >= ny else y1 + if transpose: + # window on x, bins along y: the same layout as the greedy slices + return [{"x_range": (y0, y1), "y_ranges": bins}] + return [{"y_range": (y0, y1), "x_ranges": bins}] + + +def discover_binning(sig2d, bkg2d_by_name, knobs): + """Find this category's binning, by whichever strategy the configuration names. + + bkg2d_by_name: {background_name: [hist per discovery era, ...]}. The list is usually + a single era's own histogram (standalone limit) or all of a meta-era's sub-eras + (combined limit). Yields are summed across whatever is in the list; see _bkg_yields(). + sig2d is the same discovery reference's already-summed signal histogram. + + Returns None when the category cannot be binned, which the caller reports as a skip. + """ + if knobs["strategy"] == "window": + return _discover_window(sig2d, bkg2d_by_name, knobs) + return _discover_greedy(sig2d, bkg2d_by_name, knobs) + + def bin_edges(y_axis, y_ranges): """Physical edges of the discovered y ranges, for booking the output TH1. @@ -785,7 +978,7 @@ def format_var_range(lo, hi, var): return f"{lo:.2f} < {var} < {hi:.2f}" -def slice_ranges(x_axis, slices): +def slice_ranges(x_axis, y_axis, slices): """Physical edges of the discovered slices on the sliced axis, as [[lo, hi], ...]. The slice x_ranges are bin indices, and extend_outer_edges() has already pushed the @@ -794,39 +987,67 @@ def slice_ranges(x_axis, slices): the plots can say which selection each slice actually is; nothing downstream of the datacards needs it. """ - n = x_axis.GetNbins() ranges = [] for sl in slices: - lo, hi = sl["x_range"] + (lo, hi), _, slice_axis = slice_parts(sl) + axis = x_axis if slice_axis == "x" else y_axis + n = axis.GetNbins() ranges.append( [ - None if lo < 1 else x_axis.GetBinLowEdge(lo), - None if hi > n else x_axis.GetBinUpEdge(hi), + None if lo < 1 else axis.GetBinLowEdge(lo), + None if hi > n else axis.GetBinUpEdge(hi), ] ) return ranges -def slices_to_record(slices, x_axis, y_axis): +def _slice_record(sl, sel_edges, x_axis, y_axis): + """One category as plain data; the keys say which axis is selected and which binned.""" + (lo, hi), bins, slice_axis = slice_parts(sl) + binned_axis = y_axis if slice_axis == "x" else x_axis + sel_key, bin_key = ( + ("x_range", "y_ranges") + if slice_axis == "x" + else ( + "y_range", + "x_ranges", + ) + ) + return { + sel_key: [lo, hi], + f"{slice_axis}_edges": sel_edges, + bin_key: [list(r) for r in bins], + ("y_edges" if slice_axis == "x" else "x_edges"): bin_edges(binned_axis, bins), + } + + +def slices_to_record(slices, x_axis, y_axis, objective=None): """The discovered structure as plain data, for binning.json. Both forms are written. The bin index ranges are what the code actually cuts on and are what a replay restores; the physical edges alongside them are what a human reads, and what makes the record meaningful next to a plot. + + `objective` (the binning's Z^2) is written for information only; replay reads only + the ranges. """ - return { + record = { "n_x_bins": x_axis.GetNbins(), "n_y_bins": y_axis.GetNbins(), "slices": [ - { - "x_range": list(sl["x_range"]), - "x_edges": x_edges, - "y_ranges": [list(r) for r in sl["y_ranges"]], - "y_edges": bin_edges(y_axis, sl["y_ranges"]), - } - for sl, x_edges in zip(slices, slice_ranges(x_axis, slices)) + _slice_record(sl, sel_edges, x_axis, y_axis) + for sl, sel_edges in zip(slices, slice_ranges(x_axis, y_axis, slices)) ], } + if objective is not None: + total, per_slice = objective + record["objective"] = { + "z": math.sqrt(max(total, 0.0)), + "z2": total, + "z2_per_slice": per_slice, + "n_bins": sum(len(slice_parts(sl)[1]) for sl in slices), + } + return record def record_to_slices(record, x_axis, y_axis, where): @@ -848,30 +1069,49 @@ def record_to_slices(record, x_axis, y_axis, where): "to this input." ) return [ - { - "x_range": tuple(sl["x_range"]), - "y_ranges": [tuple(r) for r in sl["y_ranges"]], - } + ( + { + "y_range": tuple(sl["y_range"]), + "x_ranges": [tuple(r) for r in sl["x_ranges"]], + } + if "y_range" in sl + else { + "x_range": tuple(sl["x_range"]), + "y_ranges": [tuple(r) for r in sl["y_ranges"]], + } + ) for sl in record["slices"] ] +def slice_parts(sl): + """(selection range, bin ranges, which axis the selection is on) for one category. + + Layouts: {"x_range", "y_ranges"} (greedy slices, window on x) and + {"y_range", "x_ranges"} (window on y). + """ + if "y_range" in sl: + return sl["y_range"], sl["x_ranges"], "y" + return sl["x_range"], sl["y_ranges"], "x" + + def rebin_hist_2d(hist2d, slices, name, naming): """Given the discovered slice structure, produce one final TH1 per slice for this specific histogram (nominal or a systematic variation).""" outputs = [] for slice_idx, sl in enumerate(slices): - xlo, xhi = sl["x_range"] + (sel_lo, sel_hi), bin_ranges, slice_axis = slice_parts(sl) # ROOT reinterprets an inverted or negative range as the full axis including # under/overflow, silently and with a plausible-looking positive yield, so an # invalid range must never reach IntegralAndError below. Checked here rather # than trusted because this is the last point where it is still cheap to say so. - if not 0 <= xlo <= xhi: + if not 0 <= sel_lo <= sel_hi: raise RuntimeError( - f"invalid x bin range ({xlo}, {xhi}) for slice {slice_idx} of {name}; " - "ROOT would read this as the whole plane." + f"invalid {slice_axis} bin range ({sel_lo}, {sel_hi}) for slice " + f"{slice_idx} of {name}; ROOT would read this as the whole plane." ) - edges = array.array("d", bin_edges(hist2d.GetYaxis(), sl["y_ranges"])) + binned_axis = hist2d.GetYaxis() if slice_axis == "x" else hist2d.GetXaxis() + edges = array.array("d", bin_edges(binned_axis, bin_ranges)) # Detached for the same reason as the projections above: Write() targets # gDirectory regardless, so nothing needs these to stay attached. h = _detach( @@ -882,15 +1122,35 @@ def rebin_hist_2d(hist2d, slices, name, naming): edges, ) ) - for bin_idx, (ylo, yhi) in enumerate(sl["y_ranges"], start=1): + for bin_idx, (blo, bhi) in enumerate(bin_ranges, start=1): err = array.array("d", [0.0]) - content = hist2d.IntegralAndError(xlo, xhi, ylo, yhi, err) + # the selection is on one axis and the bin on the other; IntegralAndError + # always wants (x_lo, x_hi, y_lo, y_hi) + if slice_axis == "x": + content = hist2d.IntegralAndError(sel_lo, sel_hi, blo, bhi, err) + else: + content = hist2d.IntegralAndError(blo, bhi, sel_lo, sel_hi, err) h.SetBinContent(bin_idx, content) h.SetBinError(bin_idx, err[0]) outputs.append(h) return outputs +def other_point_signal(key, patterns, param_name, value): + """Whether `key` is a signal histogram (nominal or variation) of a model point other + than `value`. A rebinned file is only read for its own point, so those are dropped. + """ + placeholder = "${" + param_name + "}" + for pattern in patterns: + head, found, tail = pattern.partition(placeholder) + if not found: + continue + match = re.match(re.escape(head) + r"(\d+)" + re.escape(tail) + r"(?:_|$)", key) + if match and match.group(1) != str(value): + return True + return False + + def process_category( sources, channel, @@ -922,12 +1182,15 @@ def process_category( # The resonance parameter is named by the configuration, not by this script -- it is # "MX" for bbWW but the slicing knows nothing about which parameter it is scanning. param_name = cfg["signal_param_name"] - prefix = f"{channel}/{category}/" + # read from where the input keeps this category for this mass (see InputCategories); + # written under the category's own name + in_dir = f"{channel}/{cfg['input_categories'].path(category, {param_name: mass})}" + prefix = f"{in_dir}/" # The key list comes from the first source era; a systematic that only some eras carry # is filled in from their nominal when the slices are written below. - cat_dir = sources[0][1].Get(f"{channel}/{category}") + cat_dir = sources[0][1].Get(in_dir) if not cat_dir: - print(f" [skip] {channel}/{category}: not found in {sources[0][1].GetName()}") + print(f" [skip] {in_dir}: not found in {sources[0][1].GetName()}") return signal_keys = [ @@ -989,21 +1252,7 @@ def load2d(f, key): frozen, disc_sig.GetXaxis(), disc_sig.GetYaxis(), where ) else: - slices = discover_binning( - disc_sig, - disc_bkg_by_name, - knobs["n_slices"], - knobs["max_bins_per_slice"], - knobs["min_slice_bkg_sum"], - knobs["min_bin_bkg_each"], - knobs["min_slice_bkg_neff"], - knobs["min_bkg_frac"], - knobs["min_bin_bkg_neff"], - knobs["bkg_per_bin"], - knobs["significance_mode"], - knobs["min_slice_bkg_each"], - knobs["min_slice_bkg_each_neff"], - ) + slices = discover_binning(disc_sig, disc_bkg_by_name, knobs) if slices is None: print( f" [skip] {channel}/{category} {param_name}={mass}: the " @@ -1019,7 +1268,7 @@ def load2d(f, key): # side-car path to hand a reader correctly and no key name for the two ends to agree # on -- anything that can open the shapes can already read it. naming = cfg["naming"] - ranges = slice_ranges(disc_sig.GetXaxis(), slices) + ranges = slice_ranges(disc_sig.GetXaxis(), disc_sig.GetYaxis(), slices) # One shared set of edges, applied to every source era in its own file. for source_era, in_file, out_file in sources: @@ -1029,11 +1278,15 @@ def load2d(f, key): f"{channel}/{naming.name(category, slice_idx)}", format_var_range(*ranges[slice_idx], var=cfg["slice_var"]), ) - cat_dir = in_file.Get(f"{channel}/{category}") + cat_dir = in_file.Get(in_dir) if not cat_dir: - print(f" [skip] {source_era} {channel}/{category}: not in the input") + print(f" [skip] {source_era} {in_dir}: not in the input") continue for key in [k.GetName() for k in cat_dir.GetListOfKeys()]: + if other_point_signal( + key, cfg["signal_hist_name_patterns"], param_name, mass + ): + continue hist2d = get_hist(in_file, prefix + key) if hist2d is None or hist2d.GetDimension() != 2: continue @@ -1041,7 +1294,18 @@ def load2d(f, key): out_file.cd(f"{channel}/{naming.name(category, slice_idx)}") h.Write(key) - return slices_to_record(slices, disc_sig.GetXaxis(), disc_sig.GetYaxis()) + # the Z^2 of the written binning, for the record; skipped without backgrounds (a + # replay can run on input missing one) + objective = ( + binning_objective( + build_cells(disc_sig, disc_bkg_by_name, significance), + slices, + knobs["significance_mode"], + ) + if disc_bkg_by_name + else None + ) + return slices_to_record(slices, disc_sig.GetXaxis(), disc_sig.GetYaxis(), objective) def run( @@ -1285,6 +1549,13 @@ def run( float, "minimum effective MC entries of the summed background in a bin", ), + "min_bin_bkg_each_neff": ( + float, + "minimum effective MC entries of every non-negligible background in a bin. " + "The summed test above is satisfied by any one well-measured process, so " + "this is what makes each background individually measured; bins are merged " + "into their lower-signal neighbour until it holds", + ), "bkg_per_bin": ( float, "summed background to aim for per bin; overrides max_bins_per_slice when it " @@ -1305,6 +1576,24 @@ def run( default=None, help=help_text, ) + parser.add_argument( + "--strategy", + required=False, + type=str, + default=None, + choices=list(BINNING_STRATEGIES), + help="how the edges are found: 'greedy' = each slice boundary taken on its own " + "and equal-signal quantiles inside it, 'window' = one window per category on " + "--window-axis with the other axis binned inside it by an exact partition search", + ) + parser.add_argument( + "--window-axis", + required=False, + type=str, + default=None, + choices=list(WINDOW_AXES), + help="for strategy 'window': the axis the window is cut on (default y)", + ) parser.add_argument( "--significance-mode", required=False, @@ -1318,6 +1607,8 @@ def run( overrides = {name: getattr(args, name) for name in knob_args} overrides["significance_mode"] = args.significance_mode + overrides["strategy"] = args.strategy + overrides["window_axis"] = args.window_axis knobs = load_binning_config(args.binning_config, overrides) frozen_binning = None diff --git a/common/binning_dp.py b/common/binning_dp.py new file mode 100644 index 0000000..a48d4c1 --- /dev/null +++ b/common/binning_dp.py @@ -0,0 +1,383 @@ +"""Prefix-sum tables over a 2D shape, so the window search can score any bin in O(1). + +The window search tries ~ny^2 windows, each with an exact partition search over ~nx^2 +candidate bins, which is only affordable with table lookups instead of TH2 integrals. +The search is written with the window on y; build_cells transposes the input for a +window on x. + +The gate (_bin_passes) and the figure of merit (significance) stay in rebin_2d.py and are +passed in, so each has one definition. + +Notes: +- Under/overflow are included (arrays are indexed 0..n+1), to match what the TH2 + integrals in rebin_2d.py return. +- Differencing prefix sums can turn an exact 0 into -2e-16, which the `value < 0` + positivity gate would reject. Yields are snapped to 0 below an epsilon scaled by each + process's L1 norm (~1e-8 events, far below the 0.01 floor). +- check_window_mask() re-checks a sample of every chosen window against _bin_passes and + significance on each run. +""" + +import math + +import numpy as np + + +def _hist_arrays(hist): + """(values, variances) as (nx+2, ny+2) arrays indexed [binx, biny], bins 0..n+1.""" + nx = hist.GetNbinsX() + ny = hist.GetNbinsY() + values = np.empty((nx + 2, ny + 2), dtype=np.float64) + variances = np.empty((nx + 2, ny + 2), dtype=np.float64) + for bx in range(nx + 2): + for by in range(ny + 2): + values[bx, by] = hist.GetBinContent(bx, by) + err = hist.GetBinError(bx, by) + variances[bx, by] = err * err + return values, variances + + +def _prefix(a): + """Inclusive 2D prefix sums, shape (nx+3, ny+3), with P[i, j] = sum(a[:i, :j]).""" + p = np.zeros((a.shape[0] + 1, a.shape[1] + 1), dtype=np.float64) + p[1:, 1:] = a.cumsum(axis=0).cumsum(axis=1) + return p + + +def _cancellation_eps(a, rel=1e-12): + """How large a value this array's prefix sums can invent out of rounding. + + Scaled by the L1 norm, not the sum: with negative weights a process can nearly cancel + to a small sum while its prefix sums still lose precision on the L1 scale. + """ + return rel * float(np.abs(a).sum()) + + +def _snap(value, eps): + """Zero out a value that is indistinguishable from zero at this array's precision.""" + return 0.0 if abs(value) < eps else value + + +def _rect(p, x0, x1, y0, y1): + """Sum over the inclusive bin-index rectangle [x0..x1] x [y0..y1].""" + return float(p[x1 + 1, y1 + 1] - p[x0, y1 + 1] - p[x1 + 1, y0] + p[x0, y0]) + + +class Cells: + """Every quantity the binning gates and the figure of merit need, as O(1) lookups. + + Backgrounds are summed across discovery eras (errors in quadrature), as + _bkg_yields()/_bkg_errors() do. + """ + + def __init__(self, nx, ny, sig, bkg, var, significance): + self.nx = nx + self._significance = significance + self.ny = ny + self.names = list(bkg) + self._sig = _prefix(sig) + self._bkg = {name: _prefix(bkg[name]) for name in self.names} + self._var = {name: _prefix(var[name]) for name in self.names} + # no backgrounds at all is possible (a replay whose input lacks them); keep arrays + zero = np.zeros_like(sig) + bkg_tot = sum(bkg.values()) if bkg else zero + var_tot = sum(var.values()) if var else zero + self._bkg_tot = _prefix(bkg_tot) + self._var_tot = _prefix(var_tot) + self._eps_sig = _cancellation_eps(sig) + self._eps_bkg = {name: _cancellation_eps(bkg[name]) for name in self.names} + self._eps_tot = _cancellation_eps(bkg_tot) + self._eps_var = {name: _cancellation_eps(var[name]) for name in self.names} + self._eps_var_tot = _cancellation_eps(var_tot) + + def signal(self, x0, x1, y0, y1): + return _snap(_rect(self._sig, x0, x1, y0, y1), self._eps_sig) + + def total_bkg(self, x0, x1, y0, y1): + return _snap(_rect(self._bkg_tot, x0, x1, y0, y1), self._eps_tot) + + def total_bkg_error(self, x0, x1, y0, y1): + return math.sqrt(max(_rect(self._var_tot, x0, x1, y0, y1), 0.0)) + + def yields(self, x0, x1, y0, y1): + """{background: yield}, the dict the gates take.""" + return { + n: _snap(_rect(self._bkg[n], x0, x1, y0, y1), self._eps_bkg[n]) + for n in self.names + } + + def errors(self, x0, x1, y0, y1): + """{background: MC statistical error}, the dict the gates take.""" + return { + n: math.sqrt(max(_rect(self._var[n], x0, x1, y0, y1), 0.0)) + for n in self.names + } + + def score(self, x0, x1, y0, y1, mode): + """significance()^2 for one rectangle; Z^2 adds across bins, so a partition's + value is the sum of its bins'.""" + s = self.signal(x0, x1, y0, y1) + b = self.total_bkg(x0, x1, y0, y1) + b_err = self.total_bkg_error(x0, x1, y0, y1) + return self._significance(s, b, b_err, mode) ** 2 + + def exempt(self, x0, x1, y0, y1, min_frac): + """minor_backgrounds() over the rectangle about to be subdivided.""" + if min_frac <= 0: + return set() + y = self.yields(x0, x1, y0, y1) + total = sum(y.values()) + if total <= 0: + return set() + return {n for n, v in y.items() if v < min_frac * total} + + +def build_cells(sig2d, bkg2d_by_name, significance, transpose=False): + """Cells for one channel/category. + + sig2d is the summed discovery signal; bkg2d_by_name is {background: [hist per era]}. + transpose swaps x and y, which is how the window search runs with the window on x. + """ + nx = sig2d.GetNbinsX() + ny = sig2d.GetNbinsY() + sig, _ = _hist_arrays(sig2d) + bkg = {} + var = {} + for name, hists in bkg2d_by_name.items(): + values = np.zeros((nx + 2, ny + 2), dtype=np.float64) + variances = np.zeros((nx + 2, ny + 2), dtype=np.float64) + for h in hists: + v, e2 = _hist_arrays(h) + values += v + variances += e2 + bkg[name] = values + var[name] = variances + if transpose: + nx, ny = ny, nx + sig = sig.T.copy() + bkg = {name: a.T.copy() for name, a in bkg.items()} + var = {name: a.T.copy() for name, a in var.items()} + return Cells(nx, ny, sig, bkg, var, significance) + + +NEG = -np.inf + + +def partition_dp(score, valid, max_parts): + """Best partition of [0..n-1] into exactly k contiguous valid ranges, for every k. + + Returns (values, partitions): values[k] is the best k-part total score (-inf if none) + and partitions[k] its (lo, hi) index pairs; index 0 is unused. Every k is returned + because the bin count is chosen from the whole curve. + + Exact interval DP, O(max_parts * n^2), valid because the score is additive over parts. + Invalid ranges are excluded outright, so the background gates are hard constraints. + """ + n = score.shape[0] + max_parts = max(1, min(max_parts, n)) + masked = np.where(valid, score, NEG) + + dp = np.full((max_parts + 1, n + 1), NEG) + arg = np.full((max_parts + 1, n + 1), -1, dtype=int) + dp[0, 0] = 0.0 + for k in range(1, max_parts + 1): + for e in range(k, n + 1): + # candidate: the last part is [s .. e-1], the rest is a (k-1)-part prefix + cand = dp[k - 1, :e] + masked[:e, e - 1] + s = int(np.argmax(cand)) + if cand[s] > NEG: + dp[k, e] = cand[s] + arg[k, e] = s + + values = [NEG] * (max_parts + 1) + partitions = [None] * (max_parts + 1) + for k in range(1, max_parts + 1): + if dp[k, n] == NEG: + continue + values[k] = float(dp[k, n]) + parts = [] + e = n + for kk in range(k, 0, -1): + s = int(arg[kk, e]) + parts.append((s, e - 1)) + e = s + partitions[k] = list(reversed(parts)) + return values, partitions + + +def binning_objective(cells, slices, mode): + """(total Z^2, per-slice Z^2) of a finished binning, from the recorded ranges. + + Computed for every strategy, so binnings can be compared from binning.json alone. + """ + per_slice = [] + for sl in slices: + if "y_range" in sl: # selection on y, bins along x + ylo, yhi = sl["y_range"] + per_slice.append( + sum(cells.score(a, b, ylo, yhi, mode) for a, b in sl["x_ranges"]) + ) + else: # selection on x, bins along y + xlo, xhi = sl["x_range"] + per_slice.append( + sum(cells.score(xlo, xhi, a, b, mode) for a, b in sl["y_ranges"]) + ) + return sum(per_slice), per_slice + + +def _ranges_from_column(column, n, include_outer=False): + """M[i, j] = the sum over bins (i+1)..(j+1), from one column of a prefix sum. + + The whole (n, n) table of candidate ranges is one outer subtraction. With + include_outer the first and last bins also take the under/overflow. + """ + lo = column[1 : n + 1].copy() + hi = column[2 : n + 2].copy() + if include_outer: + lo[0] = column[0] + hi[-1] = column[n + 2] + return hi[None, :] - lo[:, None] + + +def _neff_matrix(value, error): + """effective_entries() over arrays. error <= 0 gives inf if value > 0, else 0, as + the scalar does (a nan would silently fail every threshold).""" + with np.errstate(divide="ignore", invalid="ignore"): + ratio = np.where(error > 0, value / np.where(error > 0, error, 1.0), 0.0) ** 2 + degenerate = np.where(value > 0, np.inf, 0.0) + return np.where(error > 0, ratio, degenerate) + + +def _asimov_z2_matrix(s, b, b_err): + """asimov_significance(...)**2 over arrays, matching the scalar branch for branch.""" + out = np.zeros_like(s) + live = (b > 0) & (s > 0) + if not live.any(): + return out + var = b_err**2 + no_err = live & (b_err <= 0) + if no_err.any(): + ss, bb = s[no_err], b[no_err] + out[no_err] = np.maximum(2.0 * ((ss + bb) * np.log1p(ss / bb) - ss), 0.0) + with_err = live & (b_err > 0) + if with_err.any(): + ss, bb, vv = s[with_err], b[with_err], var[with_err] + term1 = (ss + bb) * np.log(((ss + bb) * (bb + vv)) / (bb * bb + (ss + bb) * vv)) + term2 = (bb * bb / vv) * np.log1p(vv * ss / (bb * (bb + vv))) + out[with_err] = np.maximum(2.0 * (term1 - term2), 0.0) + return out + + +def _sb_z2_matrix(s, b, b_err): + """significance(..., mode="sb")**2 over arrays: S^2 / (B + sigma_B^2).""" + denom = b + b_err**2 + ok = (b > 0) & (denom > 0) + with np.errstate(divide="ignore", invalid="ignore"): + return np.where(ok, s**2 / np.where(ok, denom, 1.0), 0.0) + + +def trim_by_marginal_gain(values_per_slice, counts, threshold): + """Give back every bin whose split raised Z^2 by less than `threshold`. + + Z^2 almost always grows with the bin count, so without this the search would use + every bin it is allowed. The caller sets `threshold` as a fraction of the category's + total, and the count is stepped down until the last bin pays for itself. + """ + if threshold <= 0: + return list(counts) + trimmed = [] + for values, n in zip(values_per_slice, counts): + while n > 1 and values[n] - values[n - 1] < threshold: + n -= 1 + trimmed.append(n) + return trimmed + + +def _axis_ranges(prefix, a, b, n, axis, include_outer=False): + """All ranges along `axis`, at a fixed window [a, b] on the other one.""" + if axis == "y": + column = prefix[b + 1] - prefix[a] + else: + column = prefix[:, b + 1] - prefix[:, a] + return _ranges_from_column(column, n, include_outer) + + +def window_tables(cells, y0, y1, exempt, knobs, mode): + """(score, valid) for every candidate x bin inside the window y0..y1. + + _bin_passes() and significance() restated as array operations: calling them per + candidate is far too slow for the search. check_window_mask() guards against the two + versions drifting apart. + + `exempt` (processes excused from the per-bin floors) is judged once over the whole + window. Positivity is never excused. + """ + nx = cells.nx + ranges = lambda prefix: _axis_ranges(prefix, y0, y1, nx, "x", True) + + sig = ranges(cells._sig) + b_tot = ranges(cells._bkg_tot) + b_tot = np.where(np.abs(b_tot) < cells._eps_tot, 0.0, b_tot) + v_tot = ranges(cells._var_tot) + err = np.sqrt(np.maximum(v_tot, 0.0)) + + valid = np.triu(np.ones((nx, nx), dtype=bool)) + if knobs["min_bin_bkg_neff"] > 0: + valid &= _neff_matrix(b_tot, err) >= knobs["min_bin_bkg_neff"] + for name in cells.names: + b_p = ranges(cells._bkg[name]) + b_p = np.where(np.abs(b_p) < cells._eps_bkg[name], 0.0, b_p) + valid &= b_p >= 0 + if name in exempt: + continue + magnitude = b_p > knobs["min_bin_bkg_each"] + required = knobs["min_bin_bkg_each_neff"] + if required > 0: + v_p = ranges(cells._var[name]) + magnitude &= _neff_matrix(b_p, np.sqrt(np.maximum(v_p, 0.0))) >= required + valid &= magnitude + score = ( + _asimov_z2_matrix(sig, b_tot, err) + if mode == "asimov" + else _sb_z2_matrix(sig, b_tot, err) + ) + return np.where(valid, score, 0.0), valid + + +def check_window_mask( + cells, + y0, + y1, + valid, + score, + exempt, + bin_passes, + mode, + n_samples=100, + seed=0, +): + """Re-derive a random sample of window_tables() with _bin_passes and significance(), + and stop the run on any disagreement.""" + nx = cells.nx + rng = np.random.default_rng(seed) + for _ in range(n_samples): + a = int(rng.integers(1, nx + 1)) + b = int(rng.integers(a, nx + 1)) + # the outermost bins include the x under/overflow, as in window_tables + lo = 0 if a == 1 else a + hi = nx + 1 if b == nx else b + want = bool(bin_passes(lo, hi, y0, y1, exempt)) + got = bool(valid[a - 1, b - 1]) + if got != want: + raise AssertionError( + f"binning_dp window gate disagrees with _bin_passes on bins {a}..{b} " + f"of the window {y0}..{y1}: fast path says {'valid' if got else 'invalid'}, " + f"_bin_passes says {'valid' if want else 'invalid'}." + ) + if want: + expected = cells.score(lo, hi, y0, y1, mode) + if abs(expected - score[a - 1, b - 1]) > 1e-9 * max(abs(expected), 1.0): + raise AssertionError( + f"binning_dp window score disagrees with significance() on bins " + f"{a}..{b} of the window {y0}..{y1}: {score[a - 1, b - 1]!r} vs {expected!r}." + ) diff --git a/common/tools.py b/common/tools.py index 59aeee7..e0665e5 100644 --- a/common/tools.py +++ b/common/tools.py @@ -101,6 +101,58 @@ def base(self, category): return self.split(category)[0] +class InputCategories: + """Where a category's shapes live in the input file, when that is not its own name. + + A datacard category is one name for every model point: it names the bin, the + per-category datacard directory, and what Process/Uncertainty scoping matches + against. Some inputs do not store it under one name. The HH->bbWW single-lepton + shapes hold, in the file for mass MX, a region for every mass (SR_SL_M300 ... + SR_SL_M1000), and only SR_SL_M${MX} is the one the MX hypothesis selected. Listing + those as categories would put all ten into every datacard as separate bins -- every + file has all of them, so nothing would ever report one missing. + + So the configuration maps a category to its input directory, with the model's + parameters substituted at read time: + + input_categories: + SR_SL/res2b: SR_SL_M${MX}/res2b + + Only exact names are mapped. A sliced name (SR_SL/res2b_hmebox0) is read from shapes + a preprocessing step has already written under the category's own name, so it + passes through, as does every category of a configuration that declares no map. + """ + + def __init__(self, mapping=None): + mapping = mapping or {} + if not isinstance(mapping, dict): + raise RuntimeError( + f"input_categories must map category -> input directory, got {mapping!r}" + ) + self.mapping = {str(k): str(v) for k, v in mapping.items()} + + @classmethod + def fromConfig(cls, cfg): + return cls(cfg.get("input_categories")) + + def unknown(self, categories): + """Mapped names that are not among the configuration's categories -- a typo + there would otherwise leave the category read from its own name, silently.""" + return sorted(set(self.mapping) - set(categories)) + + def path(self, category, params): + """The directory holding `category` for model point `params` ({"MX": 500}).""" + pattern = self.mapping.get(category, category) + for name, value in (params or {}).items(): + pattern = pattern.replace("${" + name + "}", str(value)) + if "${" in pattern: + raise RuntimeError( + f"input category '{pattern}' for '{category}' still has an unresolved " + f"parameter after substituting {dict(params or {})}" + ) + return pattern + + class PackageWrapper: def __init__(self, import_fn): self._package = None diff --git a/dc_make/maker.py b/dc_make/maker.py index abfd482..929ce38 100644 --- a/dc_make/maker.py +++ b/dc_make/maker.py @@ -12,6 +12,7 @@ resolveNegativeBins, getRelevantBins, CategoryNaming, + InputCategories, ) from .process import Process from .uncertainty import ( @@ -68,6 +69,23 @@ def __init__( # group the slices of one base category. self.naming = CategoryNaming.fromConfig(cfg) self.categories = list(cfg["categories"]) + # Where each category's shapes sit in the input, when that depends on the model + # point; see InputCategories. Checked against the categories as the configuration + # writes them, not self.categories: once a category list is expanded against a + # rebinning record (base "SR/res2b" -> "SR/res2b_hmebox0"), the map's keys -- the + # configured names -- no longer appear in it, and every key would read as unknown. + # A key may also be the base of a configured sliced name: that map is for + # rebin_2d.py, and the sliced shapes it writes are read under their own names. + self.input_categories = InputCategories.fromConfig(cfg) + configured = list(cfg["categories"]) + unknown = self.input_categories.unknown( + configured + [self.naming.base(c) for c in configured] + ) + if unknown: + raise RuntimeError( + f"input_categories maps {unknown}, which are neither configured " + f"categories nor their bases ({configured})" + ) self.signalFractionForRelevantBins = cfg["signalFractionForRelevantBins"] self.era_groups = cfg.get("era_groups", {}) @@ -289,6 +307,11 @@ def PPECC(self): self.processes.keys(), param_bins, self.eras, self.channels, self.categories ) + def histDir(self, channel, category, model_params): + """The input directory holding one bin's histograms. The category is the bin's + own name unless `input_categories` places it per model point.""" + return f"{channel}/{self.input_categories.path(category, model_params)}" + def getInputFile(self, era, model_params): file_name = self.model.getInputFileName(era, model_params) if file_name not in self.input_files: @@ -349,7 +372,7 @@ def getMVLnNValue(self, unc, process, era, channel, category, model_params): channel, category, model_params, - f"{channel}/{category}/{subp}", + f"{self.histDir(channel, category, model_params)}/{subp}", ) sub_yield = hist.Integral(1, hist.GetNbinsX() + 1) if sub_yield <= 0: @@ -444,7 +467,7 @@ def _getSubEraLnNVariedShapes( channel, category, model_params, - f"{channel}/{category}/{hist_name_suffix}", + f"{self.histDir(channel, category, model_params)}/{hist_name_suffix}", ) unc_value = self._getLnNValue( unc, process, proc_name_for_unc, sub_era, channel, category @@ -728,11 +751,12 @@ def _rawShape( "this bin." ) else: - hist_name = f"{channel}/{category}/{process.hist_name}" + hist_dir = self.histDir(channel, category, model_params) + hist_name = f"{hist_dir}/{process.hist_name}" hists = [] if process.subprocesses: for subp in process.subprocesses: - hist_name = f"{channel}/{category}/{subp}" + hist_name = f"{hist_dir}/{subp}" if unc_name and unc_scale: hist_name += f"_{unc_name}_{unc_scale}" subhist = self.readHist(file, hist_name) @@ -972,7 +996,9 @@ def hasNominalShape(self, process, era, channel, category): MC -- e.g. a standalone single-era limit for a sparse category/channel that only has signal statistics once combined with other eras.""" sub_eras = self.getSubEras(era) if self.isMetaEra(era) else [era] - hist_name = f"{channel}/{category}/{process.hist_name}" + hist_name = ( + f"{self.histDir(channel, category, process.params)}/{process.hist_name}" + ) for sub_era in sub_eras: _, file = self.getInputFile(sub_era, process.params) obj = self.readHist(file, hist_name) diff --git a/law/CombinedResonantLimitsTask.py b/law/CombinedResonantLimitsTask.py new file mode 100644 index 0000000..400cff3 --- /dev/null +++ b/law/CombinedResonantLimitsTask.py @@ -0,0 +1,84 @@ +import law +import luigi +import shutil +import tempfile + +from dhi.tasks.resonant import MergeResonantLimits + +from .StatInferenceTask import StatInferenceTask +from .ResonantLimitsTask import ( + ResonantLimitsTask, + cards_by_mass, + combine_cards_per_mass, + publish_dir, +) + + +class CombinedResonantLimitsTask(StatInferenceTask): + """The limit of a combination configuration (`members:`): each member's headline + cards, combined per mass and labelled by member. Same-named nuisances correlate. + """ + + # Not a workflow itself; see ResonantLimitsTask. + workflow = luigi.Parameter(default=law.parameter.NO_STR) + + def requires(self): + if not self.is_combination(): + raise RuntimeError( + f"{self.datacard_config_path()} declares no 'members', so there is " + "nothing to combine; run ResonantLimitsTask for a single configuration." + ) + return { + name: self.member_req(name, ResonantLimitsTask) for name in self.members() + } + + def store_parts(self): + return (*self.version_parts(), self.__class__.__name__) + + def output(self): + return { + "limits": self.local_target("limits.npz"), + "datacards": law.LocalDirectoryTarget(self.datacards_dir("combined")), + } + + @staticmethod + def check_members(member_cards): + """The members must share their mass points and no bin.""" + masses = { + name: sorted(cards_by_mass(cards), key=int) + for name, cards in member_cards.items() + } + reference = next(iter(masses.values())) + if any(m != reference for m in masses.values()): + raise RuntimeError(f"the members disagree on the mass points: {masses}") + + owner = {} + for name, cards in member_cards.items(): + with open(cards[0]) as f: + # the first "bin" line lists each bin once (the second, per process) + bins = next(line.split()[1:] for line in f if line.startswith("bin ")) + for b in bins: + if owner.setdefault(b, name) != name: + raise RuntimeError( + f"bin '{b}' is in both '{owner[b]}' and '{name}'; the members' " + "channels and categories must not overlap" + ) + + def run(self): + member_cards = {} + for name, req in self.requires().items(): + cards = req.headline_cards() + if not cards: + raise RuntimeError(f"member '{name}' has no datacards to combine") + member_cards[name] = cards + self.check_members(member_cards) + + with tempfile.TemporaryDirectory() as staging: + combine_cards_per_mass(list(member_cards.items()), staging, self.cmssw_env) + combined = publish_dir(staging, self.output()["datacards"]) + + merged = yield MergeResonantLimits( + version=self.version, datacards=tuple(sorted(combined)) + ) + self.output()["limits"].parent.touch() + shutil.copy2(merged.path, self.output()["limits"].path) diff --git a/law/CreateDatacardsTask.py b/law/CreateDatacardsTask.py index 1b09c9a..a63fbdb 100644 --- a/law/CreateDatacardsTask.py +++ b/law/CreateDatacardsTask.py @@ -81,7 +81,9 @@ def output(self): # fs_default. Note that combine cannot read these directly: # ResonantLimitsTask mirrors them back to datacards_dir() before handing them to # dhi -- see ResonantLimitsTask.stage_datacards. - return self.output_dir_target(self.version, "Datacards", self.datacard_era) + return self.output_dir_target( + *self.version_parts(), "Datacards", self.datacard_era + ) def run(self): statInf_entry = self.global_params["StatInference"] @@ -143,13 +145,35 @@ def run(self): base_dir_local.abspath, config, local_output.abspath ) - def plot_variable(self, variable): + @staticmethod + def _surviving_axis(base_dir): + """Which axis of the 2D input the rebinned shapes are binned along: 0 = x, 1 = y. + + Read from binning.json: a slice with a "y_range" (a window on y) is binned along x, + anything else along y. Without a record, y. + """ + import json + + path = os.path.join(base_dir, "binning.json") + try: + with open(path) as f: + record = json.load(f) + except (OSError, ValueError): + return 1 + node = record.get("binning", {}) + # era -> mass -> channel -> category -> {"slices": [...]} + while isinstance(node, dict) and "slices" not in node: + if not node: + return 1 + node = next(iter(node.values())) + slices = node.get("slices") if isinstance(node, dict) else None + return 0 if slices and "y_range" in slices[0] else 1 + + def plot_variable(self, variable, axis=1): """The variable whose axis the shapes are binned along, for histograms.yaml. - For input a 2D->1D rebinning produced, that is the y variable of the 2D entry -- - histograms.yaml records both as ``var_list: [x, y]``, so the axis metadata the - plotter needs is already described and does not have to be restated here. For - input that was always 1D, the variable is its own answer. + For 2D->1D rebinned input, the entry of the 2D variable's ``var_list: [x, y]`` + picked by `axis` (from _surviving_axis()); for 1D input, the variable itself. """ try: import FLAF.Common.Setup as Setup @@ -157,7 +181,7 @@ def plot_variable(self, variable): hists = Setup.Setup(self.ana_path(), self.period, self.version).hists var_list = hists[variable].get("var_list") if var_list and len(var_list) > 1: - return var_list[1] + return var_list[axis] except Exception as e: print(f"Warning: no var_list for {variable} ({e}); plotting it as itself") return variable @@ -332,7 +356,7 @@ def plot_rebinned_shapes(self, base_dir, config, output_dir): # binning -- so no --rebin, which would coarsen it back to the histograms.yaml # grid and undo the whole point of the rebinning. "--var", - self.plot_variable(variable), + self.plot_variable(variable, self._surviving_axis(base_dir)), # The plotted shapes are the sum over every sub-era, so the label has to # name the combination: HistPlotter reads config/plot/.yaml for the # luminosity, and a sub-era's file states that sub-era's luminosity alone. diff --git a/law/MergedHists.py b/law/MergedHists.py index 41a5709..39f523a 100644 --- a/law/MergedHists.py +++ b/law/MergedHists.py @@ -19,6 +19,12 @@ class MergedHists(StatInferenceTask, law.ExternalTask): variable = luigi.Parameter(description="Hists_merged variable (directory) name") + # The input file is the same whichever configuration reads it. + exclude_params_req = StatInferenceTask.exclude_params_req | { + "datacard_config", + "datacard_tag", + } + def output(self): return self.remote_target( self.input_hists_version, diff --git a/law/PlotPullsAndImpactsTask.py b/law/PlotPullsAndImpactsTask.py index 771fac6..15ff254 100644 --- a/law/PlotPullsAndImpactsTask.py +++ b/law/PlotPullsAndImpactsTask.py @@ -42,10 +42,16 @@ class PlotPullsAndImpactsTask(DhiPlotMixin, StatInferenceTask): def requires(self): # ResonantLimitsTask is what writes the combined cards this reads. + if self.is_combination(): + # card_for()/card_family() only know a single configuration's layout + raise RuntimeError( + f"{self.datacard_config_path()} is a combination; impacts are not yet " + "supported for it. Run them on a member with --datacard-config." + ) return [ResonantLimitsTask.req(self)] def output(self): - return self.output_dir_target(self.version, "ImpactPlots") + return self.output_dir_target(*self.version_parts(), "ImpactPlots") def entry_eras(self, entry): """The eras an entry is drawn for. diff --git a/law/PlotResonantLimitsTask.py b/law/PlotResonantLimitsTask.py index fd41998..67fdc86 100644 --- a/law/PlotResonantLimitsTask.py +++ b/law/PlotResonantLimitsTask.py @@ -18,6 +18,7 @@ from .DhiPlotMixin import DhiPlotMixin from .StatInferenceTask import StatInferenceTask from .ResonantLimitsTask import ResonantLimitsTask +from .CombinedResonantLimitsTask import CombinedResonantLimitsTask class PlotResonantLimitsTask(DhiPlotMixin, StatInferenceTask): @@ -46,12 +47,17 @@ class PlotResonantLimitsTask(DhiPlotMixin, StatInferenceTask): def requires(self): # ResonantLimitsTask is what mirrors the cards to datacards_dir(), which is what # the globs below resolve against -- the plots cannot be built before it has run. - return [ResonantLimitsTask.req(self)] + if not self.is_combination(): + return [ResonantLimitsTask.req(self)] + # a combination also draws each member's own limit_plots + return [CombinedResonantLimitsTask.req(self)] + [ + self.member_req(name, PlotResonantLimitsTask) for name in self.members() + ] def output(self): # fs_default, like CreateDatacardsTask. Holds one # / sub-directory of plots plus plots.json naming them. - return self.output_dir_target(self.version, "LimitPlots") + return self.output_dir_target(*self.version_parts(), "LimitPlots") def _resolve_config_path(self, path): return path if os.path.isabs(path) else os.path.join(self.ana_path(), path) @@ -160,18 +166,21 @@ def datacard_sequences(self, entry, era): same cards the leading curve is drawn from. """ name = entry.get("name") or "limits" - base_dir = self.datacards_dir(era) sequences, names = [], [] for card in entry["datacards"]: - pattern = os.path.join( - base_dir, Template(card["glob"]).safe_substitute(ERA=era) - ) - if not glob.glob(pattern): - raise RuntimeError( - f"limit_plots entry '{name}': no datacard matches '{pattern}'. " - f"Check the glob against the cards under {base_dir}." - ) + # several globs: dhi combines the cards they match per mass + globs = card["glob"] if isinstance(card["glob"], list) else [card["glob"]] + patterns = [] + for g in globs: + base_dir, rel = self.glob_base(g, era, name) + pattern = os.path.join(base_dir, Template(rel).safe_substitute(ERA=era)) + if not glob.glob(pattern): + raise RuntimeError( + f"limit_plots entry '{name}': no datacard matches '{pattern}'. " + f"Check the glob against the cards under {base_dir}." + ) + patterns.append(pattern) label = card["name"] if "{" in label or "}" in label: # law brace-expands datacard_names, so "fb^{-1}" would arrive as "fb^-1". @@ -179,10 +188,23 @@ def datacard_sequences(self, entry, era): f"limit_plots entry '{name}': datacard name '{label}' contains braces, " "which law brace-expands. Put the information in the campaign label instead." ) - sequences.append((pattern,)) + sequences.append(tuple(patterns)) names.append(label) return sequences, names + def glob_base(self, card_glob, era, entry_name): + """(directory, relative glob) a limit_plots glob resolves against. In a + combination the glob names its member, "SL:${ERA}/e/*.txt".""" + if not self.is_combination(): + return self.datacards_dir(era), card_glob + member, sep, rel = card_glob.partition(":") + if not sep or member not in self.members(): + raise RuntimeError( + f"limit_plots entry '{entry_name}': glob '{card_glob}' must name a member " + f"as ':'; members are {sorted(self.members())}" + ) + return self.member_req(member, ResonantLimitsTask).datacards_dir(era), rel + def limits_npz(self, sequence): """The MergeResonantLimits .npz for a datacard sequence, produced if absent. diff --git a/law/PreprocessShapesTask.py b/law/PreprocessShapesTask.py index 01235a2..56ddbfa 100644 --- a/law/PreprocessShapesTask.py +++ b/law/PreprocessShapesTask.py @@ -63,7 +63,7 @@ def create_branch_map(self): def output(self): return self.output_dir_target( - self.version, "Hists_preprocessed", self.datacard_era + *self.version_parts(), "Hists_preprocessed", self.datacard_era ) def run(self): diff --git a/law/ResonantLimitsTask.py b/law/ResonantLimitsTask.py index 646601f..7dfb578 100644 --- a/law/ResonantLimitsTask.py +++ b/law/ResonantLimitsTask.py @@ -17,6 +17,45 @@ COMBINED_KEY = "__combined__" +def cards_by_mass(cards): + """{mass: card} from datacard paths named "..._.txt".""" + out = {} + for c in cards: + m = re.search(r"_(\d+)\.txt$", c) + if m: + out[m.group(1)] = c + return out + + +def combine_cards_per_mass(labelled_cards, staging, env): + """Write staging/combined_.txt per mass from [(label, [card, ...])]; the + unique label prefixes each card's bins.""" + by_label = [(label, cards_by_mass(cards)) for label, cards in labelled_cards] + masses = sorted({m for _, cards in by_label for m in cards}, key=int) + written = [] + for mass in masses: + args = [f"{label}={cards[mass]}" for label, cards in by_label if mass in cards] + out_file = os.path.join(staging, f"combined_{mass}.txt") + with open(out_file, "w") as f: + subprocess.run(["combineCards.py"] + args, env=env, stdout=f, check=True) + written.append(out_file) + return written + + +def publish_dir(staging, target): + """Replace directory `target` with the files in `staging`, returning their paths. + Only called once everything is staged, so a failure never looks complete to law.""" + if target.exists(): + target.remove() + target.touch() + published = [] + for name in sorted(os.listdir(staging)): + dest = os.path.join(target.path, name) + shutil.copy2(os.path.join(staging, name), dest) + published.append(dest) + return published + + class ResonantLimitsTask(StatInferenceTask): # Not a workflow itself. The parameter exists so that --workflow reaches the tasks # below that are -- law's req() only forwards parameters both tasks declare -- and @@ -24,7 +63,7 @@ class ResonantLimitsTask(StatInferenceTask): workflow = luigi.Parameter(default=law.parameter.NO_STR) def store_parts(self): - return (self.version, self.__class__.__name__, "combined") + return (*self.version_parts(), self.__class__.__name__, "combined") def get_eras(self): """Every era the configuration lists: each gets its own cards and its own limit.""" @@ -53,6 +92,18 @@ def output(self): "datacards": law.LocalDirectoryTarget(self.datacards_dir("combined")), } + def headline_cards(self): + """The per-mass cards the limit is set on: the cross-era combined ones, or the + single top-level era's.""" + top_level = self.get_top_level_eras() + if len(top_level) > 1: + return sorted( + glob.glob(os.path.join(self.datacards_dir("combined"), "*.txt")) + ) + return sorted( + glob.glob(os.path.join(self.datacards_dir(top_level[0]), "*.txt")) + ) + def stage_datacards(self, era, remote_target): """Mirror an era's datacards from fs_default to a stable local path, returned. @@ -85,13 +136,6 @@ def run(self): output_dir = self.stage_datacards(e, create_dc_br0.output()) era_cards[e] = glob.glob(os.path.join(output_dir, "*.txt")) - masses = set() - for e, cards in era_cards.items(): - for c in cards: - m = re.search(r"_(\d+)\.txt$", c) - if m: - masses.add(m.group(1)) - # Built into a staging directory and published only once every card is there. # Creating the outputs first and filling them as we go left a failed # combineCards halfway through looking like a complete task -- law's default @@ -102,29 +146,12 @@ def run(self): # The cross-era combination is the one place a group era and its members # cannot both appear: the group already *is* their combination, so a card # built from both would count those events twice. - for mass in sorted(masses): - combine_args = [] - for e in self.get_top_level_eras(): - for c in era_cards.get(e, []): - if c.endswith(f"_{mass}.txt"): - combine_args.append(f"{e}={c}") - break - - if combine_args: - cmd = ["combineCards.py"] + combine_args - out_file = os.path.join(staging, f"combined_{mass}.txt") - with open(out_file, "w") as f: - subprocess.run(cmd, env=self.cmssw_env, stdout=f, check=True) - - out_dc_dir = self.output()["datacards"] - if out_dc_dir.exists(): - out_dc_dir.remove() - out_dc_dir.touch() - combined_cards = [] - for name in sorted(os.listdir(staging)): - dest = os.path.join(out_dc_dir.path, name) - shutil.copy2(os.path.join(staging, name), dest) - combined_cards.append(dest) + combine_cards_per_mass( + [(e, era_cards.get(e, [])) for e in self.get_top_level_eras()], + staging, + self.cmssw_env, + ) + combined_cards = publish_dir(staging, self.output()["datacards"]) # One merge per era, never one merge over every era's cards together. dhi groups # the datacards it is given by the mass in their file name alone and runs diff --git a/law/StatInferenceTask.py b/law/StatInferenceTask.py index 0dab693..b7a1cff 100644 --- a/law/StatInferenceTask.py +++ b/law/StatInferenceTask.py @@ -2,6 +2,7 @@ import law import luigi import os +import re import shutil import yaml @@ -29,10 +30,26 @@ class StatInferenceTask(Task): description="version of the Hists_merged tree to read; defaults to --version", ) + # set per member by a combination configuration + datacard_config = luigi.Parameter( + default="", + description="datacard configuration to use instead of StatInference.config", + ) + # empty keeps a single configuration's paths unchanged + datacard_tag = luigi.Parameter( + default="", + description="sub-directory of --version holding this configuration's products", + ) + @property def input_hists_version(self): return self.hists_version or self.version + def version_parts(self): + """--version, then the tag ("combined" for a combination's own products).""" + tag = self.datacard_tag or ("combined" if self.is_combination() else "") + return (self.version, tag) if tag else (self.version,) + def output_dir_target(self, *path): """remote_dir_target() that also works when fs_default is a local directory. @@ -50,9 +67,8 @@ def output_dir_target(self, *path): return self.remote_dir_target(*path) def datacard_config_path(self): - return os.path.join( - self.ana_path(), self.global_params["StatInference"]["config"] - ) + path = self.datacard_config or self.global_params["StatInference"]["config"] + return path if os.path.isabs(path) else os.path.join(self.ana_path(), path) def get_config_data(self): # Cached per instance: requires()/workflow_requires() are re-entered many times @@ -62,6 +78,49 @@ def get_config_data(self): self._config_data = yaml.safe_load(f) return self._config_data + def is_combination(self): + """Whether the configuration combines others (`members:`) rather than building + datacards itself.""" + return bool(self.get_config_data().get("members")) + + def members(self): + """{name: {config, hists_version, user_custom, customisations}}, validated.""" + members = self.get_config_data().get("members") or {} + known = {"config", "hists_version", "user_custom", "customisations"} + for name, spec in members.items(): + if not isinstance(spec, dict) or "config" not in spec: + raise RuntimeError( + f"{self.datacard_config_path()}: member '{name}' must be a mapping " + "with at least 'config'" + ) + unknown = sorted(set(spec) - known) + if unknown: + raise RuntimeError( + f"{self.datacard_config_path()}: member '{name}' has unknown keys " + f"{unknown}; known are {sorted(known)}" + ) + if not re.fullmatch(r"[A-Za-z0-9_]+", str(name)) or name == "combined": + raise RuntimeError( + f"{self.datacard_config_path()}: member name '{name}' must be " + "alphanumeric and not 'combined' -- it names a directory and a " + "combineCards label" + ) + return members + + def member_req(self, name, cls, **kwargs): + """`cls` for member `name`: its own configuration, tag and histograms (FLAF + builds a Setup per user_custom). Everything else is inherited.""" + spec = self.members()[name] + params = dict( + datacard_config=spec["config"], + datacard_tag=name, + hists_version=spec.get("hists_version", self.hists_version), + user_custom=spec.get("user_custom", self.user_custom), + customisations=spec.get("customisations", self.customisations), + ) + params.update(kwargs) + return cls.req(self, **params) + @property def datacard_era(self): """The era this task's products belong to: its meta-era if it has one, else @@ -85,7 +144,7 @@ def store_parts(self): Identical to the old value for every task without a meta_era. """ - return (self.version, self.__class__.__name__, self.datacard_era) + return (*self.version_parts(), self.__class__.__name__, self.datacard_era) def datacards_dir(self, era): """Local directory holding an era's datacards. @@ -94,7 +153,9 @@ def datacards_dir(self, era): combine sees them; ResonantLimitsTask mirrors them here and everything downstream (dhi's --multi-datacards globbing, the overlay plots) resolves against this path. """ - return os.path.join(self.ana_data_path(), self.version, "Datacards", era) + return os.path.join( + self.ana_data_path(), *self.version_parts(), "Datacards", era + ) def preprocess_config(self): """The datacard configuration's `preprocess:` block, or None. diff --git a/law/tasks.py b/law/tasks.py index 8e07ebf..185e606 100644 --- a/law/tasks.py +++ b/law/tasks.py @@ -11,6 +11,7 @@ from .PreprocessShapesTask import PreprocessShapesTask from .CreateDatacardsTask import CreateDatacardsTask from .ResonantLimitsTask import ResonantLimitsTask +from .CombinedResonantLimitsTask import CombinedResonantLimitsTask from .PlotResonantLimitsTask import PlotResonantLimitsTask from .PlotPullsAndImpactsTask import PlotPullsAndImpactsTask from .ResonantLimitsAndHistPlotTask import ResonantLimitsAndHistPlotTask @@ -21,6 +22,7 @@ "PreprocessShapesTask", "CreateDatacardsTask", "ResonantLimitsTask", + "CombinedResonantLimitsTask", "PlotResonantLimitsTask", "PlotPullsAndImpactsTask", "ResonantLimitsAndHistPlotTask",