From 06f3f1de0aee05fa2e774b330d99bf4bd5c9d5f7 Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Thu, 18 Jun 2026 09:35:34 -0400 Subject: [PATCH 01/10] another bug fix flavor --- scripts/histmakers/mz_5TeV.py | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index 53572dc3e..9d6e374d6 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -502,5 +502,6 @@ def build_graph(df, dataset): resultdict = narf.build_and_run(datasets, build_graph) +args.flavor = "mumu" fout = f"{os.path.basename(__file__).replace('py', 'hdf5')}" write_analysis_output(resultdict, fout, args) From cb84cce48b753f48e929fb73997b63e9ec3de76c Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Thu, 25 Jun 2026 09:08:39 -0400 Subject: [PATCH 02/10] trying again --- scripts/histmakers/mz_5TeV.py | 41 +-- .../production/include/lowpu_muonscarekit.hpp | 239 ++++++++++++++++++ 2 files changed, 263 insertions(+), 17 deletions(-) create mode 100644 wremnants/production/include/lowpu_muonscarekit.hpp diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index 9d6e374d6..d2ae0a332 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -188,12 +188,12 @@ def build_graph(df, dataset): if dataset.is_data: df = df.Define( "Muon_pt_corr", - "wrem::applyMuonScarekitData(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", + "wrem::applyMuonScarekitStep1Data(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", ) else: df = df.Define( "Muon_pt_corr", - "wrem::applyMuonScarekitMC_scaleOnly(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", + "wrem::applyMuonScarekitStep1MC(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", ) elif args.corrStep == "123": if dataset.is_data: @@ -204,7 +204,7 @@ def build_graph(df, dataset): else: df = df.Define( "Muon_pt_corr", - "wrem::applyMuonScarekitMC_noKFactor(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers)", + "wrem::applyMuonScarekitMC_scaleOnly(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", ) else: # "1234" if dataset.is_data: @@ -314,8 +314,9 @@ def build_graph(df, dataset): df = df.DefinePerSample("central_pdf_weight", "1.0") df = df.Alias("nominal_weight_uncorr", "exp_weight") df = df.DefinePerSample("theory_weight_truncate", "10.0") + applied_theory_corrs = [] for theory_corr_name in theory_corrs: - if theory_corr_name not in corr_helpers[dataset.name]: + if theory_corr_name not in corr_helpers.get(dataset.name, {}): continue df = theory_corrections.define_theory_corr_weight_column( df, theory_corr_name @@ -331,9 +332,14 @@ def build_graph(df, dataset): f"{theory_corr_name}_corr_weight", ], ) + applied_theory_corrs.append(theory_corr_name) - theory_corr_name = theory_corrs[0] - df = df.Define("nominal_weight", f"{theory_corr_name}Weight_tensor[0]") + if applied_theory_corrs: + df = df.Define( + "nominal_weight", f"{applied_theory_corrs[0]}Weight_tensor[0]" + ) + else: + df = df.Alias("nominal_weight", "exp_weight") # ---- Fill histograms ---- hist_nLepton = df.HistoBoost( @@ -456,17 +462,18 @@ def build_graph(df, dataset): ) results.append(hist_mutraileta_prefire) - systematics.add_theory_corr_hists( - results, - df, - [axis_ptll, axis_absYll, axis_cosThetaStarll], - ["ptll", "absYll", "cosThetaStarll"], - corr_helpers[dataset.name], - theory_corrs, - modify_central_weight=True, - isW=False, - base_name="ptll", - ) + if applied_theory_corrs: + systematics.add_theory_corr_hists( + results, + df, + [axis_ptll, axis_absYll, axis_cosThetaStarll], + ["ptll", "absYll", "cosThetaStarll"], + corr_helpers[dataset.name], + theory_corrs, + modify_central_weight=True, + isW=False, + base_name="ptll", + ) results += [ hist_mll, diff --git a/wremnants/production/include/lowpu_muonscarekit.hpp b/wremnants/production/include/lowpu_muonscarekit.hpp new file mode 100644 index 000000000..2f4d8ac31 --- /dev/null +++ b/wremnants/production/include/lowpu_muonscarekit.hpp @@ -0,0 +1,239 @@ +#ifndef WREMNANTS_LOWPU_MUONSCAREKIT_H +#define WREMNANTS_LOWPU_MUONSCAREKIT_H + +#include "defines.hpp" +#include +#include +#include +#include +#include +#include + +namespace wrem { + +struct MuonScarekitCB { + static const double pi; + static const double sqrtPiOver2; + static const double sqrt2; + + double m, s, a, n; + double B, C, D, N, NA, Ns, NC, F, G, k; + double cdfMa, cdfPa; + + MuonScarekitCB(double mean, double sigma, double alpha, double nn) + : m(mean), s(sigma), a(alpha), n(nn) { + init(); + } + + void init() { + double fa = fabs(a); + double ex = exp(-fa * fa / 2); + double A = pow(n / fa, n) * ex; + double C1 = n / fa / (n - 1) * ex; + double D1 = 2 * sqrtPiOver2 * boost::math::erf(fa / sqrt2); + B = n / fa - fa; + C = (D1 + 2 * C1) / C1; + D = (D1 + 2 * C1) / 2; + N = 1.0 / s / (D1 + 2 * C1); + k = 1.0 / (n - 1); + NA = N * A; + Ns = N * s; + NC = Ns * C1; + F = 1 - fa * fa / n; + G = s * n / fa; + cdfMa = cdf(m - a * s); + cdfPa = cdf(m + a * s); + } + + double cdf(double x) const { + double d = (x - m) / s; + if (d < -a) + return NC / pow(F - s * d / G, n - 1); + if (d > a) + return NC * (C - pow(F + s * d / G, 1 - n)); + return Ns * (D - sqrtPiOver2 * boost::math::erf(-d / sqrt2)); + } + + double invcdf(double u) const { + if (u < cdfMa) + return m + G * (F - pow(NC / u, k)); + if (u > cdfPa) + return m - G * (F - pow(C - u / NC, -k)); + return m - sqrt2 * s * boost::math::erf_inv((D - u / Ns) / sqrtPiOver2); + } +}; + +const double MuonScarekitCB::pi = 3.14159265358979; +const double MuonScarekitCB::sqrtPiOver2 = sqrt(MuonScarekitCB::pi / 2.0); +const double MuonScarekitCB::sqrt2 = sqrt(2.0); + +namespace muonscarekit_step1_impl { +TFile *tf_scale = + TFile::Open("wremnants-data/data/lowPU/muonscarekit/step1_C.root", "READ"); +TH2D *h_M_DATA = (TH2D *)tf_scale->Get("M_DATA"); +TH2D *h_A_DATA = (TH2D *)tf_scale->Get("A_DATA"); +TH2D *h_M_SIG = (TH2D *)tf_scale->Get("M_SIG"); +TH2D *h_A_SIG = (TH2D *)tf_scale->Get("A_SIG"); +} // namespace muonscarekit_step1_impl + +namespace muonscarekit_impl { +TFile *tf_scale = TFile::Open( + "wremnants-data/data/lowPU/muonscarekit/step3_correction.root", "READ"); +TH2D *h_M_DATA = (TH2D *)tf_scale->Get("M_DATA"); +TH2D *h_A_DATA = (TH2D *)tf_scale->Get("A_DATA"); +TH2D *h_M_SIG = (TH2D *)tf_scale->Get("M_SIG"); +TH2D *h_A_SIG = (TH2D *)tf_scale->Get("A_SIG"); + +TFile *tf_cb = TFile::Open( + "wremnants-data/data/lowPU/muonscarekit/step2_fitresults.root", "READ"); +TH3D *h_cb = (TH3D *)tf_cb->Get("h_results_cb"); +TH3D *h_poly = (TH3D *)tf_cb->Get("h_results_poly"); + +TFile *tf_k = + TFile::Open("wremnants-data/data/lowPU/muonscarekit/step4_k.root", "READ"); +TH2D *h_k_data = (TH2D *)tf_k->Get("k_hist_DATA"); +TH2D *h_k_sig = (TH2D *)tf_k->Get("k_hist_SIG"); +} // namespace muonscarekit_impl + +Vec_f applyMuonScarekitData(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { + using namespace muonscarekit_impl; + unsigned int size = pt.size(); + Vec_f res(size); + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_DATA->GetBinContent(h_M_DATA->FindBin(eta[i], phi[i])); + double A = h_A_DATA->GetBinContent(h_A_DATA->FindBin(eta[i], phi[i])); + res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); + } + return res; +} + +Vec_f applyMuonScarekitMC_scaleOnly(Vec_f pt, Vec_f eta, Vec_f phi, + Vec_i charge) { + using namespace muonscarekit_impl; + unsigned int size = pt.size(); + Vec_f res(size); + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); + double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); + res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); + } + return res; +} + +Vec_f applyMuonScarekitStep1Data(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { + using namespace muonscarekit_step1_impl; + unsigned int size = pt.size(); + Vec_f res(size); + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_DATA->GetBinContent(h_M_DATA->FindBin(eta[i], phi[i])); + double A = h_A_DATA->GetBinContent(h_A_DATA->FindBin(eta[i], phi[i])); + res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); + } + return res; +} + +Vec_f applyMuonScarekitStep1MC(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { + using namespace muonscarekit_step1_impl; + unsigned int size = pt.size(); + Vec_f res(size); + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); + double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); + res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); + } + return res; +} + +Vec_f applyMuonScarekitMC_noKFactor(Vec_f pt, Vec_f eta, Vec_f phi, + Vec_i charge, Vec_i nTrackerLayers) { + using namespace muonscarekit_impl; + unsigned int size = pt.size(); + Vec_f res(size); + + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); + double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); + double pt_scale = 1.0 / (M / pt[i] + charge[i] * A); + + Int_t etabin = h_cb->GetXaxis()->FindBin(fabs((double)eta[i])); + Int_t nlbin = h_cb->GetYaxis()->FindBin((double)nTrackerLayers[i]); + + double mean_cb = h_cb->GetBinContent(etabin, nlbin, 1); + double sig_cb = h_cb->GetBinContent(etabin, nlbin, 2); + double n_cb = h_cb->GetBinContent(etabin, nlbin, 3); + double alpha_cb = h_cb->GetBinContent(etabin, nlbin, 4); + + double a_poly = h_poly->GetBinContent(etabin, nlbin, 1); + double b_poly = h_poly->GetBinContent(etabin, nlbin, 2); + double c_poly = h_poly->GetBinContent(etabin, nlbin, 3); + double sigma_poly = + a_poly + b_poly * pt_scale + c_poly * pt_scale * pt_scale; + if (sigma_poly < 0.0) + sigma_poly = 0.0; + + if (sigma_poly == 0.0 || n_cb <= 1.0 + 1e-6 || sig_cb <= 0.0 || + alpha_cb <= 0.0) { + res[i] = static_cast(pt_scale); + continue; + } + + MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); + double rndm_cb = cb.invcdf(gRandom->Rndm()); + + // step123: scale + smearing model, but no k-factor rescaling (k_mc = 1) + res[i] = static_cast(pt_scale * (1.0 + sigma_poly * rndm_cb)); + } + return res; +} + +Vec_f applyMuonScarekitMC(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge, + Vec_i nTrackerLayers, unsigned int run, + unsigned int lumi) { + using namespace muonscarekit_impl; + unsigned int size = pt.size(); + Vec_f res(size); + + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); + double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); + double pt_scale = 1.0 / (M / pt[i] + charge[i] * A); + + Int_t etabin = h_cb->GetXaxis()->FindBin(fabs((double)eta[i])); + Int_t nlbin = h_cb->GetYaxis()->FindBin((double)nTrackerLayers[i]); + + double mean_cb = h_cb->GetBinContent(etabin, nlbin, 1); + double sig_cb = h_cb->GetBinContent(etabin, nlbin, 2); + double n_cb = h_cb->GetBinContent(etabin, nlbin, 3); + double alpha_cb = h_cb->GetBinContent(etabin, nlbin, 4); + + double a_poly = h_poly->GetBinContent(etabin, nlbin, 1); + double b_poly = h_poly->GetBinContent(etabin, nlbin, 2); + double c_poly = h_poly->GetBinContent(etabin, nlbin, 3); + double sigma_poly = + a_poly + b_poly * pt_scale + c_poly * pt_scale * pt_scale; + if (sigma_poly < 0.0) + sigma_poly = 0.0; + + Int_t absetabin = h_k_data->GetXaxis()->FindBin(fabs((double)eta[i])); + double k_data_v = h_k_data->GetBinContent(absetabin, 3); + double k_sig_v = h_k_sig->GetBinContent(absetabin, 3); + double k_mc = (k_sig_v < k_data_v) + ? sqrt(k_data_v * k_data_v - k_sig_v * k_sig_v) + : 0.0; + + if (k_mc == 0.0 || sigma_poly == 0.0 || n_cb <= 1.0 + 1e-6 || + sig_cb <= 0.0 || alpha_cb <= 0.0) { + res[i] = static_cast(pt_scale); + continue; + } + + MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); + double rndm_cb = cb.invcdf(gRandom->Rndm()); + + res[i] = static_cast(pt_scale * (1.0 + k_mc * sigma_poly * rndm_cb)); + } + return res; +} + +} // namespace wrem +#endif From e29fc862e78bd01521284b30c4acd220fd03754d Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Thu, 25 Jun 2026 09:14:16 -0400 Subject: [PATCH 03/10] removing corrstep arg --- scripts/histmakers/mz_5TeV.py | 53 +++-------- .../production/include/lowpu_muonscarekit.hpp | 88 ------------------- 2 files changed, 10 insertions(+), 131 deletions(-) diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index d2ae0a332..632b63a9c 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -12,14 +12,6 @@ choices=["none", "rochester", "scarekit"], help="Muon momentum correction to apply", ) -parser.add_argument( - "--corrStep", - default="1234", - choices=["0", "1", "123", "1234"], - help="Scarekit calibration stage (only with --muonCorr scarekit): " - "0 = no correction, 1 = scale only, 123 = scale+smearing, 1234 = full", -) - args = parser.parse_args() logger = logging.setup_logger(__file__, args.verbose, args.noColorLogger) @@ -182,41 +174,16 @@ def build_graph(df, dataset): "wrem::applyRochesterMC(Muon_pt, Muon_eta, Muon_phi, ROOT::VecOps::RVec(Muon_charge.begin(), Muon_charge.end()), Muon_genPartIdx, GenPart_pt, Muon_nTrackerLayers)", ) elif args.muonCorr == "scarekit": - if args.corrStep == "0": - df = df.Alias("Muon_pt_corr", "Muon_pt") - elif args.corrStep == "1": - if dataset.is_data: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitStep1Data(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", - ) - else: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitStep1MC(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", - ) - elif args.corrStep == "123": - if dataset.is_data: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitData(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", - ) - else: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitMC_scaleOnly(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", - ) - else: # "1234" - if dataset.is_data: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitData(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", - ) - else: - df = df.Define( - "Muon_pt_corr", - "wrem::applyMuonScarekitMC(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers, run, luminosityBlock)", - ) + if dataset.is_data: + df = df.Define( + "Muon_pt_corr", + "wrem::applyMuonScarekitData(Muon_pt, Muon_eta, Muon_phi, Muon_charge)", + ) + else: + df = df.Define( + "Muon_pt_corr", + "wrem::applyMuonScarekitMC(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers, run, luminosityBlock)", + ) else: # "none" df = df.Alias("Muon_pt_corr", "Muon_pt") diff --git a/wremnants/production/include/lowpu_muonscarekit.hpp b/wremnants/production/include/lowpu_muonscarekit.hpp index 2f4d8ac31..241ca1e34 100644 --- a/wremnants/production/include/lowpu_muonscarekit.hpp +++ b/wremnants/production/include/lowpu_muonscarekit.hpp @@ -67,15 +67,6 @@ const double MuonScarekitCB::pi = 3.14159265358979; const double MuonScarekitCB::sqrtPiOver2 = sqrt(MuonScarekitCB::pi / 2.0); const double MuonScarekitCB::sqrt2 = sqrt(2.0); -namespace muonscarekit_step1_impl { -TFile *tf_scale = - TFile::Open("wremnants-data/data/lowPU/muonscarekit/step1_C.root", "READ"); -TH2D *h_M_DATA = (TH2D *)tf_scale->Get("M_DATA"); -TH2D *h_A_DATA = (TH2D *)tf_scale->Get("A_DATA"); -TH2D *h_M_SIG = (TH2D *)tf_scale->Get("M_SIG"); -TH2D *h_A_SIG = (TH2D *)tf_scale->Get("A_SIG"); -} // namespace muonscarekit_step1_impl - namespace muonscarekit_impl { TFile *tf_scale = TFile::Open( "wremnants-data/data/lowPU/muonscarekit/step3_correction.root", "READ"); @@ -107,85 +98,6 @@ Vec_f applyMuonScarekitData(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { return res; } -Vec_f applyMuonScarekitMC_scaleOnly(Vec_f pt, Vec_f eta, Vec_f phi, - Vec_i charge) { - using namespace muonscarekit_impl; - unsigned int size = pt.size(); - Vec_f res(size); - for (unsigned int i = 0; i < size; ++i) { - double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); - double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); - res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); - } - return res; -} - -Vec_f applyMuonScarekitStep1Data(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { - using namespace muonscarekit_step1_impl; - unsigned int size = pt.size(); - Vec_f res(size); - for (unsigned int i = 0; i < size; ++i) { - double M = h_M_DATA->GetBinContent(h_M_DATA->FindBin(eta[i], phi[i])); - double A = h_A_DATA->GetBinContent(h_A_DATA->FindBin(eta[i], phi[i])); - res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); - } - return res; -} - -Vec_f applyMuonScarekitStep1MC(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { - using namespace muonscarekit_step1_impl; - unsigned int size = pt.size(); - Vec_f res(size); - for (unsigned int i = 0; i < size; ++i) { - double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); - double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); - res[i] = static_cast(1.0 / (M / pt[i] + charge[i] * A)); - } - return res; -} - -Vec_f applyMuonScarekitMC_noKFactor(Vec_f pt, Vec_f eta, Vec_f phi, - Vec_i charge, Vec_i nTrackerLayers) { - using namespace muonscarekit_impl; - unsigned int size = pt.size(); - Vec_f res(size); - - for (unsigned int i = 0; i < size; ++i) { - double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); - double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); - double pt_scale = 1.0 / (M / pt[i] + charge[i] * A); - - Int_t etabin = h_cb->GetXaxis()->FindBin(fabs((double)eta[i])); - Int_t nlbin = h_cb->GetYaxis()->FindBin((double)nTrackerLayers[i]); - - double mean_cb = h_cb->GetBinContent(etabin, nlbin, 1); - double sig_cb = h_cb->GetBinContent(etabin, nlbin, 2); - double n_cb = h_cb->GetBinContent(etabin, nlbin, 3); - double alpha_cb = h_cb->GetBinContent(etabin, nlbin, 4); - - double a_poly = h_poly->GetBinContent(etabin, nlbin, 1); - double b_poly = h_poly->GetBinContent(etabin, nlbin, 2); - double c_poly = h_poly->GetBinContent(etabin, nlbin, 3); - double sigma_poly = - a_poly + b_poly * pt_scale + c_poly * pt_scale * pt_scale; - if (sigma_poly < 0.0) - sigma_poly = 0.0; - - if (sigma_poly == 0.0 || n_cb <= 1.0 + 1e-6 || sig_cb <= 0.0 || - alpha_cb <= 0.0) { - res[i] = static_cast(pt_scale); - continue; - } - - MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); - double rndm_cb = cb.invcdf(gRandom->Rndm()); - - // step123: scale + smearing model, but no k-factor rescaling (k_mc = 1) - res[i] = static_cast(pt_scale * (1.0 + sigma_poly * rndm_cb)); - } - return res; -} - Vec_f applyMuonScarekitMC(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge, Vec_i nTrackerLayers, unsigned int run, unsigned int lumi) { From cd59f5823606f2903e00e54b07d4f779c77c5427 Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Tue, 30 Jun 2026 12:37:11 -0400 Subject: [PATCH 04/10] setup rabbit like script for Z --- scripts/rabbit/new_make_tensor.py | 305 ++++++++++++++++++++++++++++++ 1 file changed, 305 insertions(+) create mode 100644 scripts/rabbit/new_make_tensor.py diff --git a/scripts/rabbit/new_make_tensor.py b/scripts/rabbit/new_make_tensor.py new file mode 100644 index 000000000..0e3f422c1 --- /dev/null +++ b/scripts/rabbit/new_make_tensor.py @@ -0,0 +1,305 @@ +import argparse +import os +import sys + +import h5py + + +def print_flush(*args, **kwargs): + print(*args, **kwargs) + sys.stdout.flush() + + +wremnants_base = os.environ.get("WREM_BASE", None) +if wremnants_base is None: + script_dir = os.path.dirname(os.path.abspath(__file__)) + wremnants_base = os.path.abspath(os.path.join(script_dir, "../..")) + if os.path.exists(os.path.join(wremnants_base, "wums")): + sys.path.insert(0, wremnants_base) + +from rabbit import tensorwriter +from wremnants.utilities.io_tools import base_io + +parser = argparse.ArgumentParser() +parser.add_argument("infile", help="Input HDF5 file with histograms") +parser.add_argument("-o", "--output", default="./", help="output directory") +parser.add_argument("--outname", default="my_tensor", help="output file name") +parser.add_argument( + "--histName", default="ptll", help="Histogram name to use (default: ptll)" +) +parser.add_argument( + "--procFilters", + nargs="*", + default=["Zmumu"], + help="Processes to include (default: Zmumu)", +) +parser.add_argument( + "--sparse", default=False, action="store_true", help="Make sparse tensor" +) +parser.add_argument( + "--systematicType", + choices=["log_normal", "normal"], + default="log_normal", + help="probability density for systematic variations", +) +args = parser.parse_args() + +# Load data from HDF5 file +infile_path = os.path.abspath(args.infile) +if not os.path.exists(infile_path): + raise FileNotFoundError(f"Input file not found: {infile_path}") + +h5file = h5py.File(infile_path, "r") +results = base_io.load_results_h5py(h5file) + +hist_name = args.histName +all_procs = [p for p in results.keys() if p != "meta_info"] +print_flush(f"All processes found in file: {all_procs}") + +if args.procFilters: + procs_to_use = [p for p in all_procs if any(filt in p for filt in args.procFilters)] +else: + procs_to_use = all_procs + +print_flush(f"Processes after filtering: {procs_to_use}") + +# All processes are treated as MC (no data loading) +mc_procs = procs_to_use +print_flush(f"MC processes found: {mc_procs}") + +# Load MC histograms +h_mc_dict = {} +for proc in mc_procs: + if "output" in results[proc] and hist_name in results[proc]["output"]: + h_proxy = results[proc]["output"][hist_name] + h_mc_dict[proc] = h_proxy.get() if hasattr(h_proxy, "get") else h_proxy + +# Identify signal processes before closing file +signal_procs = [p for p in mc_procs if "Zmumu" in p or "Ztautau" in p] + + +def _theory_corr_hist_name(results, proc, tag): + """Match histmaker output: __Corr (e.g. ptll_..._pdfas_Corr).""" + for name in results[proc]["output"]: + if name.endswith("_Corr") and tag in name: + return name + return None + + +h_theory_corr_pdfas = None +h_theory_corr_pdfvars = None +h_theory_corr_pdfas_dict = {} +h_theory_corr_pdfvars_dict = {} + +if signal_procs: + first_sig_proc = signal_procs[0] + if first_sig_proc in results and "output" in results[first_sig_proc]: + theory_corr_hist_name_pdfas = _theory_corr_hist_name( + results, first_sig_proc, "pdfas" + ) + theory_corr_hist_name_pdfvars = _theory_corr_hist_name( + results, first_sig_proc, "pdfvars" + ) + if not theory_corr_hist_name_pdfas: + print_flush( + f"Warning: no pdfas theory hist in {first_sig_proc}; " + f"keys: {list(results[first_sig_proc]['output'].keys())}" + ) + if not theory_corr_hist_name_pdfvars: + print_flush(f"Warning: no pdfvars theory hist in {first_sig_proc}") + # Load pdfas correction (for alpha_s variations) + if ( + theory_corr_hist_name_pdfas + and theory_corr_hist_name_pdfas in results[first_sig_proc]["output"] + ): + h_proxy = results[first_sig_proc]["output"][theory_corr_hist_name_pdfas] + h_theory_corr_pdfas = h_proxy.get() if hasattr(h_proxy, "get") else h_proxy + print_flush( + f"Found theory correction histogram for alpha_s variations: {theory_corr_hist_name_pdfas}" + ) + print_flush( + f"Histogram axes: {[ax.name for ax in h_theory_corr_pdfas.axes]}" + ) + + # Load for all signal processes + for proc in signal_procs: + if proc in results and "output" in results[proc]: + if theory_corr_hist_name_pdfas in results[proc]["output"]: + h_proxy = results[proc]["output"][theory_corr_hist_name_pdfas] + h_theory_corr_pdfas_dict[proc] = ( + h_proxy.get() if hasattr(h_proxy, "get") else h_proxy + ) + + # Load pdfvars correction (for PDF variations) + if ( + theory_corr_hist_name_pdfvars + and theory_corr_hist_name_pdfvars in results[first_sig_proc]["output"] + ): + h_proxy = results[first_sig_proc]["output"][theory_corr_hist_name_pdfvars] + h_theory_corr_pdfvars = ( + h_proxy.get() if hasattr(h_proxy, "get") else h_proxy + ) + print_flush( + f"Found theory correction histogram for PDF variations: {theory_corr_hist_name_pdfvars}" + ) + print_flush( + f"Histogram axes: {[ax.name for ax in h_theory_corr_pdfvars.axes]}" + ) + + # Load for all signal processes + for proc in signal_procs: + if proc in results and "output" in results[proc]: + if theory_corr_hist_name_pdfvars in results[proc]["output"]: + h_proxy = results[proc]["output"][theory_corr_hist_name_pdfvars] + h_theory_corr_pdfvars_dict[proc] = ( + h_proxy.get() if hasattr(h_proxy, "get") else h_proxy + ) + +h5file.close() + +# Use first MC process as expected data (Asimov data) +if h_mc_dict: + first_mc_proc = list(h_mc_dict.keys())[0] + h_data = h_mc_dict[first_mc_proc].copy() + print_flush(f"Using MC process '{first_mc_proc}' as expected data (Asimov)") +else: + raise RuntimeError("No MC processes found to use as data") + +# Handle vars axis if present +axis_names = [ax.name for ax in h_data.axes] if hasattr(h_data, "axes") else [] +has_vars = "vars" in axis_names + +if has_vars: + h_data_base = h_data[{"vars": 0}] + h_mc_base = {proc: h[{"vars": 0}] for proc, h in h_mc_dict.items()} +else: + h_data_base = h_data + h_mc_base = h_mc_dict + +# Build tensor +writer = tensorwriter.TensorWriter( + sparse=args.sparse, + systematic_type=args.systematicType, +) + +channel_name = "ch0" +writer.add_channel(h_data_base.axes, channel_name) +writer.add_data(h_data_base, channel_name) + +background_procs = [p for p in mc_procs if p not in signal_procs] + +# Add processes +for proc in signal_procs: + if proc in h_mc_base: + writer.add_process(h_mc_base[proc], proc, channel_name, signal=False) + +for proc in background_procs: + if proc in h_mc_base: + writer.add_process(h_mc_base[proc], proc, channel_name, signal=False) + +# Extract PDF variations from pdfvars correction histogram (noi=False, default) +if h_theory_corr_pdfvars is not None and signal_procs: + proc_name = signal_procs[0] + h_sig_corr_pdfvars = h_theory_corr_pdfvars_dict.get( + proc_name, h_theory_corr_pdfvars + ) + + if "vars" in h_sig_corr_pdfvars.axes.name: + vars_axis = h_sig_corr_pdfvars.axes["vars"] + print_flush( + f"Found pdfvars vars axis with entries: {[str(v) for v in vars_axis]}" + ) + + # Extract all PDF variations (exclude alpha_s variations which have "_as_" in the name) + pdf_variations = [] + for var_name in vars_axis: + var_str = str(var_name) + # Skip central (index 0) and alpha_s variations + if var_str != "central" and "_as_" not in var_str: + pdf_variations.append(var_str) + + if pdf_variations: + print_flush( + f"Found {len(pdf_variations)} PDF variations: {pdf_variations[:5]}..." + ) # Print first 5 + + for var_name_up, var_name_down in zip( + pdf_variations[1::2], pdf_variations[2::2] + ): + h_var_up = h_sig_corr_pdfvars[{"vars": var_name_up}] + h_var_down = h_sig_corr_pdfvars[{"vars": var_name_down}] + + h_pdf_var_up = h_var_up.project(*h_mc_base[proc_name].axes.name) + h_pdf_var_down = h_var_down.project(*h_mc_base[proc_name].axes.name) + + syst_name = f"{var_name_up}_{var_name_down}" + + writer.add_systematic( + [h_pdf_var_up, h_pdf_var_down], # List of [up, down] histograms + syst_name, + proc_name, + "ch0", + constrained=True, + # mirror=False + ) + + num_pairs = len(pdf_variations) // 2 + print_flush( + f"Added {num_pairs} PDF systematic pairs for process {proc_name} (noi=False)" + ) + else: + print_flush(f"Warning: No PDF variations found in pdfvars histogram") + else: + print_flush( + f"Warning: pdfvars correction histogram does not have vars axis. Axes: {[ax.name for ax in h_sig_corr_pdfvars.axes]}" + ) + +# Extract alphas variations from pdfas correction histogram (noi=True) +if h_theory_corr_pdfas is not None and signal_procs: + proc_name = signal_procs[0] + h_sig_corr_pdfas = h_theory_corr_pdfas_dict.get(proc_name, h_theory_corr_pdfas) + + if "vars" in h_sig_corr_pdfas.axes.name: + vars_axis = h_sig_corr_pdfas.axes["vars"] + print_flush( + f"Found pdfas vars axis with entries: {[str(v) for v in vars_axis]}" + ) + + var_name_0120 = "pdfCT18ZNNLO_as_0120" + var_name_0116 = "pdfCT18ZNNLO_as_0116" + + h_alphas_0120 = None + h_alphas_0116 = None + + if var_name_0120 in vars_axis: + h_var = h_sig_corr_pdfas[{"vars": var_name_0120}] + h_alphas_0120 = h_var.project(*h_mc_base[proc_name].axes.name) + print_flush(f"Found alphas 0120 variation: {var_name_0120}") + + if var_name_0116 in vars_axis: + h_var = h_sig_corr_pdfas[{"vars": var_name_0116}] + h_alphas_0116 = h_var.project(*h_mc_base[proc_name].axes.name) + print_flush(f"Found alphas 0116 variation: {var_name_0116}") + + if h_alphas_0120 is not None and h_alphas_0116 is not None: + writer.add_systematic( + [h_alphas_0120, h_alphas_0116], + "pdfAlphaS", + proc_name, + "ch0", + constrained=False, + noi=True, # Only alpha_s variations have noi=True + ) + print_flush( + f"Added pdfAlphaS systematic for process {proc_name} (noi=True)" + ) + else: + print_flush(f"Warning: Could not find both alpha_s variations") + else: + print_flush( + f"Warning: pdfas correction histogram does not have vars axis. Axes: {[ax.name for ax in h_sig_corr_pdfas.axes]}" + ) + +# Write output +writer.write(outfolder=args.output, outfilename=args.outname) +print_flush(f"Tensor written to: {args.output}/{args.outname}.hdf5") From f1a873eb3762acebbcdbedb4181c14639d29747b Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Wed, 1 Jul 2026 06:17:16 -0400 Subject: [PATCH 05/10] calibrations finished, have not yet uploaded root files --- scripts/histmakers/mz_5TeV.py | 14 ++- .../production/include/lowpu_muonscarekit.hpp | 115 +++++++++++------- 2 files changed, 82 insertions(+), 47 deletions(-) diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index 632b63a9c..d07824b6b 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -17,6 +17,7 @@ logger = logging.setup_logger(__file__, args.verbose, args.noColorLogger) import hist +import ROOT import narf from wremnants.production import ( @@ -35,6 +36,9 @@ elif args.muonCorr == "scarekit": narf.clingutils.Load("libROOTDataFrame") narf.clingutils.Declare('#include "lowpu_muonscarekit.hpp"') + scarekit_mc_helper = ROOT.wrem.MuonScarekitMCHelper( + args.randomSeedForToys, ROOT.ROOT.GetThreadPoolSize() + ) datasets = getDatasets( maxFiles=args.maxFiles, @@ -182,7 +186,15 @@ def build_graph(df, dataset): else: df = df.Define( "Muon_pt_corr", - "wrem::applyMuonScarekitMC(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers, run, luminosityBlock)", + scarekit_mc_helper, + [ + "rdfslot_", + "Muon_pt", + "Muon_eta", + "Muon_phi", + "Muon_charge", + "Muon_nTrackerLayers", + ], ) else: # "none" df = df.Alias("Muon_pt_corr", "Muon_pt") diff --git a/wremnants/production/include/lowpu_muonscarekit.hpp b/wremnants/production/include/lowpu_muonscarekit.hpp index 241ca1e34..da5f3e690 100644 --- a/wremnants/production/include/lowpu_muonscarekit.hpp +++ b/wremnants/production/include/lowpu_muonscarekit.hpp @@ -6,8 +6,10 @@ #include #include #include -#include #include +#include +#include +#include namespace wrem { @@ -98,54 +100,75 @@ Vec_f applyMuonScarekitData(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { return res; } -Vec_f applyMuonScarekitMC(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge, - Vec_i nTrackerLayers, unsigned int run, - unsigned int lumi) { - using namespace muonscarekit_impl; - unsigned int size = pt.size(); - Vec_f res(size); - - for (unsigned int i = 0; i < size; ++i) { - double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); - double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); - double pt_scale = 1.0 / (M / pt[i] + charge[i] * A); - - Int_t etabin = h_cb->GetXaxis()->FindBin(fabs((double)eta[i])); - Int_t nlbin = h_cb->GetYaxis()->FindBin((double)nTrackerLayers[i]); - - double mean_cb = h_cb->GetBinContent(etabin, nlbin, 1); - double sig_cb = h_cb->GetBinContent(etabin, nlbin, 2); - double n_cb = h_cb->GetBinContent(etabin, nlbin, 3); - double alpha_cb = h_cb->GetBinContent(etabin, nlbin, 4); - - double a_poly = h_poly->GetBinContent(etabin, nlbin, 1); - double b_poly = h_poly->GetBinContent(etabin, nlbin, 2); - double c_poly = h_poly->GetBinContent(etabin, nlbin, 3); - double sigma_poly = - a_poly + b_poly * pt_scale + c_poly * pt_scale * pt_scale; - if (sigma_poly < 0.0) - sigma_poly = 0.0; - - Int_t absetabin = h_k_data->GetXaxis()->FindBin(fabs((double)eta[i])); - double k_data_v = h_k_data->GetBinContent(absetabin, 3); - double k_sig_v = h_k_sig->GetBinContent(absetabin, 3); - double k_mc = (k_sig_v < k_data_v) - ? sqrt(k_data_v * k_data_v - k_sig_v * k_sig_v) - : 0.0; - - if (k_mc == 0.0 || sigma_poly == 0.0 || n_cb <= 1.0 + 1e-6 || - sig_cb <= 0.0 || alpha_cb <= 0.0) { - res[i] = static_cast(pt_scale); - continue; +class MuonScarekitMCHelper { + +public: + MuonScarekitMCHelper(const std::size_t seed = 0, + const unsigned int nslots = 1) { + const unsigned int nslotsactual = std::max(nslots, 1U); + rng_.reserve(nslotsactual); + auto const hash = std::hash()("MuonScarekitMCHelper"); + for (std::size_t islot = 0; islot < nslotsactual; ++islot) { + std::seed_seq seq{hash, seed, islot}; + rng_.emplace_back(seq); } + } - MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); - double rndm_cb = cb.invcdf(gRandom->Rndm()); - - res[i] = static_cast(pt_scale * (1.0 + k_mc * sigma_poly * rndm_cb)); + Vec_f operator()(const unsigned int slot, Vec_f pt, Vec_f eta, Vec_f phi, + Vec_i charge, Vec_i nTrackerLayers) { + using namespace muonscarekit_impl; + auto &rngslot = rng_[slot]; + std::uniform_real_distribution unif(0., 1.); + + unsigned int size = pt.size(); + Vec_f res(size); + + for (unsigned int i = 0; i < size; ++i) { + double M = h_M_SIG->GetBinContent(h_M_SIG->FindBin(eta[i], phi[i])); + double A = h_A_SIG->GetBinContent(h_A_SIG->FindBin(eta[i], phi[i])); + double pt_scale = 1.0 / (M / pt[i] + charge[i] * A); + + Int_t etabin = h_cb->GetXaxis()->FindBin(fabs((double)eta[i])); + Int_t nlbin = h_cb->GetYaxis()->FindBin((double)nTrackerLayers[i]); + + double mean_cb = h_cb->GetBinContent(etabin, nlbin, 1); + double sig_cb = h_cb->GetBinContent(etabin, nlbin, 2); + double n_cb = h_cb->GetBinContent(etabin, nlbin, 3); + double alpha_cb = h_cb->GetBinContent(etabin, nlbin, 4); + + double a_poly = h_poly->GetBinContent(etabin, nlbin, 1); + double b_poly = h_poly->GetBinContent(etabin, nlbin, 2); + double c_poly = h_poly->GetBinContent(etabin, nlbin, 3); + double sigma_poly = + a_poly + b_poly * pt_scale + c_poly * pt_scale * pt_scale; + if (sigma_poly < 0.0) + sigma_poly = 0.0; + + Int_t absetabin = h_k_data->GetXaxis()->FindBin(fabs((double)eta[i])); + double k_data_v = h_k_data->GetBinContent(absetabin, 3); + double k_sig_v = h_k_sig->GetBinContent(absetabin, 3); + double k_mc = (k_sig_v < k_data_v) + ? sqrt(k_data_v * k_data_v - k_sig_v * k_sig_v) + : 0.0; + + if (k_mc == 0.0 || sigma_poly == 0.0 || n_cb <= 1.0 + 1e-6 || + sig_cb <= 0.0 || alpha_cb <= 0.0) { + res[i] = static_cast(pt_scale); + continue; + } + + MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); + double rndm_cb = cb.invcdf(unif(rngslot)); + + res[i] = + static_cast(pt_scale * (1.0 + k_mc * sigma_poly * rndm_cb)); + } + return res; } - return res; -} + +private: + std::vector rng_; +}; } // namespace wrem #endif From b516e6ee3ba77a10bb85f30a741897a28b26ec8d Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Wed, 1 Jul 2026 12:21:44 -0400 Subject: [PATCH 06/10] actually thread safe now --- scripts/histmakers/mz_5TeV.py | 8 ++--- .../production/include/lowpu_muonscarekit.hpp | 31 +++++++++---------- 2 files changed, 19 insertions(+), 20 deletions(-) diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index d07824b6b..7d22a3a4a 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -36,9 +36,7 @@ elif args.muonCorr == "scarekit": narf.clingutils.Load("libROOTDataFrame") narf.clingutils.Declare('#include "lowpu_muonscarekit.hpp"') - scarekit_mc_helper = ROOT.wrem.MuonScarekitMCHelper( - args.randomSeedForToys, ROOT.ROOT.GetThreadPoolSize() - ) + scarekit_mc_helper = ROOT.wrem.MuonScarekitMCHelper(args.randomSeedForToys) datasets = getDatasets( maxFiles=args.maxFiles, @@ -188,7 +186,9 @@ def build_graph(df, dataset): "Muon_pt_corr", scarekit_mc_helper, [ - "rdfslot_", + "run", + "luminosityBlock", + "event", "Muon_pt", "Muon_eta", "Muon_phi", diff --git a/wremnants/production/include/lowpu_muonscarekit.hpp b/wremnants/production/include/lowpu_muonscarekit.hpp index da5f3e690..2ec3a752d 100644 --- a/wremnants/production/include/lowpu_muonscarekit.hpp +++ b/wremnants/production/include/lowpu_muonscarekit.hpp @@ -103,21 +103,19 @@ Vec_f applyMuonScarekitData(Vec_f pt, Vec_f eta, Vec_f phi, Vec_i charge) { class MuonScarekitMCHelper { public: - MuonScarekitMCHelper(const std::size_t seed = 0, - const unsigned int nslots = 1) { - const unsigned int nslotsactual = std::max(nslots, 1U); - rng_.reserve(nslotsactual); - auto const hash = std::hash()("MuonScarekitMCHelper"); - for (std::size_t islot = 0; islot < nslotsactual; ++islot) { - std::seed_seq seq{hash, seed, islot}; - rng_.emplace_back(seq); - } - } - - Vec_f operator()(const unsigned int slot, Vec_f pt, Vec_f eta, Vec_f phi, - Vec_i charge, Vec_i nTrackerLayers) { + MuonScarekitMCHelper(const std::size_t seed = 0) + : hash_(std::hash()("MuonScarekitMCHelper")), seed_(seed) {} + + // Per-event seeding (run, lumi, event): reproducible across runs and thread + // counts, and thread-safe (RNG is local to each call). Mirrors + // wrem::SmearingHelper in muon_calibration.hpp. + Vec_f operator()(const unsigned int run, const unsigned int lumi, + const unsigned long long event, Vec_f pt, Vec_f eta, + Vec_f phi, Vec_i charge, Vec_i nTrackerLayers) const { using namespace muonscarekit_impl; - auto &rngslot = rng_[slot]; + std::seed_seq seq{hash_, seed_, std::size_t(run), std::size_t(lumi), + std::size_t(event)}; + std::mt19937 rng(seq); std::uniform_real_distribution unif(0., 1.); unsigned int size = pt.size(); @@ -158,7 +156,7 @@ class MuonScarekitMCHelper { } MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb); - double rndm_cb = cb.invcdf(unif(rngslot)); + double rndm_cb = cb.invcdf(unif(rng)); res[i] = static_cast(pt_scale * (1.0 + k_mc * sigma_poly * rndm_cb)); @@ -167,7 +165,8 @@ class MuonScarekitMCHelper { } private: - std::vector rng_; + const std::size_t hash_; + std::size_t seed_; }; } // namespace wrem From 92402362db4ce2fe18e923b75c9305e76c7ef0b2 Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Thu, 2 Jul 2026 04:40:40 -0400 Subject: [PATCH 07/10] fixes --- scripts/histmakers/mz_5TeV.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index 7d22a3a4a..f904067d4 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -12,6 +12,8 @@ choices=["none", "rochester", "scarekit"], help="Muon momentum correction to apply", ) +# This is the 5 TeV low-PU analysis: default to the 5 TeV era (2017G) +parser.set_defaults(era="2017G") args = parser.parse_args() logger = logging.setup_logger(__file__, args.verbose, args.noColorLogger) From a86b86a035ef00f5cf93124b468b2741c35704e0 Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Fri, 3 Jul 2026 04:32:01 -0400 Subject: [PATCH 08/10] stop tracking new_make_tensor --- scripts/rabbit/new_make_tensor.py | 305 ------------------------------ 1 file changed, 305 deletions(-) delete mode 100644 scripts/rabbit/new_make_tensor.py diff --git a/scripts/rabbit/new_make_tensor.py b/scripts/rabbit/new_make_tensor.py deleted file mode 100644 index 0e3f422c1..000000000 --- a/scripts/rabbit/new_make_tensor.py +++ /dev/null @@ -1,305 +0,0 @@ -import argparse -import os -import sys - -import h5py - - -def print_flush(*args, **kwargs): - print(*args, **kwargs) - sys.stdout.flush() - - -wremnants_base = os.environ.get("WREM_BASE", None) -if wremnants_base is None: - script_dir = os.path.dirname(os.path.abspath(__file__)) - wremnants_base = os.path.abspath(os.path.join(script_dir, "../..")) - if os.path.exists(os.path.join(wremnants_base, "wums")): - sys.path.insert(0, wremnants_base) - -from rabbit import tensorwriter -from wremnants.utilities.io_tools import base_io - -parser = argparse.ArgumentParser() -parser.add_argument("infile", help="Input HDF5 file with histograms") -parser.add_argument("-o", "--output", default="./", help="output directory") -parser.add_argument("--outname", default="my_tensor", help="output file name") -parser.add_argument( - "--histName", default="ptll", help="Histogram name to use (default: ptll)" -) -parser.add_argument( - "--procFilters", - nargs="*", - default=["Zmumu"], - help="Processes to include (default: Zmumu)", -) -parser.add_argument( - "--sparse", default=False, action="store_true", help="Make sparse tensor" -) -parser.add_argument( - "--systematicType", - choices=["log_normal", "normal"], - default="log_normal", - help="probability density for systematic variations", -) -args = parser.parse_args() - -# Load data from HDF5 file -infile_path = os.path.abspath(args.infile) -if not os.path.exists(infile_path): - raise FileNotFoundError(f"Input file not found: {infile_path}") - -h5file = h5py.File(infile_path, "r") -results = base_io.load_results_h5py(h5file) - -hist_name = args.histName -all_procs = [p for p in results.keys() if p != "meta_info"] -print_flush(f"All processes found in file: {all_procs}") - -if args.procFilters: - procs_to_use = [p for p in all_procs if any(filt in p for filt in args.procFilters)] -else: - procs_to_use = all_procs - -print_flush(f"Processes after filtering: {procs_to_use}") - -# All processes are treated as MC (no data loading) -mc_procs = procs_to_use -print_flush(f"MC processes found: {mc_procs}") - -# Load MC histograms -h_mc_dict = {} -for proc in mc_procs: - if "output" in results[proc] and hist_name in results[proc]["output"]: - h_proxy = results[proc]["output"][hist_name] - h_mc_dict[proc] = h_proxy.get() if hasattr(h_proxy, "get") else h_proxy - -# Identify signal processes before closing file -signal_procs = [p for p in mc_procs if "Zmumu" in p or "Ztautau" in p] - - -def _theory_corr_hist_name(results, proc, tag): - """Match histmaker output: __Corr (e.g. ptll_..._pdfas_Corr).""" - for name in results[proc]["output"]: - if name.endswith("_Corr") and tag in name: - return name - return None - - -h_theory_corr_pdfas = None -h_theory_corr_pdfvars = None -h_theory_corr_pdfas_dict = {} -h_theory_corr_pdfvars_dict = {} - -if signal_procs: - first_sig_proc = signal_procs[0] - if first_sig_proc in results and "output" in results[first_sig_proc]: - theory_corr_hist_name_pdfas = _theory_corr_hist_name( - results, first_sig_proc, "pdfas" - ) - theory_corr_hist_name_pdfvars = _theory_corr_hist_name( - results, first_sig_proc, "pdfvars" - ) - if not theory_corr_hist_name_pdfas: - print_flush( - f"Warning: no pdfas theory hist in {first_sig_proc}; " - f"keys: {list(results[first_sig_proc]['output'].keys())}" - ) - if not theory_corr_hist_name_pdfvars: - print_flush(f"Warning: no pdfvars theory hist in {first_sig_proc}") - # Load pdfas correction (for alpha_s variations) - if ( - theory_corr_hist_name_pdfas - and theory_corr_hist_name_pdfas in results[first_sig_proc]["output"] - ): - h_proxy = results[first_sig_proc]["output"][theory_corr_hist_name_pdfas] - h_theory_corr_pdfas = h_proxy.get() if hasattr(h_proxy, "get") else h_proxy - print_flush( - f"Found theory correction histogram for alpha_s variations: {theory_corr_hist_name_pdfas}" - ) - print_flush( - f"Histogram axes: {[ax.name for ax in h_theory_corr_pdfas.axes]}" - ) - - # Load for all signal processes - for proc in signal_procs: - if proc in results and "output" in results[proc]: - if theory_corr_hist_name_pdfas in results[proc]["output"]: - h_proxy = results[proc]["output"][theory_corr_hist_name_pdfas] - h_theory_corr_pdfas_dict[proc] = ( - h_proxy.get() if hasattr(h_proxy, "get") else h_proxy - ) - - # Load pdfvars correction (for PDF variations) - if ( - theory_corr_hist_name_pdfvars - and theory_corr_hist_name_pdfvars in results[first_sig_proc]["output"] - ): - h_proxy = results[first_sig_proc]["output"][theory_corr_hist_name_pdfvars] - h_theory_corr_pdfvars = ( - h_proxy.get() if hasattr(h_proxy, "get") else h_proxy - ) - print_flush( - f"Found theory correction histogram for PDF variations: {theory_corr_hist_name_pdfvars}" - ) - print_flush( - f"Histogram axes: {[ax.name for ax in h_theory_corr_pdfvars.axes]}" - ) - - # Load for all signal processes - for proc in signal_procs: - if proc in results and "output" in results[proc]: - if theory_corr_hist_name_pdfvars in results[proc]["output"]: - h_proxy = results[proc]["output"][theory_corr_hist_name_pdfvars] - h_theory_corr_pdfvars_dict[proc] = ( - h_proxy.get() if hasattr(h_proxy, "get") else h_proxy - ) - -h5file.close() - -# Use first MC process as expected data (Asimov data) -if h_mc_dict: - first_mc_proc = list(h_mc_dict.keys())[0] - h_data = h_mc_dict[first_mc_proc].copy() - print_flush(f"Using MC process '{first_mc_proc}' as expected data (Asimov)") -else: - raise RuntimeError("No MC processes found to use as data") - -# Handle vars axis if present -axis_names = [ax.name for ax in h_data.axes] if hasattr(h_data, "axes") else [] -has_vars = "vars" in axis_names - -if has_vars: - h_data_base = h_data[{"vars": 0}] - h_mc_base = {proc: h[{"vars": 0}] for proc, h in h_mc_dict.items()} -else: - h_data_base = h_data - h_mc_base = h_mc_dict - -# Build tensor -writer = tensorwriter.TensorWriter( - sparse=args.sparse, - systematic_type=args.systematicType, -) - -channel_name = "ch0" -writer.add_channel(h_data_base.axes, channel_name) -writer.add_data(h_data_base, channel_name) - -background_procs = [p for p in mc_procs if p not in signal_procs] - -# Add processes -for proc in signal_procs: - if proc in h_mc_base: - writer.add_process(h_mc_base[proc], proc, channel_name, signal=False) - -for proc in background_procs: - if proc in h_mc_base: - writer.add_process(h_mc_base[proc], proc, channel_name, signal=False) - -# Extract PDF variations from pdfvars correction histogram (noi=False, default) -if h_theory_corr_pdfvars is not None and signal_procs: - proc_name = signal_procs[0] - h_sig_corr_pdfvars = h_theory_corr_pdfvars_dict.get( - proc_name, h_theory_corr_pdfvars - ) - - if "vars" in h_sig_corr_pdfvars.axes.name: - vars_axis = h_sig_corr_pdfvars.axes["vars"] - print_flush( - f"Found pdfvars vars axis with entries: {[str(v) for v in vars_axis]}" - ) - - # Extract all PDF variations (exclude alpha_s variations which have "_as_" in the name) - pdf_variations = [] - for var_name in vars_axis: - var_str = str(var_name) - # Skip central (index 0) and alpha_s variations - if var_str != "central" and "_as_" not in var_str: - pdf_variations.append(var_str) - - if pdf_variations: - print_flush( - f"Found {len(pdf_variations)} PDF variations: {pdf_variations[:5]}..." - ) # Print first 5 - - for var_name_up, var_name_down in zip( - pdf_variations[1::2], pdf_variations[2::2] - ): - h_var_up = h_sig_corr_pdfvars[{"vars": var_name_up}] - h_var_down = h_sig_corr_pdfvars[{"vars": var_name_down}] - - h_pdf_var_up = h_var_up.project(*h_mc_base[proc_name].axes.name) - h_pdf_var_down = h_var_down.project(*h_mc_base[proc_name].axes.name) - - syst_name = f"{var_name_up}_{var_name_down}" - - writer.add_systematic( - [h_pdf_var_up, h_pdf_var_down], # List of [up, down] histograms - syst_name, - proc_name, - "ch0", - constrained=True, - # mirror=False - ) - - num_pairs = len(pdf_variations) // 2 - print_flush( - f"Added {num_pairs} PDF systematic pairs for process {proc_name} (noi=False)" - ) - else: - print_flush(f"Warning: No PDF variations found in pdfvars histogram") - else: - print_flush( - f"Warning: pdfvars correction histogram does not have vars axis. Axes: {[ax.name for ax in h_sig_corr_pdfvars.axes]}" - ) - -# Extract alphas variations from pdfas correction histogram (noi=True) -if h_theory_corr_pdfas is not None and signal_procs: - proc_name = signal_procs[0] - h_sig_corr_pdfas = h_theory_corr_pdfas_dict.get(proc_name, h_theory_corr_pdfas) - - if "vars" in h_sig_corr_pdfas.axes.name: - vars_axis = h_sig_corr_pdfas.axes["vars"] - print_flush( - f"Found pdfas vars axis with entries: {[str(v) for v in vars_axis]}" - ) - - var_name_0120 = "pdfCT18ZNNLO_as_0120" - var_name_0116 = "pdfCT18ZNNLO_as_0116" - - h_alphas_0120 = None - h_alphas_0116 = None - - if var_name_0120 in vars_axis: - h_var = h_sig_corr_pdfas[{"vars": var_name_0120}] - h_alphas_0120 = h_var.project(*h_mc_base[proc_name].axes.name) - print_flush(f"Found alphas 0120 variation: {var_name_0120}") - - if var_name_0116 in vars_axis: - h_var = h_sig_corr_pdfas[{"vars": var_name_0116}] - h_alphas_0116 = h_var.project(*h_mc_base[proc_name].axes.name) - print_flush(f"Found alphas 0116 variation: {var_name_0116}") - - if h_alphas_0120 is not None and h_alphas_0116 is not None: - writer.add_systematic( - [h_alphas_0120, h_alphas_0116], - "pdfAlphaS", - proc_name, - "ch0", - constrained=False, - noi=True, # Only alpha_s variations have noi=True - ) - print_flush( - f"Added pdfAlphaS systematic for process {proc_name} (noi=True)" - ) - else: - print_flush(f"Warning: Could not find both alpha_s variations") - else: - print_flush( - f"Warning: pdfas correction histogram does not have vars axis. Axes: {[ax.name for ax in h_sig_corr_pdfas.axes]}" - ) - -# Write output -writer.write(outfolder=args.output, outfilename=args.outname) -print_flush(f"Tensor written to: {args.output}/{args.outname}.hdf5") From 8e13a831aaa79401ff3b4fadd4290bcec10b5cc6 Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Fri, 3 Jul 2026 04:33:07 -0400 Subject: [PATCH 09/10] restore gitignore --- .gitignore | 1 - 1 file changed, 1 deletion(-) diff --git a/.gitignore b/.gitignore index 8afe8a6f5..dce054ca1 100644 --- a/.gitignore +++ b/.gitignore @@ -158,4 +158,3 @@ cython_debug/ *.pdf *.png *.hdf5 -*.root \ No newline at end of file From 5f58f8393dce7708e8ad2b263afa1640f6de7e3e Mon Sep 17 00:00:00 2001 From: Hayden Tedrake Date: Fri, 3 Jul 2026 04:35:05 -0400 Subject: [PATCH 10/10] try again --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index dce054ca1..8afe8a6f5 100644 --- a/.gitignore +++ b/.gitignore @@ -158,3 +158,4 @@ cython_debug/ *.pdf *.png *.hdf5 +*.root \ No newline at end of file