|
|
#include "TH1D.h"
#include "TH2D.h"
#include "THStack.h"
#include "TGraph.h"
#include "TTree.h"
#include "TChain.h"
#include "TFile.h"
#include "TCanvas.h"
#include "TROOT.h"
#include "TStyle.h"
#include "TColor.h"
#include "TLorentzVector.h"
#include "TRandom3.h"
#include "TLorentzVector.h"
#include "TString.h"
#include "RooDataHist.h"
#include "RooRealVar.h"
#include "RooPlot.h"
#include "RooGaussian.h"
#include "RooExponential.h"
#include "RooRealConstant.h"
#include "RooAddPdf.h"
#include "RooFitResult.h"
#include "RooProduct.h"
#include "RooCrystalBall.h"
#include "RooBreitWigner.h"
#include "RooArgSet.h"
#include "RooFFTConvPdf.h"
#include "RooNovosibirsk.h"
#include <string>
#include <iostream>
#include <cmath>
#include <vector>
template <typename T> class Cut { public: virtual bool Check(T y) = 0; virtual std::string ToString() = 0; };
template <typename T> class LT : public Cut<T> { private: T value;
public: LT(T val) : value{val} { }
bool Check(T y) { return y < value; }
std::string ToString() { return TString::Format("(x < %s)", std::to_string(value).c_str()).Data(); } };
template <typename T> class GT : public Cut<T> { private: T value;
public: GT(T val) : value{val} {}
bool Check(T y) { return y > value; }
std::string ToString() { return TString::Format("(x > %s)", std::to_string(value).c_str()).Data(); } };
template <typename T> class EQ : public Cut<T> { private: T value;
public: EQ(T val) : value{val} {}
bool Check(T y) { return y == value; }
std::string ToString() { return TString::Format("(x == %s)", std::to_string(value).c_str()).Data(); } };
template <typename T> class BT : public Cut<T> { private: T low; T up;
public: BT(T lower, T upper) : low{lower}, up{upper} {}
bool Check(T y) { return low < y && y < up; }
std::string ToString() { return TString::Format("(x in [%s, %s])", std::to_string(low).c_str(), std::to_string(up).c_str()).Data(); } };
template <typename T> class Variable { private: const char *name; std::vector<Cut<T> *> cuts;
T value;
public: Variable() { }
Variable(const char *name, std::vector<Cut<T> *> cuts) { this->name = name; this->cuts = cuts; }
T *GetValuePtr() { return &value; }
T GetValue() const { return value; }
bool CheckCuts() const { for (int i = 0; i < cuts.size(); i++) { if (!cuts[i]->Check(value)) { return false; } } return true; } };
class VariableCollection { private: std::map<std::string, Variable<Double_t>> doubles; std::map<std::string, Variable<Float_t>> floats; std::map<std::string, Variable<Int_t>> ints; std::map<std::string, Variable<Bool_t>> bools; std::vector<std::vector<std::string>> trigger_cuts;
TLorentzVector ComputeFourMomentum(const char *name, bool do_mass_sub, double subbed_mass) { std::string pxname = TString::Format("%s_PX", name).Data(); std::string pyname = TString::Format("%s_PY", name).Data(); std::string pzname = TString::Format("%s_PZ", name).Data(); std::string ename = TString::Format("%s_ENERGY", name).Data();
auto itx = floats.find(pxname); if (itx == floats.end()) { return TLorentzVector{}; }
auto ity = floats.find(pyname); if (ity == floats.end()) { return TLorentzVector{}; }
auto itz = floats.find(pzname); if (itz == floats.end()) { return TLorentzVector{}; }
TVector3 momentum(itx->second.GetValue(), ity->second.GetValue(), itz->second.GetValue()); double energy = 0;
if (do_mass_sub) { energy = TMath::Sqrt(TMath::Sq(subbed_mass) + momentum.Mag2()); } else { auto ite = floats.find(ename); if (ite == floats.end()) { return TLorentzVector{}; } energy = ite->second.GetValue(); }
return TLorentzVector(momentum, energy); }
public: VariableCollection() { doubles = std::map<std::string, Variable<Double_t>>{}; floats = std::map<std::string, Variable<Float_t>>{}; ints = std::map<std::string, Variable<Int_t>>{}; bools = std::map<std::string, Variable<Bool_t>>{}; trigger_cuts = std::vector<std::vector<std::string>>{}; }
void Connect(TChain *chain) { for (auto &[key, var] : doubles) { chain->SetBranchAddress(key.c_str(), var.GetValuePtr()); }
for (auto &[key, var] : floats) { chain->SetBranchAddress(key.c_str(), var.GetValuePtr()); }
for (auto &[key, var] : ints) { chain->SetBranchAddress(key.c_str(), var.GetValuePtr()); }
for (auto &[key, var] : bools) { chain->SetBranchAddress(key.c_str(), var.GetValuePtr()); } }
void AddDouble(const char *name, const std::vector<Cut<Double_t> *> &cuts) { doubles[name] = Variable<Double_t>(name, cuts); }
void AddFloat(const char *name, const std::vector<Cut<Float_t> *> &cuts) { floats[name] = Variable<Float_t>(name, cuts); }
void AddInt(const char *name, const std::vector<Cut<Int_t> *> &cuts) { ints[name] = Variable<Int_t>(name, cuts); }
void AddBool(const char *name, const std::vector<Cut<Bool_t> *> &cuts) { bools[name] = Variable<Bool_t>(name, cuts); }
void AddTriggerCut(const std::vector<std::string> &triggers_list) { for (const auto &trig : triggers_list) { bools[trig.c_str()] = Variable<Bool_t>(trig.c_str(), {}); } trigger_cuts.push_back(triggers_list); }
void AddFourMomentumFor(const char *name) { std::string pxname = TString::Format("%s_PX", name).Data(); std::string pyname = TString::Format("%s_PY", name).Data(); std::string pzname = TString::Format("%s_PZ", name).Data(); std::string ename = TString::Format("%s_ENERGY", name).Data();
floats[pxname] = Variable<Float_t>(pxname.c_str(), {}); floats[pyname] = Variable<Float_t>(pyname.c_str(), {}); floats[pzname] = Variable<Float_t>(pzname.c_str(), {}); floats[ename] = Variable<Float_t>(ename.c_str(), {}); }
Double_t GetDouble(const char *name) const { auto it = doubles.find(name); if (it != doubles.end()) { return it->second.GetValue(); } return 0; }
Float_t GetFloat(const char *name) const { auto it = floats.find(name); if (it != floats.end()) { return it->second.GetValue(); } return 0; }
Int_t GetInt(const char *name) const { auto it = ints.find(name); if (it != ints.end()) { return it->second.GetValue(); } return 0; }
Bool_t GetBool(const char *name) const { auto it = bools.find(name); if (it != bools.end()) { return it->second.GetValue(); } return 0; }
TLorentzVector GetFourMomentum(const char *name) { return ComputeFourMomentum(name, false, 0.); }
TLorentzVector GetFourMomentumWithMassSub(const char *name, double subbed_mass) { return ComputeFourMomentum(name, true, subbed_mass); }
bool CheckCuts() const { for (const auto &[key, var] : doubles) { if (!var.CheckCuts()) { return false; } }
for (const auto &[key, var] : floats) { if (!var.CheckCuts()) { return false; } }
for (const auto &[key, var] : ints) { if (!var.CheckCuts()) { return false; } }
for (const auto &[key, var] : bools) { if (!var.CheckCuts()) { return false; } }
for (const auto &triggers : trigger_cuts) { bool cut_okay = false; for (const auto &trigger : triggers) { cut_okay = cut_okay | this->GetBool(trigger.c_str()); }
if (!cut_okay) { return false; } }
return true; }
void PrintValues() const { for (const auto &[key, var] : doubles) { std::cout << key << ": " << var.GetValue() << std::endl; }
for (const auto &[key, var] : floats) { std::cout << key << ": " << var.GetValue() << std::endl; }
for (const auto &[key, var] : ints) { std::cout << key << ": " << var.GetValue() << std::endl; }
for (const auto &[key, var] : bools) { std::cout << key << ": " << var.GetValue() << std::endl; } } };
void BuToHpMuMu() {
TChain *data_chain = new TChain("BuToHpMuMu/DecayTree"); data_chain->Add("/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0r1_tuple_90000000_2023_v0r0p6288631.root"); data_chain->Add("/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/spruce_magdown_2023_v0_tuple_90000000_v0r0p6288631.root");
VariableCollection col{}; col.AddFourMomentumFor("L1"); col.AddFourMomentumFor("L2"); col.AddFourMomentumFor("Hp");
col.AddTriggerCut({"Hlt2_InclDetDiMuonDecision", "Hlt2_InclDetDiMuon_3BodyDecision", "Hlt2_InclDetDiMuon_4BodyDecision"}); col.AddTriggerCut({"Hlt1TrackMVADecision", "Hlt1TwoTrackMVADecision"});
col.Connect(data_chain);
const Double_t MASS_HIST_MIN = 5100.; const Double_t MASS_HIST_MAX = 6000.; const int N_BINS = 70; TH1D *h1_B_M = new TH1D("h1_B_M", "B Mass", N_BINS, MASS_HIST_MIN, MASS_HIST_MAX);
unsigned int entries = data_chain->GetEntries(); for (unsigned int i = 0; i < entries; i++) { data_chain->GetEntry(i);
if (col.CheckCuts()) { auto l1_4v = col.GetFourMomentum("L1"); auto l2_4v = col.GetFourMomentum("L2"); auto K_4v = col.GetFourMomentumWithMassSub("Hp", 493.677);
h1_B_M->Fill((l1_4v + l2_4v + K_4v).M()); }
if ((i + 1) % 10000 == 0 || i + 1 == entries) { std::cout << "[" << "BuToHpMuMu" << "] Processed event: " << i + 1 << " (" << TString::Format("%.2f", ((double)(i + 1) / (double)entries) * 100.) << "%)" << std::endl; } }
TCanvas *c1 = new TCanvas("c1", "c1", 0, 0, 800, 600); h1_B_M->Draw(); c1->Draw();
return 0; }
int do_all() { BuToHpMuMu(); return 0; }
|