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