From 2d45d2d0d1bdf82371f80702c8ebdddabd02a17d Mon Sep 17 00:00:00 2001 From: "Marcel R." Date: Tue, 21 Jan 2025 18:41:14 +0100 Subject: [PATCH] Fix data shape in qcd hook. --- hbt/config/hist_hooks.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/hbt/config/hist_hooks.py b/hbt/config/hist_hooks.py index ccc51791..55ce07c4 100644 --- a/hbt/config/hist_hooks.py +++ b/hbt/config/hist_hooks.py @@ -110,6 +110,20 @@ def qcd_estimation(task, hists): ss_noniso_data = hist_to_num(get_hist(data_hist, "ss_noniso"), "ss_noniso_data") ss_iso_data = hist_to_num(get_hist(data_hist, "ss_iso"), "ss_iso_data") + # data will always have a single shift whereas mc might have multiple, + # broadcast numbers in-place manually if necessary + if (n_shifts := mc_hist.axes["shift"].size) > 1: + def broadcast_data_num(num: sn.Number) -> None: + num._nominal = np.repeat(num.nominal, n_shifts, axis=0) + for name, (unc_up, unc_down) in num._uncertainties.items(): + num._uncertainties[name] = ( + np.repeat(unc_up, n_shifts, axis=0), + np.repeat(unc_down, n_shifts, axis=0), + ) + broadcast_data_num(os_noniso_data) + broadcast_data_num(ss_noniso_data) + broadcast_data_num(ss_iso_data) + # estimate qcd shapes in the three sideband regions # shapes: (SHIFT, VAR) os_noniso_qcd = os_noniso_data - os_noniso_mc