LightFlavorRatios/truth_ratio/calculate_truth_ratio.C

0001 #include "../util/binning.h"
0002 #include <Math/Vector4D.h>
0003
0004 struct Hists
0005 {
0006   TH1F* h_pt;
0007   TH1F* h_eta;
0008   TH1F* h_phi;
0009   TH1F* h_y;
0010   Hists(std::string basename, std::string basetitle)
0011   {
0012     h_pt = makeHistogram(basename,basetitle,BinInfo::final_pt_bins);
0013     h_eta = makeHistogram(basename,basetitle,BinInfo::final_eta_bins);
0014     h_phi = makeHistogram(basename,basetitle,BinInfo::final_phi_bins);
0015     h_y = makeHistogram(basename,basetitle,BinInfo::final_rapidity_bins);
0016
0017     h_pt->Sumw2();
0018     h_eta->Sumw2();
0019     h_phi->Sumw2();
0020     h_y->Sumw2();
0021   }
0022   void divide(Hists& numerator, Hists& denominator, float scale_factor=1.)
0023   {
0024     h_pt->Divide(numerator.h_pt,denominator.h_pt,scale_factor);
0025     h_eta->Divide(numerator.h_eta,denominator.h_eta,scale_factor);
0026     h_phi->Divide(numerator.h_phi,denominator.h_phi,scale_factor);
0027     h_y->Divide(numerator.h_y,denominator.h_y,scale_factor);
0028   }
0029   void write()
0030   {
0031     h_pt->Write();
0032     h_eta->Write();
0033     h_phi->Write();
0034     h_y->Write();
0035   }
0036 };
0037
0038 std::vector<Hists> process_tree(TTree* t, std::string basename, std::string basetitle, int mother_flavor, std::vector<int> required_daughter_flavors, bool parity_inclusive)
0039 {
0040   const float mother_mass = TDatabasePDG::Instance()->GetParticle(mother_flavor)->Mass();
0041
0042   std::vector<Hists> vh;
0043   vh.emplace_back(basename,basetitle);
0044   vh.emplace_back(basename+"_wd",basetitle+" with daughters");
0045
0046   float fevent;
0047   float gflavor;
0048   float gtrackID;
0049   float rtrackID;
0050   float gparentflavor;
0051   float gparentID;
0052   float gvz;
0053   float gpt;
0054   float gpz;
0055   float geta;
0056   float gphi;
0057   float gprimary;
0058   float gprimaryid;
0059
0060   t->SetBranchAddress("event",&fevent);
0061   t->SetBranchAddress("gflavor",&gflavor);
0062   t->SetBranchAddress("gtrackID",&gtrackID);
0063   t->SetBranchAddress("trackID",&rtrackID);
0064   t->SetBranchAddress("gparentflavor",&gparentflavor);
0065   t->SetBranchAddress("gparentid",&gparentID);
0066   t->SetBranchAddress("gvz",&gvz);
0067   t->SetBranchAddress("gpt",&gpt);
0068   t->SetBranchAddress("gpz",&gpz);
0069   t->SetBranchAddress("geta",&geta);
0070   t->SetBranchAddress("gphi",&gphi);
0071   t->SetBranchAddress("gprimary",&gprimary);
0072   t->SetBranchAddress("gprimaryid",&gprimaryid);
0073
0074   size_t current_event = 0;
0075   size_t current_event_start_entry = 0;
0076   std::map<int,std::pair<ROOT::Math::PtEtaPhiMVector,std::set<int>>> current_daughter_map;
0077   // key: parent ID, value: (mother Lorentz vector, set of daughter flavors)
0078
0079   for(size_t i=0; i<t->GetEntries(); i++)
0080   {
0081     t->GetEntry(i);
0082     int event = round(fevent);
0083
0084     if(event != current_event)
0085     {
0086       if(event % 100 == 0) std::cout << "event " << current_event << std::endl;
0087       // gather all reconstructible daughters for each mother
0088       for(size_t j=current_event_start_entry; j<i; j++)
0089       {
0090         t->GetEntry(j);
0091         // don't round reco trackID, this gets rid of NaN-ness on conversion to int
0092         int parentID = round(gparentID);
0093         //std::cout << "parentID = " << parentID << std::endl;
0094         int flavor = round(gflavor);
0095         //std::cout << "flavor = " << flavor << std::endl;
0096         //std::cout << "trackID = " << rtrackID << std::endl;
0097         //std::cout << "map count = " << current_daughter_map.count(parentID) << std::endl;
0098         if(!std::isnan(rtrackID) && current_daughter_map.count(parentID)>0)
0099         {
0100           current_daughter_map[parentID].second.insert(flavor);
0101         }
0102       }
0103       //std::cout << "----------------------------------------------" << std::endl;
0104       // check list of mothers for correct set of daughter flavors
0105       for(auto [motherID, sv_info] : current_daughter_map)
0106       {
0107         ROOT::Math::PtEtaPhiMVector mother_lorentzvector = sv_info.first;
0108         std::set<int> daughter_flavors = sv_info.second;
0109
0110         //std::cout << "mother iD = " << motherID << std::endl;
0111         //std::cout << "daughters: ";
0112         //for(int df : daughter_flavors) std::cout << df << " ";
0113         //std::cout << std::endl;
0114         bool has_all_daughters = std::all_of(required_daughter_flavors.begin(),required_daughter_flavors.end(),[&](int flavor){ return daughter_flavors.contains(flavor); });
0115         bool has_all_opposite_daughters = std::all_of(required_daughter_flavors.begin(),required_daughter_flavors.end(),[&](int flavor){ return daughter_flavors.contains(-1*flavor); });
0116         //if(has_all_daughters) std::cout << "has all daughters" << std::endl;
0117         //if(has_all_opposite_daughters) std::cout << "has all opposite daughters" << std::endl;
0118         if(has_all_daughters || (parity_inclusive && has_all_opposite_daughters))
0119         {
0120           vh[1].h_pt->Fill(mother_lorentzvector.Pt());
0121           vh[1].h_eta->Fill(mother_lorentzvector.Eta());
0122           vh[1].h_phi->Fill(mother_lorentzvector.Phi());
0123           vh[1].h_y->Fill(mother_lorentzvector.Rapidity());
0124         }
0125       }
0126       // clean up
0127       current_daughter_map.clear();
0128       current_event = event;
0129       current_event_start_entry = i;
0130       t->GetEntry(i);
0131     }
0132
0133     // check if this is a valid mother
0134     int primary = round(gprimary);
0135     int flavor = round(gflavor);
0136     int parentflavor = round(gparentflavor);
0137     int trackID = round(gtrackID);
0138     int parentID = round(gparentID);
0139     int primaryID = round(gprimaryid);
0140
0141     //std::cout << "gtrackID = " << gtrackID << " flavor = " << flavor << " parentID = " << parentID << std::endl;
0142     bool correct_flavor = (flavor == mother_flavor || (parity_inclusive && abs(flavor)==mother_flavor));
0143     bool is_duplicate = (parentflavor == flavor && parentID != trackID);
0144     bool is_primary = (primary==1);
0145     // handle K0->Ks oscillation
0146     // sometimes they're primary, other times they get "decayed" from a primary K0
0147     if(flavor==310) is_primary = (primary==1) || (abs(parentflavor)==311 && parentID == primaryID);
0148     //if(correct_flavor) std::cout << "correct flavor" << std::endl;
0149     //if(is_duplicate) std::cout << "is duplicate, parentID = " << parentID << ", parentflavor = " << parentflavor << std::endl;
0150     if(is_primary && correct_flavor)// && !is_duplicate)
0151     {
0152       ROOT::Math::PtEtaPhiMVector mother_lorentzvector(gpt, geta, gphi, mother_mass);
0153       float rapidity = mother_lorentzvector.Rapidity();
0154
0155       // add to mother histograms and create new daughter map entry
0156       if(gpt>=0.6 && gpt<=4. && fabs(geta)<=0.8 && fabs(rapidity)<=0.8 && fabs(gvz)<=10.)
0157       {
0158         vh[0].h_pt->Fill(gpt);
0159         vh[0].h_eta->Fill(geta);
0160         vh[0].h_phi->Fill(gphi);
0161         vh[0].h_y->Fill(rapidity);
0162         current_daughter_map.insert({trackID,std::make_pair(mother_lorentzvector,std::set<int>())});
0163       }
0164     }
0165   }
0166
0167   return vh;
0168 }
0169
0170 void calculate_truth_ratio(int numerator_flavor, std::vector<int> numerator_daughter_flavors,
0171                            int denominator_flavor, std::vector<int> denominator_daughter_flavors,
0172                            std::string numerator_particlename, std::string denominator_particlename,
0173                            std::string numerator_basefile, std::string denominator_basefile,
0174                            float scale_factor, std::string outfilebase, int process, bool numerator_abs = true, bool denominator_abs = true)
0175 {
0176   size_t ndigits = 6;
0177   std::string process_str = std::string(ndigits - std::to_string(process).length(),'0') + std::to_string(process);
0178
0179   TFile* f_num = TFile::Open((numerator_basefile+"_"+process_str+".root").c_str());
0180   TFile* f_denom = TFile::Open((denominator_basefile+"_"+process_str+".root").c_str());
0181
0182   TTree* t_num = (TTree*)f_num->Get("ntp_gtrack");
0183   TTree* t_denom = (TTree*)f_denom->Get("ntp_gtrack");
0184
0185   std::string outfname = outfilebase+"_"+process_str+".root";
0186   TFile* fout = new TFile(outfname.c_str(),"RECREATE");
0187
0188   std::string title_r = std::to_string(scale_factor) + numerator_particlename + "/" + denominator_particlename;
0189
0190   Hists hists_r("hr",title_r);
0191
0192   Hists eff_n("eff_n",numerator_particlename + " reconstruction efficiency");
0193   Hists eff_d("eff_d",denominator_particlename + " reconstruction efficiency");
0194
0195   std::vector<Hists> vhists_n = process_tree(t_num,"hn",numerator_particlename+" yield",numerator_flavor,numerator_daughter_flavors,numerator_abs);
0196   std::vector<Hists> vhists_d = process_tree(t_denom,"hd",denominator_particlename+" yield",denominator_flavor,denominator_daughter_flavors,denominator_abs);
0197
0198   hists_r.divide(vhists_n[0],vhists_d[0],scale_factor);
0199   eff_n.divide(vhists_n[1],vhists_n[0],scale_factor);
0200   eff_d.divide(vhists_d[1],vhists_d[0],scale_factor);
0201
0202   hists_r.h_pt->Divide(vhists_n[0].h_pt,vhists_d[0].h_pt,scale_factor);
0203   hists_r.h_eta->Divide(vhists_n[0].h_eta,vhists_d[0].h_eta,scale_factor);
0204   hists_r.h_phi->Divide(vhists_n[0].h_phi,vhists_d[0].h_phi,scale_factor);
0205   hists_r.h_y->Divide(vhists_n[0].h_y,vhists_d[0].h_y,scale_factor);
0206
0207   for(Hists& h : vhists_n) h.write();
0208   for(Hists& h : vhists_d) h.write();
0209   hists_r.write();
0210   eff_n.write();
0211   eff_d.write();
0212
0213 }