diff --git a/scripts/histmakers/mz_5TeV.py b/scripts/histmakers/mz_5TeV.py index 53572dc3e..f904067d4 100644 --- a/scripts/histmakers/mz_5TeV.py +++ b/scripts/histmakers/mz_5TeV.py @@ -12,19 +12,14 @@ 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", -) - +# 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) import hist +import ROOT import narf from wremnants.production import ( @@ -43,6 +38,7 @@ elif args.muonCorr == "scarekit": narf.clingutils.Load("libROOTDataFrame") narf.clingutils.Declare('#include "lowpu_muonscarekit.hpp"') + scarekit_mc_helper = ROOT.wrem.MuonScarekitMCHelper(args.randomSeedForToys) datasets = getDatasets( maxFiles=args.maxFiles, @@ -182,41 +178,26 @@ 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::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)", - ) - 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_noKFactor(Muon_pt, Muon_eta, Muon_phi, Muon_charge, Muon_nTrackerLayers)", - ) - 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", + scarekit_mc_helper, + [ + "run", + "luminosityBlock", + "event", + "Muon_pt", + "Muon_eta", + "Muon_phi", + "Muon_charge", + "Muon_nTrackerLayers", + ], + ) else: # "none" df = df.Alias("Muon_pt_corr", "Muon_pt") @@ -314,8 +295,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 +313,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 +443,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, @@ -502,5 +490,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) diff --git a/wremnants/production/include/lowpu_muonscarekit.hpp b/wremnants/production/include/lowpu_muonscarekit.hpp new file mode 100644 index 000000000..2ec3a752d --- /dev/null +++ b/wremnants/production/include/lowpu_muonscarekit.hpp @@ -0,0 +1,173 @@ +#ifndef WREMNANTS_LOWPU_MUONSCAREKIT_H +#define WREMNANTS_LOWPU_MUONSCAREKIT_H + +#include "defines.hpp" +#include +#include +#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_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; +} + +class MuonScarekitMCHelper { + +public: + 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; + 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(); + 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(rng)); + + res[i] = + static_cast(pt_scale * (1.0 + k_mc * sigma_poly * rndm_cb)); + } + return res; + } + +private: + const std::size_t hash_; + std::size_t seed_; +}; + +} // namespace wrem +#endif