inclusive_detached_dilepton/mapmc_b02hphmmumu.cpp

#include <string>
#include <iostream>
#include <cmath>
#include <algorithm>
#include <filesystem>
#include <string_view>

#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 "TLegend.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 "constants.h"
#include "basic_analysis.h"

int mapmc_b02hphmmumu()
{

  const char *sim_tree_name = "B0ToHpHmMuMu_noPID";

  TChain *sim_chain = new TChain(TString::Format("%s/DecayTree", sim_tree_name));
  sim_chain->Add("/auto/data/pfeiffer/inclusive_detached_dilepton/MC/rd_btoxll_simulation_fullstream_v0r0p6919194_B0ToKpPimMuMu_11144002_magdown.root");

  FourVect *l14v_sim = FourVect::Init(sim_chain, "L1");
  FourVect *l24v_sim = FourVect::Init(sim_chain, "L2");
  FourVect *hp4v_sim = FourVect::Init(sim_chain, "Hp");
  FourVect *hm4v_sim = FourVect::Init(sim_chain, "Hm");

  Int_t B_BKGCAT, L1_TRUEID, L2_TRUEID, Hp_TRUEID, Hm_TRUEID;

  sim_chain->SetBranchAddress("L1_TRUEID", &L1_TRUEID);
  sim_chain->SetBranchAddress("L2_TRUEID", &L2_TRUEID);
  sim_chain->SetBranchAddress("Hp_TRUEID", &Hp_TRUEID);
  sim_chain->SetBranchAddress("Hm_TRUEID", &Hm_TRUEID);
  sim_chain->SetBranchAddress("B0_BKGCAT", &B_BKGCAT);

  Float_t B0_PT_in, B0_BPVFDCHI2_in, B0_BPVDIRA_in, 
        Jpsi_BPVIPCHI2_in, Jpsi_PT_in, Jpsi_BPVDIRA_in, Jpsi_BPVFDCHI2_in,
        Res_BPVIPCHI2_in, Res_PT_in, Res_BPVFDCHI2_in,
        Hp_BPVIPCHI2_in, Hp_PT_in, 
        Hm_BPVIPCHI2_in, Hm_PT_in, 
        L1_BPVIPCHI2_in, L1_PT_in, 
        L2_BPVIPCHI2_in, L2_PT_in;
  Double_t B0_CHI2_in, B0_CHI2DOF_in, 
          Jpsi_CHI2_in, Jpsi_CHI2DOF_in,
          Hp_PROBNN_K_in, Hp_PID_K_in,
          Hm_PROBNN_K_in, Hm_PID_K_in,
          L1_PID_MU_in, 
          L2_PID_MU_in;

  sim_chain->SetBranchAddress("B0_PT", &B0_PT_in);
  sim_chain->SetBranchAddress("B0_BPVFDCHI2", &B0_BPVFDCHI2_in);
  sim_chain->SetBranchAddress("B0_BPVDIRA", &B0_BPVDIRA_in);
  sim_chain->SetBranchAddress("B0_CHI2", &B0_CHI2_in);
  sim_chain->SetBranchAddress("B0_CHI2VXNDOF", &B0_CHI2DOF_in);

  sim_chain->SetBranchAddress("Jpsi_BPVIPCHI2", &Jpsi_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("Jpsi_CHI2", &Jpsi_CHI2_in);
  sim_chain->SetBranchAddress("Jpsi_CHI2DOF", &Jpsi_CHI2DOF_in);
  sim_chain->SetBranchAddress("Jpsi_PT", &Jpsi_PT_in);
  sim_chain->SetBranchAddress("Jpsi_BPVFDCHI2", &Jpsi_BPVFDCHI2_in);

  sim_chain->SetBranchAddress("Res_BPVIPCHI2", &Res_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("Res_PT", &Res_PT_in);
  sim_chain->SetBranchAddress("Res_BPVFDCHI2", &Res_BPVFDCHI2_in);

  sim_chain->SetBranchAddress("Hp_BPVIPCHI2", &Hp_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("Hp_PT", &Hp_PT_in);
  sim_chain->SetBranchAddress("Hp_PROBNN_K", &Hp_PROBNN_K_in);
  sim_chain->SetBranchAddress("Hp_PID_K", &Hp_PID_K_in);

  sim_chain->SetBranchAddress("Hm_BPVIPCHI2", &Hm_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("Hm_PT", &Hm_PT_in);
  sim_chain->SetBranchAddress("Hm_PROBNN_K", &Hm_PROBNN_K_in);
  sim_chain->SetBranchAddress("Hm_PID_K", &Hm_PID_K_in);

  sim_chain->SetBranchAddress("L1_BPVIPCHI2", &L1_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("L1_PT", &L1_PT_in);
  sim_chain->SetBranchAddress("L1_PID_MU", &L1_PID_MU_in);

  sim_chain->SetBranchAddress("L2_BPVIPCHI2", &L2_BPVIPCHI2_in);
  sim_chain->SetBranchAddress("L2_PT", &L2_PT_in);
  sim_chain->SetBranchAddress("L2_PID_MU", &L2_PID_MU_in);

  TTree *output_tree = new TTree("DecayTree", "DecayTree");

  Float_t B0_PT_out, B0_BPVFDCHI2_out, B0_BPVDIRA_out,    
      Jpsi_BPVIPCHI2_out, Jpsi_PT_out, Jpsi_BPVDIRA_out, Jpsi_BPVFDCHI2_out,
      Kst0_BPVIPCHI2_out, Kst0_PT_out, Kst0_BPVFDCHI2_out,
      Kplus_BPVIPCHI2_out, Kplus_PT_out, 
      piminus_BPVIPCHI2_out, piminus_PT_out, muplus_BPVIPCHI2_out, muplus_PT_out,
      muminus_BPVIPCHI2_out, muminus_PT_out,
      B0_PX_out, B0_PY_out, B0_PZ_out, B0_ENERGY_out,
      muplus_PX_out, muplus_PY_out, muplus_PZ_out, muplus_ENERGY_out,
      muminus_PX_out, muminus_PY_out, muminus_PZ_out, muminus_ENERGY_out,
      Kplus_PX_out, Kplus_PY_out, Kplus_PZ_out, Kplus_ENERGY_out,
      piminus_PX_out, piminus_PY_out, piminus_PZ_out, piminus_ENERGY_out;

  Double_t B0_M_out, B0_CHI2_out, B0_CHI2DOF_out,
          Jpsi_CHI2_out, Jpsi_CHI2DOF_out,
          muplus_M_out, 
          muminus_M_out, 
          Kplus_M_out, Kplus_PROBNN_K_out, Kplus_PID_K_out,
          piminus_M_out, piminus_PROBNN_K_out;
  Int_t muplus_ID_out, muminus_ID_out, Kplus_ID_out, piminus_ID_out;

  output_tree->Branch("B0_PT", &B0_PT_out, "B0_PT/F");
  output_tree->Branch("B0_BPVFDCHI2", &B0_BPVFDCHI2_out, "B0_BPVFDCHI2/F");
  output_tree->Branch("B0_BPVDIRA", &B0_BPVDIRA_out, "B0_BPVDIRA/F");
  output_tree->Branch("B0_CHI2", &B0_CHI2_out, "B0_CHI2/D");
  output_tree->Branch("B0_CHI2DOF", &B0_CHI2DOF_out, "B0_CHI2DOF/D");
  output_tree->Branch("B0_PX", &B0_PX_out, "0_PX/F");
  output_tree->Branch("B0_PY", &B0_PY_out, "0_PY/F");
  output_tree->Branch("B0_PZ", &B0_PZ_out, "0_PZ/F");
  output_tree->Branch("B0_ENERGY", &B0_ENERGY_out, "B0_ENERGY/F");
  output_tree->Branch("B0_M", &B0_M_out, "B0_M/D");

  output_tree->Branch("Jpsi_BPVIPCHI2", &Jpsi_BPVIPCHI2_out, "Jpsi_BPVIPCHI2/F");
  output_tree->Branch("Jpsi_PT", &Jpsi_PT_out, "Jpsi_PT/F");
  output_tree->Branch("Jpsi_BPVDIRA", &Jpsi_BPVDIRA_out, "Jpsi_BPVDIRA/F");
  output_tree->Branch("Jpsi_CHI2", &Jpsi_CHI2_out, "Jpsi_CHI2/D");
  output_tree->Branch("Jpsi_CHI2DOF", &Jpsi_CHI2DOF_out, "Jpsi_CHI2DOF/D");
  output_tree->Branch("Jpsi_BPVFDCHI2", &Jpsi_BPVFDCHI2_out, "Jpsi_BPVFDCHI2/F");

  output_tree->Branch("Kst0_BPVIPCHI2", &Kst0_BPVIPCHI2_out, "Kst0_BPVIPCHI2/F");
  output_tree->Branch("Kst0_PT", &Kst0_PT_out, "Kst0_PT/F");
  output_tree->Branch("Kst0_BPVFDCHI2", &Kst0_BPVFDCHI2_out, "Kst0_BPVFDCHI2/F");

  output_tree->Branch("Kplus_BPVIPCHI2", &Kplus_BPVIPCHI2_out, "Kplus_BPVIPCHI2/F");
  output_tree->Branch("Kplus_PROBNN_K", &Kplus_PROBNN_K_out, "Kplus_PROBNN_K/D");
  output_tree->Branch("Kplus_PID_K", &Kplus_PID_K_out, "Kplus_PID_K/D");
  output_tree->Branch("Kplus_PT", &Kplus_PT_out, "Kplus_PT/F");
  output_tree->Branch("Kplus_PX", &Kplus_PX_out, "Kplus_PX/F");
  output_tree->Branch("Kplus_PY", &Kplus_PY_out, "Kplus_PY/F");
  output_tree->Branch("Kplus_PZ", &Kplus_PZ_out, "Kplus_PZ/F");
  output_tree->Branch("Kplus_ENERGY", &Kplus_ENERGY_out, "Kplus_ENERGY/F");
  output_tree->Branch("Kplus_M", &Kplus_M_out, "Kplus_M/D");
  output_tree->Branch("Kplus_ID", &Kplus_ID_out, "Kplus_ID/I");

  output_tree->Branch("piminus_BPVIPCHI2", &piminus_BPVIPCHI2_out, "piminus_BPVIPCHI2/F");
  output_tree->Branch("piminus_PROBNN_K", &piminus_PROBNN_K_out, "piminus_PROBNN_K/D");
  output_tree->Branch("piminus_PT", &piminus_PT_out, "piminus_PT/F");
  output_tree->Branch("piminus_PX", &piminus_PX_out, "piminus_PX/F");
  output_tree->Branch("piminus_PY", &piminus_PY_out, "piminus_PY/F");
  output_tree->Branch("piminus_PZ", &piminus_PZ_out, "piminus_PZ/F");
  output_tree->Branch("piminus_ENERGY", &piminus_ENERGY_out, "piminus_ENERGY/F");
  output_tree->Branch("piminus_M", &piminus_M_out, "piminus_M/D");
  output_tree->Branch("piminus_ID", &piminus_ID_out, "piminus_ID/I");

  output_tree->Branch("muminus_BPVIPCHI2", &muminus_BPVIPCHI2_out, "muminus_BPVIPCHI2/F");
  output_tree->Branch("muminus_PT", &muminus_PT_out, "muminus_PT/F");
  output_tree->Branch("muminus_PX", &muminus_PX_out, "muminus_PX/F");
  output_tree->Branch("muminus_PY", &muminus_PY_out, "muminus_PY/F");
  output_tree->Branch("muminus_PZ", &muminus_PZ_out, "muminus_PZ/F");
  output_tree->Branch("muminus_ENERGY", &muminus_ENERGY_out, "muminus_ENERGY/F");
  output_tree->Branch("muminus_M", &muminus_M_out, "muminus_M/D");
  output_tree->Branch("muminus_ID", &muminus_ID_out, "muminus_ID/I");

  output_tree->Branch("muplus_BPVIPCHI2", &muplus_BPVIPCHI2_out, "muplus_BPVIPCHI2/F");
  output_tree->Branch("muplus_PT", &muplus_PT_out, "muplus_PT/F");
  output_tree->Branch("muplus_PX", &muplus_PX_out, "muplus_PX/F");
  output_tree->Branch("muplus_PY", &muplus_PY_out, "muplus_PY/F");
  output_tree->Branch("muplus_PZ", &muplus_PZ_out, "muplus_PZ/F");
  output_tree->Branch("muplus_ENERGY", &muplus_ENERGY_out, "muplus_ENERGY/F");
  output_tree->Branch("muplus_M", &muplus_M_out, "muplus_M/D");
  output_tree->Branch("muplus_ID", &muplus_ID_out, "muplus_ID/I");

  unsigned int sim_entries = sim_chain->GetEntries();
  unsigned int selected_entries = 0;

  for (unsigned int i = 0; i < sim_entries; i++)
  {
    sim_chain->GetEntry(i);
    if (B_BKGCAT == 30 && TMath::Abs(L1_TRUEID) == PID_MUON && L2_TRUEID == -L1_TRUEID && ((TMath::Abs(Hp_TRUEID) == PID_KAON && TMath::Abs(Hm_TRUEID) == PID_PION) || (TMath::Abs(Hp_TRUEID) == PID_PION && TMath::Abs(Hm_TRUEID) == PID_KAON)) && L1_PID_MU_in > -3. && L2_PID_MU_in > -3.)
    {
      B0_PT_out = B0_PT_in;
      B0_BPVFDCHI2_out = B0_BPVFDCHI2_in;
      B0_BPVDIRA_out = B0_BPVDIRA_in;
      B0_CHI2_out = B0_CHI2_in;
      B0_CHI2DOF_out = B0_CHI2DOF_in;

      Jpsi_BPVIPCHI2_out = Jpsi_BPVIPCHI2_in;
      Jpsi_PT_out = Jpsi_PT_in;
      Jpsi_BPVDIRA_out = Jpsi_BPVDIRA_in;
      Jpsi_CHI2_out = Jpsi_CHI2_in;
      Jpsi_CHI2DOF_out = Jpsi_CHI2DOF_in;
      Jpsi_BPVFDCHI2_out = Jpsi_BPVFDCHI2_in;

      Kst0_BPVIPCHI2_out = Res_BPVIPCHI2_in;
      Kst0_BPVFDCHI2_out = Res_BPVFDCHI2_in;
      Kst0_PT_out = Res_PT_in;

      muminus_BPVIPCHI2_out = L1_BPVIPCHI2_in;
      muminus_ID_out = L1_TRUEID;
      muminus_PT_out = L1_PT_in;
      auto muminus_lv = l14v_sim->LorentzVector();
      muminus_PX_out = muminus_lv.Px();
      muminus_PY_out = muminus_lv.Py();
      muminus_PZ_out = muminus_lv.Pz();
      muminus_ENERGY_out = muminus_lv.E();
      muminus_M_out = muminus_lv.M();

      muplus_BPVIPCHI2_out = L2_BPVIPCHI2_in;
      muplus_ID_out = L2_TRUEID;
      muplus_PT_out = L2_PT_in;
      auto muplus_lv = l24v_sim->LorentzVector();
      muplus_PX_out = muplus_lv.Px();
      muplus_PY_out = muplus_lv.Py();
      muplus_PZ_out = muplus_lv.Pz();
      muplus_ENERGY_out = muplus_lv.E();
      muplus_M_out = muplus_lv.M();

      if (TMath::Abs(Hp_TRUEID) == PID_KAON && TMath::Abs(Hm_TRUEID) == PID_PION)
      {
        Kplus_BPVIPCHI2_out = Hp_BPVIPCHI2_in;
        Kplus_PT_out = Hp_PT_in;
        Kplus_PROBNN_K_out = Hp_PROBNN_K_in;
        Kplus_PID_K_out = Hp_PID_K_in;
        Kplus_ID_out = Hp_TRUEID;
        piminus_BPVIPCHI2_out = Hm_BPVIPCHI2_in;
        piminus_PT_out = Hm_PT_in;
        piminus_PROBNN_K_out = Hm_PROBNN_K_in;
        piminus_ID_out = Hm_TRUEID;

        auto Kplus_lv = hp4v_sim->LorentzVector(K_MASS);
        auto piminus_lv = hm4v_sim->LorentzVector();

        Kplus_PX_out = Kplus_lv.Px();
        Kplus_PY_out = Kplus_lv.Py();
        Kplus_PZ_out = Kplus_lv.Pz();
        Kplus_ENERGY_out = Kplus_lv.E();
        Kplus_M_out = Kplus_lv.M();

        piminus_PX_out = piminus_lv.Px();
        piminus_PY_out = piminus_lv.Py();
        piminus_PZ_out = piminus_lv.Pz();
        piminus_ENERGY_out = piminus_lv.E();
        piminus_M_out = piminus_lv.M();

        auto B0_lv = muminus_lv + muplus_lv + Kplus_lv + piminus_lv;

        B0_PX_out = B0_lv.Px();
        B0_PY_out = B0_lv.Py();
        B0_PZ_out = B0_lv.Pz();
        B0_ENERGY_out = B0_lv.E();
        B0_M_out = B0_lv.M();
      }
      else if (TMath::Abs(Hp_TRUEID) == PID_PION && TMath::Abs(Hm_TRUEID) == PID_KAON)
      {
        Kplus_BPVIPCHI2_out = Hm_BPVIPCHI2_in;
        Kplus_PT_out = Hm_PROBNN_K_in;
        Kplus_PROBNN_K_out = Hm_PROBNN_K_in;
        Kplus_PID_K_out = Hm_PID_K_in;
        Kplus_ID_out = Hm_TRUEID;
        piminus_BPVIPCHI2_out = Hp_BPVIPCHI2_in;
        piminus_PT_out = Hp_PT_in;
        piminus_PROBNN_K_out = Hp_PROBNN_K_in;
        piminus_ID_out = Hp_TRUEID;

        auto Kplus_lv = hm4v_sim->LorentzVector(K_MASS);
        auto piminus_lv = hp4v_sim->LorentzVector();

        Kplus_PX_out = Kplus_lv.Px();
        Kplus_PY_out = Kplus_lv.Py();
        Kplus_PZ_out = Kplus_lv.Pz();
        Kplus_ENERGY_out = Kplus_lv.E();
        Kplus_M_out = Kplus_lv.M();

        piminus_PX_out = piminus_lv.Px();
        piminus_PY_out = piminus_lv.Py();
        piminus_PZ_out = piminus_lv.Pz();
        piminus_ENERGY_out = piminus_lv.E();
        piminus_M_out = piminus_lv.M();

        auto B0_lv = muminus_lv + muplus_lv + Kplus_lv + piminus_lv;

        B0_PX_out = B0_lv.Px();
        B0_PY_out = B0_lv.Py();
        B0_PZ_out = B0_lv.Pz();
        B0_ENERGY_out = B0_lv.E();
        B0_M_out = B0_lv.M();
      }

      output_tree->Fill();
      selected_entries++;
    }

    PrintProgress(TString::Format("%s Mapping", sim_tree_name), sim_entries, 10000, i);
  }

  TFile* output_file = TFile::Open("/auto/data/pfeiffer/inclusive_detached_dilepton/MC/B0ToHpHmMuMu_mapped_mc.root", "RECREATE");
  output_file->cd();
  output_file->mkdir("B0ToHpHmMuMu_noPID_mapped/");
  output_file->cd("B0ToHpHmMuMu_noPID_mapped/");
  output_tree->SetDirectory(gDirectory);
  output_tree->Write();
  output_file->Close();

  std::cout << "Triggered Entries: " << sim_entries << " Selected Entries: " << selected_entries << std::endl;

  return 0;
}