inclusive_detached_dilepton/pre_selection_cuts.cpp

#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 "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 <string>
#include <iostream>
#include <cmath>

const int nBins = 150;

RooPlot *CreateRooFit(TH1D *hist);

int pre_selection_cuts()
{
    // just some plotting options
    gROOT->SetStyle("Plain");
    TPad foo;
    const int NRGBs = 5;
    const int NCont = 250;
    double stops[NRGBs] = {0.00, 0.34, 0.61, 0.84, 1.00};
    double red[NRGBs] = {0.00, 0.00, 0.87, 1.00, 0.51};
    double green[NRGBs] = {0.00, 0.81, 1.00, 0.20, 0.00};
    double blue[NRGBs] = {0.51, 1.00, 0.12, 0.00, 0.00};
    TColor::CreateGradientColorTable(NRGBs, stops, red, green, blue, NCont);
    gStyle->SetNumberContours(NCont);
    gStyle->SetLabelFont(132, "xyz");
    gStyle->SetTitleFont(132, "xyz");
    gStyle->SetLegendFont(132);
    gStyle->SetStatFont(132);
    gStyle->SetEndErrorSize(10.0);
    gStyle->SetOptStat(0);
    gStyle->SetOptFit(0);

    // files to load
    std::vector<std::string> data_filenames =
        {
            "/auto/data/pfeiffer/inclusive_detached_dilepton/data_samples/BuToKpMuMu_Collision23_Beam6800GeV-VeloClosed-MagDown-Excl-UT_RealData_Sprucing23r1_90000000_RD.root"};

    TChain *data_chain = new TChain("SpruceRD_BuToKpMuMu/DecayTree");
    for (unsigned int i = 0; i < data_filenames.size(); i++)
    {
        data_chain->Add(data_filenames.at(i).c_str());
    }

    // files to load
    std::vector<std::string> mc_filenames =
        {
            "/auto/data/pfeiffer/inclusive_detached_dilepton/MC/BuToKpMuMu_rd_btoxll_simulation_12143001_MagDown_v0r0p6316987.root"};

    TChain *mc_chain = new TChain("BuToKpMuMu23_noPID/DecayTree");
    for (unsigned int i = 0; i < mc_filenames.size(); i++)
    {
        mc_chain->Add(mc_filenames.at(i).c_str());
    }

    Double_t Bplus_M, // 4400 -> 8200
        Jpsi_M,       // 200 -> 6600
        Kplus_M;      // 493.6769 -> 493.6779

    Float_t Bplus_PT,    // 0 -> 30 * 10^3
        Bplus_BPVFDCHI2, // 0 -> 7000
        Jpsi_BPVIPCHI2,  // 0 -> 3
        Jpsi_PT,         // 0 -> 26 * 10^3
        Kplus_BPVIPCHI2, // 0 -> 7
        Kplus_PT;        // 0 -> 11 * 10^3

    Double_t Bplus_M_mc, // 4400 -> 8200
        Jpsi_M_mc,       // 200 -> 6600
        Kplus_M_mc;      // 493.6769 -> 493.6779

    Float_t Bplus_PT_mc,    // 0 -> 30 * 10^3
        Bplus_BPVFDCHI2_mc, // 0 -> 7000
        Jpsi_BPVIPCHI2_mc,  // 0 -> 3
        Jpsi_PT_mc,         // 0 -> 26 * 10^3
        Kplus_BPVIPCHI2_mc, // 0 -> 7
        Kplus_PT_mc;        // 0 -> 11 * 10^3

    data_chain->SetBranchAddress("Bplus_M", &Bplus_M);
    data_chain->SetBranchAddress("Bplus_PT", &Bplus_PT);
    data_chain->SetBranchAddress("Bplus_BPVFDCHI2", &Bplus_BPVFDCHI2);
    data_chain->SetBranchAddress("Jpsi_M", &Jpsi_M);
    data_chain->SetBranchAddress("Jpsi_BPVIPCHI2", &Jpsi_BPVIPCHI2);
    data_chain->SetBranchAddress("Jpsi_PT", &Jpsi_PT);
    data_chain->SetBranchAddress("Kplus_M", &Kplus_M);
    data_chain->SetBranchAddress("Kplus_BPVIPCHI2", &Kplus_BPVIPCHI2);
    data_chain->SetBranchAddress("Kplus_PT", &Kplus_PT);

    mc_chain->SetBranchAddress("B_M", &Bplus_M_mc);
    mc_chain->SetBranchAddress("B_PT", &Bplus_PT_mc);
    mc_chain->SetBranchAddress("B_BPVFDCHI2", &Bplus_BPVFDCHI2_mc);
    mc_chain->SetBranchAddress("Jpsi_M", &Jpsi_M_mc);
    mc_chain->SetBranchAddress("Jpsi_BPVIPCHI2", &Jpsi_BPVIPCHI2_mc);
    mc_chain->SetBranchAddress("Jpsi_PT", &Jpsi_PT_mc);
    mc_chain->SetBranchAddress("K_M", &Kplus_M_mc);
    mc_chain->SetBranchAddress("K_BPVIPCHI2", &Kplus_BPVIPCHI2_mc);
    mc_chain->SetBranchAddress("K_PT", &Kplus_PT_mc);

    TH1D *h1_Bplus_M = new TH1D("h1_Bplus_M", "B+ Mass", nBins, 4500, 7500);
    TH1D *h1_Bplus_PT = new TH1D("h1_Bplus_PT", "B+ P_{T}", nBins, 0, 15000);
    TH1D *h1_Bplus_BPVFDCHI2 = new TH1D("h1_Bplus_BPVFDCHI2", "B+ FD #chi^{2}", nBins, 0, 1500);
    TH1D *h1_Jpsi_M = new TH1D("h1_Jpsi_M", "J/#psi Mass", nBins, 200, 6600);
    TH1D *h1_Jpsi_BPVIPCHI2 = new TH1D("h1_Jpsi_BPVIPCHI2", "J/#psi IP #chi^{2}", nBins, 0, 100);
    TH1D *h1_Jpsi_PT = new TH1D("h1_Jpsi_PT", "J/#psi P_{T}", nBins, 0, 15000);
    TH1D *h1_Kplus_M = new TH1D("h1_Kplus_M", "K^{+} Mass", nBins, 493.676999999, 493.677000001);
    TH1D *h1_Kplus_BPVIPCHI2 = new TH1D("h1_Kplus_BPVIPCHI2", "K^{+} IP #chi^{2}", nBins, 0, 100);
    TH1D *h1_Kplus_PT = new TH1D("h1_Kplus_PT", "K^{+} P_{T}", nBins, 0, 6500);

    TH1D *h1_bkg_Bplus_M = new TH1D("h1_bkg_Bplus_M", "B+ Mass", nBins, 4500, 7500);
    TH1D *h1_bkg_Bplus_PT = new TH1D("h1_bkg_Bplus_PT", "B+ P_{T}", nBins, 0, 15000);
    TH1D *h1_bkg_Bplus_BPVFDCHI2 = new TH1D("h1_bkg_Bplus_BPVFDCHI2", "B+ FD #chi^{2}", nBins, 0, 1500);
    TH1D *h1_bkg_Jpsi_M = new TH1D("h1_bkg_Jpsi_M", "J/#psi Mass", nBins, 200, 6600);
    TH1D *h1_bkg_Jpsi_BPVIPCHI2 = new TH1D("h1_bkg_Jpsi_BPVIPCHI2", "J/#psi IP #chi^{2}", nBins, 0, 100);
    TH1D *h1_bkg_Jpsi_PT = new TH1D("h1_bkg_Jpsi_PT", "J/#psi P_{T}", nBins, 0, 15000);
    TH1D *h1_bkg_Kplus_M = new TH1D("h1_bkg_Kplus_M", "K^{+} Mass", nBins, 493.676999999, 493.677000001);
    TH1D *h1_bkg_Kplus_BPVIPCHI2 = new TH1D("h1_bkg_Kplus_BPVIPCHI2", "K^{+} IP #chi^{2}", nBins, 0, 100);
    TH1D *h1_bkg_Kplus_PT = new TH1D("h1_bkg_Kplus_PT", "K^{+} P_{T}", nBins, 0, 6500);

    unsigned int entries = data_chain->GetEntries();
    // loop over all entries in the tree
    for (unsigned int i = 0; i < entries; i++)
    {
        data_chain->GetEntry(i);
        const double JSPI_MASS = 3096.;

        if (!std::isfinite(Jpsi_BPVIPCHI2) || !std::isfinite(Jpsi_M)) {
            std::cout << "INF Value (" << i << ") Jpsi_BPVIPCHI2: " << Jpsi_BPVIPCHI2 << ", Jpsi_M: " << Jpsi_M << std::endl;
        }

        /*
        B+ PT > 1500.
        B+ FDCHI2 > 100.
        J/Psi IP > 0.05
        J/Psi PT > 1000.
        K+ IP > 0.2
        K+ PT > 650.
        eher auf IP CHI2 schneiden, statt IP. da eher signifikanz
        */
        if (Bplus_PT > 1200. && Bplus_BPVFDCHI2 > 90. && /*Jpsi_IP_OWNPV > 0.05 && */
            Jpsi_PT > 800. && Kplus_BPVIPCHI2 > 0.17 && Kplus_PT > 500. && (JSPI_MASS - 150. < Jpsi_M && Jpsi_M < JSPI_MASS + 150.))
        {
            h1_Bplus_M->Fill(Bplus_M);
        }
        else
        {
            h1_bkg_Bplus_M->Fill(Bplus_M);
        }

        h1_bkg_Bplus_PT->Fill(Bplus_PT);
        h1_bkg_Bplus_BPVFDCHI2->Fill(Bplus_BPVFDCHI2);
        h1_bkg_Jpsi_M->Fill(Jpsi_M);
        h1_bkg_Jpsi_BPVIPCHI2->Fill(Jpsi_BPVIPCHI2);
        h1_bkg_Jpsi_PT->Fill(Jpsi_PT);
        h1_bkg_Kplus_M->Fill(Kplus_M);
        h1_bkg_Kplus_BPVIPCHI2->Fill(Kplus_BPVIPCHI2);
        h1_bkg_Kplus_PT->Fill(Kplus_PT);
    }

    unsigned int mc_entries = mc_chain->GetEntries();

    for (size_t i = 0; i < mc_entries; i++)
    {
        mc_chain->GetEntry(i);
        h1_Bplus_PT->Fill(Bplus_PT_mc);
        h1_Bplus_BPVFDCHI2->Fill(Bplus_BPVFDCHI2_mc);
        h1_Jpsi_M->Fill(Jpsi_M_mc);
        h1_Jpsi_BPVIPCHI2->Fill(Jpsi_BPVIPCHI2_mc);
        h1_Jpsi_PT->Fill(Jpsi_PT_mc);
        h1_Kplus_M->Fill(Kplus_M);
        h1_Kplus_BPVIPCHI2->Fill(Kplus_BPVIPCHI2_mc);
        h1_Kplus_PT->Fill(Kplus_PT_mc);
    }

    TCanvas *c1 = new TCanvas("c1", "Canvas 1", 0, 0, 800, 600);

    h1_bkg_Bplus_M->SetLineColor(kRed);
    h1_bkg_Bplus_M->Draw();
    h1_Bplus_M->Draw();

    c1->Draw();

    std::vector<TH1 *> small_histos{h1_Bplus_PT, h1_Bplus_BPVFDCHI2, h1_Jpsi_M, h1_Jpsi_BPVIPCHI2, h1_Jpsi_PT, h1_Kplus_M, h1_Kplus_BPVIPCHI2, h1_Kplus_PT};
    std::vector<TH1 *> small_bkg_histos{h1_bkg_Bplus_PT, h1_bkg_Bplus_BPVFDCHI2, h1_bkg_Jpsi_M, h1_bkg_Jpsi_BPVIPCHI2, h1_bkg_Jpsi_PT, h1_bkg_Kplus_M, h1_bkg_Kplus_BPVIPCHI2, h1_bkg_Kplus_PT};

    TCanvas *c2 = new TCanvas("c2", "Canvas 2", 0, 0, 1200, 600);
    c2->Divide(4, 2);

    for (size_t i = 0; i < small_histos.size(); i++)
    {
        c2->cd(i + 1);

        small_histos[i]->SetLineWidth(1);
        small_histos[i]->SetMinimum(0);
        small_histos[i]->SetLineColor(kRed);
        small_histos[i]->SetFillColor(kRed);
        small_histos[i]->SetFillStyle(3004);
        small_histos[i]->Scale(1. / small_histos[i]->Integral(), "width");

        small_bkg_histos[i]->SetLineColor(kBlue);
        small_bkg_histos[i]->SetLineWidth(2);
        small_bkg_histos[i]->SetMinimum(0);
        small_bkg_histos[i]->Scale(1. / small_bkg_histos[i]->Integral(), "width");

        if (small_histos[i]->GetMaximum() > small_bkg_histos[i]->GetMaximum())
        {
            small_histos[i]->Draw("HIST");
            small_bkg_histos[i]->Draw("HIST SAME");
        }
        else
        {
            small_bkg_histos[i]->Draw("HIST");
            small_histos[i]->Draw("HIST SAME");
        }
    }

    c2->Draw();

    /*
    TCanvas *c3 = new TCanvas("c3", "Canvas 3", 0, 0, 800, 600);

    auto fitFrame = CreateRooFit(h1_Bplus_M);
    fitFrame->Draw();

    c3->Draw();
    */

    return 0;
}

