LightFlavorRatios/swimming_correction/pythia_hepmc_origin.C

0001 struct TrackData
0002 {
0003   int gtrackID;
0004   int gflavor;
0005   std::array<float,3> gvertex;
0006   std::array<float,3> gp;
0007 };
0008
0009 void pythia_origin(int process=1)
0010 {
0011
0012   std::string hepmc_file;
0013
0014   HepMC3::ReaderAscii reader(hepmc_file);
0015   HepMC3::GenEvent event;
0016
0017   const int mother_pdgid = 3122;
0018 /*
0019   TChain* f = new TChain("ntp_gtrack");
0020
0021   std::string file_base = "/sphenix/tg/tg01/hf/frawley/sims/pythia_fromhits_detroit/pythia_reco_3122_";
0022
0023   if(process==0)
0024   {
0025     f->Add((file_base+"[0-9]_g4svtx_eval.root").c_str());
0026     f->Add((file_base+"[0-9][0-9]_g4svtx_eval.root").c_str());
0027   }
0028   else
0029   {
0030     f->Add((file_base+std::to_string(process)+"[0-9][0-9]_g4svtx_eval.root").c_str());
0031   }
0032
0033   //f->Add("/sphenix/tg/tg01/hf/frawley/sims/pythia_fromhits_detroit/pythia_reco_3122_[0-9]_g4svtx_eval.root");
0034   //f->Add("/sphenix/tg/tg01/hf/frawley/sims/pythia_fromhits_detroit_volcut/pythia_reco_310_[0-9]_g4svtx_eval.root");
0035   //f->Add("/sphenix/tg/tg01/hf/cdean/LF_analysis/evaluatorFiles/Lambda2ppi/20260504/evalFile_Lambda2ppi_000[0-9][0-9].root");
0036   //f->Add("/sphenix/tg/tg01/hf/cdean/LF_analysis/evaluatorFiles/Kshort2pipi/20260504/evalFile_Kshort2pipi_000[0-9][0-9].root");
0037
0038 //  const int mother_pdgid = 310;
0039   const int mother_pdgid = 3122;
0040
0041   float f_event;
0042   float f_gtrackID;
0043   float f_gflavor;
0044   float f_gparentflavor;
0045   float f_gparentid;
0046   float f_gprimary;
0047   float gvx;
0048   float gvy;
0049   float gvz;
0050   float gpx;
0051   float gpy;
0052   float gpz;
0053
0054   f->SetBranchAddress("event",&f_event);
0055   f->SetBranchAddress("gtrackID",&f_gtrackID);
0056   f->SetBranchAddress("gflavor",&f_gflavor);
0057   f->SetBranchAddress("gparentflavor",&f_gparentflavor);
0058   f->SetBranchAddress("gparentid",&f_gparentid);
0059   f->SetBranchAddress("gprimary",&f_gprimary);
0060   f->SetBranchAddress("gvx",&gvx);
0061   f->SetBranchAddress("gvy",&gvy);
0062   f->SetBranchAddress("gvz",&gvz);
0063   f->SetBranchAddress("gpx",&gpx);
0064   f->SetBranchAddress("gpy",&gpy);
0065   f->SetBranchAddress("gpz",&gpz);
0066
0067   std::map<int,int> parent_flavor_count;
0068   std::multimap<int,int> thisevent_nextID_targets;
0069   std::map<int,std::vector<TrackData>> thisevent_history;
0070
0071   size_t current_file=0;
0072   size_t current_event=0;
0073   size_t current_event_start=0;
0074 */
0075   size_t total_parents=0;
0076
0077   int file_out;
0078   int event_out;
0079   std::vector<int> gflavor_out;
0080   std::vector<int> gtrackID_out;
0081   std::vector<float> gvx_out;
0082   std::vector<float> gvy_out;
0083   std::vector<float> gvz_out;
0084   std::vector<float> gpx_out;
0085   std::vector<float> gpy_out;
0086   std::vector<float> gpz_out;
0087
0088   TFile* fout = new TFile(("/sphenix/user/mjpeters/analysis/LightFlavorRatios/swimming_correction/output/histout_"+std::to_string(process)+".root").c_str(),"RECREATE");
0089   TTree* tout = new TTree("history","Candidate history");
0090   tout->Branch("file",&file_out);
0091   tout->Branch("event",&event_out);
0092   tout->Branch("gtrackID_history",&gtrackID_out);
0093   tout->Branch("gflavor_history",&gflavor_out);
0094   tout->Branch("gvx_history",&gvx_out);
0095   tout->Branch("gvy_history",&gvy_out);
0096   tout->Branch("gvz_history",&gvz_out);
0097   tout->Branch("gpx_history",&gpx_out);
0098   tout->Branch("gpy_history",&gpy_out);
0099   tout->Branch("gpz_history",&gpz_out);
0100
0101   int event = 0;
0102
0103   while(!reader.failed())
0104   {
0105     reader.read_event(event);
0106     if(reader.failed())
0107     {
0108       break;
0109     }
0110
0111     for(const auto& p : event.particles())
0112     {
0113       if(std::abs(p->pid()) == mother_pdgid)
0114       {
0115         if(parent_flavor_count.contains(gparentflavor))
0116         {
0117           parent_flavor_count[gparentflavor] += 1;
0118         }
0119         else
0120         {
0121           parent_flavor_count[gparentflavor] = 1;
0122         }
0123         thisevent_history[gtrackID].push_back({.gtrackID = gtrackID, .gflavor = gflavor, .gvertex = {gvx,gvy,gvz}, .gp = {gpx,gpy,gpz}});
0124         //std::cout << "gtrackID " << gtrackID << " gprimary " << gprimary << " gparentid " << gparentid << std::endl;
0125         if(!gprimary)
0126         {
0127           thisevent_nextID_targets.insert({gparentid,gtrackID});
0128         }
0129         total_parents++;
0130       }
0131     }
0132
0133     event++;
0134   }
0135
0136   for(size_t i=0;i<f->GetEntries();i++)
0137   {
0138     //if(i % 10000 == 0) std::cout << "entry " << i << std::endl;
0139     f->GetEntry(i);
0140
0141     int event = round(f_event);
0142     int gtrackID = round(f_gtrackID);
0143     int gflavor = round(f_gflavor);
0144     int gparentflavor = round(f_gparentflavor);
0145     int gparentid = round(f_gparentid);
0146     bool gprimary = ((int)round(f_gprimary) == 1);
0147
0148     if(event!=current_event)
0149     {
0150       //std::cout << "event " << event << std::endl;
0151       while(thisevent_nextID_targets.size()!=0)
0152       {
0153         std::multimap<int,int> new_thisevent_nextID_targets;
0154         for(size_t j=current_event_start;j<i;j++)
0155         {
0156           f->GetEntry(j);
0157
0158           int event_j = round(f_event);
0159           int gtrackID_j = round(f_gtrackID);
0160           int gflavor_j = round(f_gflavor);
0161           int gparentflavor_j = round(f_gparentflavor);
0162           int gparentid_j = round(f_gparentid);
0163           bool gprimary_j = ((int)round(f_gprimary) == 1);
0164
0165           auto target_range = thisevent_nextID_targets.equal_range(gtrackID_j);
0166
0167           if(target_range.first!=target_range.second)
0168           {
0169             for(auto it = target_range.first; it!=target_range.second; ++it)
0170             {
0171               int target_gtrackID = (*it).second;
0172               thisevent_history[target_gtrackID].push_back({.gtrackID = gtrackID_j, .gflavor = gflavor_j, .gvertex = {gvx,gvy,gvz}, .gp = {gpx,gpy,gpz}});
0173               if((gprimary_j && (gparentid_j!=gtrackID_j)) || (!gprimary_j && (gparentid_j==gtrackID_j)))
0174               {
0175                 std::cout << "WARNING: gprimary condition doesn't make sense!!" << std::endl;
0176                 std::cout << "gprimary " << gprimary_j << " gparentid " << gparentid_j << " gtrackID " << gtrackID_j << std::endl;
0177               }
0178               if(!gprimary_j)
0179               {
0180                 new_thisevent_nextID_targets.insert({gparentid_j,target_gtrackID});
0181               }
0182             }
0183           }
0184         }
0185         thisevent_nextID_targets = new_thisevent_nextID_targets;
0186       }
0187 /*
0188       if(thisevent_history.size()>0)
0189       {
0190         std::cout << "event " << current_event << " history:" << std::endl;
0191         for(auto const& [trackid, history] : thisevent_history)
0192         {
0193           std::cout << "track ID " << trackid << ": ";
0194           for(const TrackData& h : history)
0195           {
0196             std::cout << h.gflavor << " (" << h.gtrackID << ") <- ";
0197           }
0198           std::cout << std::endl;
0199         }
0200       }
0201 */
0202       file_out = current_file;
0203       event_out = current_event;
0204
0205       for(auto const& [id, history] : thisevent_history)
0206       {
0207         for(const TrackData& h : history)
0208         {
0209           gflavor_out.push_back(h.gflavor);
0210           gtrackID_out.push_back(h.gtrackID);
0211           gvx_out.push_back(h.gvertex[0]);
0212           gvy_out.push_back(h.gvertex[1]);
0213           gvz_out.push_back(h.gvertex[2]);
0214           gpx_out.push_back(h.gp[0]);
0215           gpy_out.push_back(h.gp[1]);
0216           gpz_out.push_back(h.gp[2]);
0217         }
0218
0219         tout->Fill();
0220
0221         gflavor_out.clear();
0222         gtrackID_out.clear();
0223         gvx_out.clear();
0224         gvy_out.clear();
0225         gvz_out.clear();
0226         gpx_out.clear();
0227         gpy_out.clear();
0228         gpz_out.clear();
0229       }
0230
0231       current_event = event;
0232       current_event_start = i;
0233
0234       thisevent_nextID_targets.clear();
0235       thisevent_history.clear();
0236
0237       f->GetEntry(i);
0238     }
0239
0240     if(event==0 && i==current_event_start)
0241     {
0242       std::cout << "file " << current_file << std::endl;
0243       current_file++;
0244     }
0245
0246     if(fabs(gflavor)==mother_pdgid)
0247     {
0248       if(parent_flavor_count.contains(gparentflavor))
0249       {
0250         parent_flavor_count[gparentflavor] += 1;
0251       }
0252       else
0253       {
0254         parent_flavor_count[gparentflavor] = 1;
0255       }
0256       thisevent_history[gtrackID].push_back({.gtrackID = gtrackID, .gflavor = gflavor, .gvertex = {gvx,gvy,gvz}, .gp = {gpx,gpy,gpz}});
0257       //std::cout << "gtrackID " << gtrackID << " gprimary " << gprimary << " gparentid " << gparentid << std::endl;
0258       if(!gprimary)
0259       {
0260         thisevent_nextID_targets.insert({gparentid,gtrackID});
0261       }
0262       total_parents++;
0263     }
0264   }
0265
0266   std::cout << "origin counts:" << std::endl;
0267   for(auto const& [pdgid, ct] : parent_flavor_count)
0268   {
0269     std::cout << pdgid << ": " << ct << " / " << total_parents << " (" << (float)ct/total_parents*100. << "%)" << std::endl;
0270   }
0271
0272   fout->Write();
0273
0274 }