Implementing muonscarekit corrections

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<float>(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)

wremnants/production/include/lowpu_muonscarekit.hpp

@@ -0,0 +1,173 @@
+#ifndef WREMNANTS_LOWPU_MUONSCAREKIT_H
+#define WREMNANTS_LOWPU_MUONSCAREKIT_H
+
+#include "defines.hpp"
+#include <TFile.h>
+#include <TH2D.h>
+#include <TH3D.h>
+#include <TMath.h>
+#include <boost/math/special_functions/erf.hpp>
+#include <random>
+#include <string>
+#include <vector>
+
+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<float>(1.0 / (M / pt[i] + charge[i] * A));
+  }
+  return res;
+}
+
+class MuonScarekitMCHelper {
+
+public:
+  MuonScarekitMCHelper(const std::size_t seed = 0)
+      : hash_(std::hash<std::string>()("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<double> 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<float>(pt_scale);
+        continue;
+      }
+
+      MuonScarekitCB cb(mean_cb, sig_cb, alpha_cb, n_cb);
+      double rndm_cb = cb.invcdf(unif(rng));
+
+      res[i] =
+          static_cast<float>(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