RooPlot *CreateRooFit(TH1D *hist)
{
    RooRealVar roo_var_mass("roo_var_mass", "B+ Mass Variable", 4500., 7500.);
    roo_var_mass.setRange("fitting_range", 4700., 6500.);

    TString hist_name = "roohist_bplus_M";
    RooDataHist roohist_bplus_M(hist_name, "B Plus Mass Histogram", roo_var_mass, RooFit::Import(*hist));

    RooRealVar roo_sig_gauss_mean("roo_sig_gauss_mean", "Mass Gauss Mean", 5250., 5100., 5400.);
    RooRealVar roo_sig_gauss_sigma("roo_sig_gauss_sigma", "Mass Gauss Sigma", 60., 0., 150.);

    RooGaussian roo_sig_gauss("roo_sig_gauss", "B+ Mass Signal Gaussian", roo_var_mass, roo_sig_gauss_mean, roo_sig_gauss_sigma);

    RooRealVar roo_bkg_exp_c("roo_bkg_exp_c", "Background C", -0.00117202, -0.003, -0.001);
    RooExponential roo_bkg_exp("roo_bkg_exp", "B+ Mass Background Exp", roo_var_mass, roo_bkg_exp_c);

    RooRealVar roo_var_mass_nsig("roo_var_mass_nsig", "B+ Mass N Signal", 0., hist->GetEntries());
    RooRealVar roo_var_mass_nbkg("roo_var_mass_nbkg", "B+ Mass N Background", 0., hist->GetEntries());

    TString pdf_name = "roo_pdf_sig_plus_bkg";
    RooAddPdf roo_pdf_sig_plus_bkg(pdf_name, "Sig + Bkg PDF",
                                   RooArgList(roo_sig_gauss, roo_bkg_exp),
                                   RooArgList(roo_var_mass_nsig, roo_var_mass_nbkg));

    RooPlot *roo_frame_mass = roo_var_mass.frame();
    roohist_bplus_M.plotOn(roo_frame_mass, RooFit::Binning(nBins), RooFit::Name(hist_name));

    RooFitResult *fitres = roo_pdf_sig_plus_bkg.fitTo(roohist_bplus_M, RooFit::Save(), RooFit::PrintLevel(1), RooFit::Range("fitting_range"));

    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::VisualizeError(*fitres, 1), RooFit::FillColor(kOrange + 1), RooFit::FillStyle(3144));
    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::LineColor(kRed), RooFit::LineStyle(kSolid), RooFit::Range("fitting_range"), RooFit::Name(pdf_name));
    roo_pdf_sig_plus_bkg.plotOn(roo_frame_mass, RooFit::Components(RooArgSet(roo_bkg_exp)), RooFit::LineColor(kBlue), RooFit::LineStyle(kDashed), RooFit::Range("fitting_range"));

    roo_sig_gauss.paramOn(roo_frame_mass, RooFit::Layout(0.60, 0.99, 0.90));
    roo_frame_mass->getAttText()->SetTextSize(0.027);

    return roo_frame_mass;
}