Back to home page

sPhenix code displayed by LXR
 
 

    

File indexing completed on 2026-07-16 08:11:28 

0001 #include <limits>
0002 #include <vector>
0003 #include <any>
0004 #include <cmath>
0005 
0006 using namespace std;
0007 using namespace ROOT::Math;
0008 
0009 struct track_params
0010 {
0011   float event_ID  = numeric_limits<float>::quiet_NaN();
0012   float track_ID  = numeric_limits<float>::quiet_NaN();
0013   float parent_ID = numeric_limits<float>::quiet_NaN();
0014   float PID       = numeric_limits<float>::quiet_NaN();
0015   float true_pT   = numeric_limits<float>::quiet_NaN();
0016   float true_eta  = numeric_limits<float>::quiet_NaN();
0017   float true_phi  = numeric_limits<float>::quiet_NaN();
0018   float reco_pT   = numeric_limits<float>::quiet_NaN();
0019   float reco_eta  = numeric_limits<float>::quiet_NaN();
0020   float reco_phi  = numeric_limits<float>::quiet_NaN();
0021   float reco_maps = numeric_limits<float>::quiet_NaN();
0022   float true_vx   = numeric_limits<float>::quiet_NaN();
0023   float true_vy   = numeric_limits<float>::quiet_NaN();
0024   float true_vz   = numeric_limits<float>::quiet_NaN();
0025 };
0026 
0027 struct mother_params
0028 {
0029   track_params mother;
0030   track_params track_1;
0031   track_params track_2;
0032   bool hasTrack1 = false;
0033   bool hasTrack2 = false;
0034 };
0035 
0036 
0037 void parseEvaluators(string species = "Kshort", int nFiles = 2, int startNumber = 0)
0038 {
0039   int endNumber = startNumber + nFiles - 1;
0040   int mother_PDG = species == "Kshort" ? 310 : 3122;
0041   int track_1_PDG = species == "Kshort" ? 211 : 2212;
0042   int track_2_PDG = species == "Kshort" ? -211 : 211;
0043 
0044   string header = "/sphenix/tg/tg01/hf/frawley/sims/pythia_fromhits_detroit_volcut/pythia_reco_" + to_string(mother_PDG) + "_";
0045   string trailer = "_g4svtx_eval.root";
0046 
0047   string gtrack_string = "ntp_gtrack";
0048 
0049   vector<mother_params> accepted_mothers;
0050 
0051   int tmp = 0;
0052   int barWidth = 50;
0053 
0054   for (int i = 0; i < nFiles; ++i)
0055   {
0056     if (tmp != (int)100*i/nFiles)
0057     {
0058       tmp = (int)100*i/nFiles;
0059       if ((tmp%1)  == 0)
0060       {
0061         cout << "[";
0062         int pos = barWidth * tmp/100;
0063         for (int k = 0; k < barWidth; ++k)
0064         {
0065           if (k < pos) cout << "=";
0066           else if (k == pos) cout << ">";
0067           else cout << " ";
0068         }
0069         cout << "] " << tmp << " %\r";
0070         cout.flush();
0071       }
0072     } 
0073     vector<track_params> all_track_2s, all_track_1s;
0074     vector<mother_params> mothers;  
0075  
0076     string inString = header + to_string(i + startNumber) + trailer;
0077     TFile* inFile = new TFile(inString.c_str());
0078 
0079     TTree* gtrackTree = (TTree*)inFile->Get(gtrack_string.c_str());
0080     float event;         gtrackTree->SetBranchAddress("event",         &event);
0081     float gprimary;      gtrackTree->SetBranchAddress("gprimary",      &gprimary);
0082     float gflavor;       gtrackTree->SetBranchAddress("gflavor",       &gflavor);
0083     float gtrackID;      gtrackTree->SetBranchAddress("gtrackID",      &gtrackID);
0084     float gparentflavor; gtrackTree->SetBranchAddress("gparentflavor", &gparentflavor);
0085     float gparentid;     gtrackTree->SetBranchAddress("gparentid",     &gparentid);
0086     float gnmaps;        gtrackTree->SetBranchAddress("gnmaps",        &gnmaps);
0087     float gpt;           gtrackTree->SetBranchAddress("gpt",           &gpt);
0088     float geta;          gtrackTree->SetBranchAddress("geta",          &geta);
0089     float gphi;          gtrackTree->SetBranchAddress("gphi",          &gphi);
0090     float pt;            gtrackTree->SetBranchAddress("pt",            &pt);
0091     float eta;           gtrackTree->SetBranchAddress("eta",           &eta);
0092     float phi;           gtrackTree->SetBranchAddress("phi",           &phi);
0093     float nmaps;         gtrackTree->SetBranchAddress("nmaps",         &nmaps);
0094     float gvx;           gtrackTree->SetBranchAddress("gvx",           &gvx);
0095     float gvy;           gtrackTree->SetBranchAddress("gvy",           &gvy);
0096     float gvz;           gtrackTree->SetBranchAddress("gvz",           &gvz);
0097 
0098     unsigned int gtrack_entries = gtrackTree->GetEntries();
0099 
0100     for (unsigned int j = 0; j < gtrack_entries; ++j)
0101     {
0102       gtrackTree->GetEntry(j);
0103 
0104       bool isMother = abs(gflavor) == mother_PDG;
0105       bool isTrack1 = species == "Kshort" ? gflavor == track_1_PDG : abs(gflavor) == track_1_PDG;
0106       bool isTrack2 = species == "Kshort" ? gflavor == track_2_PDG : abs(gflavor) == track_2_PDG;
0107 
0108       if (isMother || isTrack1 || isTrack2) //Only care about the mother, track_2s and track_1s
0109       {
0110         if (isMother && (gprimary != 1)) continue; //Skip particles that are the mother but are NOT prompt
0111 
0112         if (isMother)
0113         {
0114           mother_params particle;
0115           particle.mother.event_ID = event;
0116           particle.mother.track_ID = gtrackID;
0117           particle.mother.PID = gflavor;
0118           particle.mother.true_pT = gpt;
0119           particle.mother.true_eta = geta;
0120           particle.mother.true_phi = gphi;
0121           particle.mother.true_vx = gvx;
0122           particle.mother.true_vy = gvy;
0123           particle.mother.true_vz = gvz;
0124           mothers.push_back(particle);
0125           continue;
0126         }
0127 
0128         track_params particle;
0129         particle.event_ID = event;
0130         particle.track_ID = gtrackID;
0131         particle.parent_ID = gparentid;
0132         particle.PID = gflavor;
0133         particle.true_pT = gpt;
0134         particle.true_eta = geta;
0135         particle.true_phi = gphi;
0136         particle.true_vx = gvx;
0137         particle.true_vy = gvy;
0138         particle.true_vz = gvz;
0139         particle.reco_maps = nmaps;
0140         if (nmaps > 0)
0141         {
0142           particle.reco_pT = pt;
0143           particle.reco_eta = eta;
0144           particle.reco_phi = phi;
0145         }
0146 
0147         if (isTrack1) all_track_1s.push_back(particle);
0148         else all_track_2s.push_back(particle);
0149       }
0150     }
0151 
0152     //OK, we have all mothers, track_1s and track_2s. Let's match!
0153     pair<string, vector<track_params>> track_pairs[2];
0154     track_pairs[0] = {"track_1", all_track_1s};
0155     track_pairs[1] = {"track_2", all_track_2s};
0156     for (auto &data : track_pairs)
0157     {
0158       for (auto &track : data.second)
0159       {
0160         auto it = std::find_if(mothers.begin(), mothers.end(), [&](const mother_params& candidate)
0161         {
0162             return candidate.mother.event_ID == track.event_ID && candidate.mother.track_ID == track.parent_ID;
0163         });
0164 
0165         if (it != mothers.end())
0166         {
0167           if (data.first == "track_1")
0168           {
0169             it->hasTrack1 = true;
0170             it->track_1 = track;
0171           }
0172           else
0173           {
0174             it->hasTrack2 = true;
0175             it->track_2 = track;
0176           }
0177         }
0178       }
0179     }
0180 
0181     inFile->Close();
0182 
0183     for (auto &mother : mothers)
0184     {
0185       if (mother.hasTrack1 && mother.hasTrack2) accepted_mothers.push_back(mother);
0186     }
0187   }
0188 
0189   cout << "[";
0190   int pos = barWidth * tmp;
0191   for (int i = 0; i < barWidth; ++i)
0192   {
0193     if (i < pos) cout << "=";
0194     else if (i == pos) cout << ">";
0195     else cout << " ";
0196   }
0197   cout << "] 100 %\r";
0198   cout.flush();
0199   cout<<endl;
0200 
0201   string outString = "/sphenix/user/cdean/private/files/LightFlavor/evaluator_skimmed_files/prompt_" + species + "_" + to_string(startNumber) + "_to_" + to_string(endNumber) + ".root";
0202   TFile* outFile = new TFile(outString.c_str(), "RECREATE");
0203   TTree* outTree = new TTree("TruthInfo", "TruthInfo");
0204   outTree->OptimizeBaskets();
0205   outTree->SetAutoSave(-5e6);  // Save the output file every 5MB
0206 
0207   float mother_PID;          outTree->Branch("mother_PID",         &mother_PID,         "mother_PID/F");
0208   float mother_true_vx;      outTree->Branch("mother_true_vx",     &mother_true_vx,     "mother_true_vx/F");
0209   float mother_true_vy;      outTree->Branch("mother_true_vy",     &mother_true_vy,     "mother_true_vy/F");
0210   float mother_true_vz;      outTree->Branch("mother_true_vz",     &mother_true_vz,     "mother_true_vz/F");
0211   float mother_true_mass;    outTree->Branch("mother_true_mass",   &mother_true_mass,   "mother_true_mass/F");
0212   float mother_reco_mass;    outTree->Branch("mother_reco_mass",   &mother_reco_mass,   "mother_reco_mass/F");
0213   float mother_true_pT;      outTree->Branch("mother_true_pT",     &mother_true_pT,     "mother_true_pT/F");
0214   float mother_reco_pT;      outTree->Branch("mother_reco_pT",     &mother_reco_pT,     "mother_reco_pT/F");
0215   float mother_true_eta;     outTree->Branch("mother_true_eta",    &mother_true_eta,    "mother_true_eta/F");
0216   float mother_true_phi;     outTree->Branch("mother_true_phi",    &mother_true_phi,    "mother_true_phi/F");
0217   float mother_true_rap;     outTree->Branch("mother_true_rap",    &mother_true_rap,    "mother_true_rap/F");
0218   float track_1_PID;         outTree->Branch("track_1_PID",        &track_1_PID,        "track_1_PID/F");
0219   float track_1_true_vx;     outTree->Branch("track_1_true_vx",    &track_1_true_vx,    "track_1_true_vx/F");
0220   float track_1_true_vy;     outTree->Branch("track_1_true_vy",    &track_1_true_vy,    "track_1_true_vy/F");
0221   float track_1_true_vz;     outTree->Branch("track_1_true_vz",    &track_1_true_vz,    "track_1_true_vz/F");
0222   float track_1_true_pT;     outTree->Branch("track_1_true_pT",    &track_1_true_pT,    "track_1_true_pT/F");
0223   float track_1_true_eta;    outTree->Branch("track_1_true_eta",   &track_1_true_eta,   "track_1_true_eta/F");
0224   float track_1_true_phi;    outTree->Branch("track_1_true_phi",   &track_1_true_phi,   "track_1_true_phi/F");
0225   float track_1_reco_pT;     outTree->Branch("track_1_reco_pT",    &track_1_reco_pT,    "track_1_reco_pT/F");
0226   float track_1_reco_eta;    outTree->Branch("track_1_reco_eta",   &track_1_reco_eta,   "track_1_reco_eta/F");
0227   float track_1_reco_phi;    outTree->Branch("track_1_reco_phi",   &track_1_reco_phi,   "track_1_reco_phi/F");
0228   float track_1_reco_nmaps;  outTree->Branch("track_1_reco_nmaps", &track_1_reco_nmaps, "track_1_reco_nmaps/F");
0229   float track_2_PID;         outTree->Branch("track_2_PID",        &track_2_PID,        "track_2_PID/F");
0230   float track_2_true_vx;     outTree->Branch("track_2_true_vx",    &track_2_true_vx,    "track_2_true_vx/F");
0231   float track_2_true_vy;     outTree->Branch("track_2_true_vy",    &track_2_true_vy,    "track_2_true_vy/F");
0232   float track_2_true_vz;     outTree->Branch("track_2_true_vz",    &track_2_true_vz,    "track_2_true_vz/F");
0233   float track_2_true_pT;     outTree->Branch("track_2_true_pT",    &track_2_true_pT,    "track_2_true_pT/F");
0234   float track_2_true_eta;    outTree->Branch("track_2_true_eta",   &track_2_true_eta,   "track_2_true_eta/F");
0235   float track_2_true_phi;    outTree->Branch("track_2_true_phi",   &track_2_true_phi,   "track_2_true_phi/F");
0236   float track_2_reco_pT;     outTree->Branch("track_2_reco_pT",    &track_2_reco_pT,    "track_2_reco_pT/F");
0237   float track_2_reco_eta;    outTree->Branch("track_2_reco_eta",   &track_2_reco_eta,   "track_2_reco_eta/F");
0238   float track_2_reco_phi;    outTree->Branch("track_2_reco_phi",   &track_2_reco_phi,   "track_2_reco_phi/F");
0239   float track_2_reco_nmaps;  outTree->Branch("track_2_reco_nmaps", &track_2_reco_nmaps, "track_2_reco_nmaps/F");
0240 
0241   cout << "Writing accepted mothers to TTree" << endl;
0242 
0243   for (auto &candidate : accepted_mothers)
0244   {
0245     mother_PID         = candidate.mother.PID;
0246     mother_true_pT     = candidate.mother.true_pT;
0247     mother_true_eta    = candidate.mother.true_eta;
0248     mother_true_phi    = candidate.mother.true_phi;
0249     mother_true_vx     = candidate.mother.true_vx;
0250     mother_true_vy     = candidate.mother.true_vy;
0251     mother_true_vz     = candidate.mother.true_vz;
0252 
0253     track_1_PID        = candidate.track_1.PID;
0254     track_1_true_pT    = candidate.track_1.true_pT;
0255     track_1_true_eta   = candidate.track_1.true_eta;
0256     track_1_true_phi   = candidate.track_1.true_phi;
0257     track_1_true_vx    = candidate.track_1.true_vx;
0258     track_1_true_vy    = candidate.track_1.true_vy;
0259     track_1_true_vz    = candidate.track_1.true_vz;
0260     track_1_reco_nmaps = candidate.track_1.reco_maps;
0261     track_1_reco_pT    = candidate.track_1.reco_pT; 
0262     track_1_reco_eta   = candidate.track_1.reco_eta; 
0263     track_1_reco_phi   = candidate.track_1.reco_phi; 
0264 
0265     track_2_PID        = candidate.track_2.PID;
0266     track_2_true_pT    = candidate.track_2.true_pT;
0267     track_2_true_eta   = candidate.track_2.true_eta;
0268     track_2_true_phi   = candidate.track_2.true_phi;
0269     track_2_true_vx    = candidate.track_2.true_vx;
0270     track_2_true_vy    = candidate.track_2.true_vy;
0271     track_2_true_vz    = candidate.track_2.true_vz;
0272     track_2_reco_nmaps = candidate.track_2.reco_maps;
0273     track_2_reco_pT    = candidate.track_2.reco_pT; 
0274     track_2_reco_eta   = candidate.track_2.reco_eta; 
0275     track_2_reco_phi   = candidate.track_2.reco_phi;
0276 
0277     float track_1_mass = TDatabasePDG::Instance()->GetParticle(track_1_PID)->Mass();
0278     float track_2_mass = TDatabasePDG::Instance()->GetParticle(track_2_PID)->Mass();
0279     PtEtaPhiMVector track_1_true_LV = PtEtaPhiMVector(track_1_true_pT, track_1_true_eta, track_1_true_phi, track_1_mass);
0280     PtEtaPhiMVector track_2_true_LV = PtEtaPhiMVector(track_2_true_pT, track_2_true_eta, track_2_true_phi, track_2_mass);
0281 
0282     PtEtaPhiMVector mother_true_LV = track_1_true_LV + track_2_true_LV;
0283 
0284     mother_true_mass = mother_true_LV.M();
0285     mother_true_rap  = mother_true_LV.Rapidity();
0286 
0287     if (track_1_reco_nmaps > 0 && track_2_reco_nmaps > 0)
0288     {
0289       PtEtaPhiMVector track_1_reco_LV = PtEtaPhiMVector(track_1_reco_pT, track_1_reco_eta, track_1_reco_phi, track_1_mass);
0290       PtEtaPhiMVector track_2_reco_LV = PtEtaPhiMVector(track_2_reco_pT, track_2_reco_eta, track_2_reco_phi, track_2_mass);
0291 
0292       PtEtaPhiMVector mother_reco_LV = track_1_reco_LV + track_2_reco_LV;
0293       mother_reco_mass = mother_reco_LV.M();
0294       mother_reco_pT = mother_reco_LV.Pt();
0295     }
0296     else
0297     {
0298       mother_reco_mass = numeric_limits<float>::quiet_NaN(); 
0299       mother_reco_pT = numeric_limits<float>::quiet_NaN();
0300     }
0301 
0302     outTree->Fill(); 
0303   }
0304 
0305   outTree->Write();
0306   outFile->Close();
0307 }
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.3.7 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!