Back to home page

sPhenix code displayed by LXR

 
 

    


File indexing completed on 2025-08-05 08:11:22

0001 #include <TDatabasePDG.h>
0002 #include <TFile.h>
0003 #include <TMath.h>
0004 #include <TObjString.h>
0005 #include <TParticle.h>
0006 #include <TParticlePDG.h>
0007 #include <TRandom3.h>
0008 #include <TTree.h>
0009 #include <TTreeIndex.h>
0010 
0011 #include <fstream>
0012 #include <iostream>
0013 #include <sstream>
0014 #include <string>
0015 #include <vector>
0016 
0017 #include "GenHadron.h"
0018 #include "Tracklet.h"
0019 #include "Vertex.h"
0020 #include "pdgidfunc.h"
0021 
0022 bool verbose_debug = true;
0023 bool useetadepadccut = false;
0024 bool savedetail = true;
0025 bool dotruthmatching = false;
0026 
0027 int main(int argc, char *argv[])
0028 {
0029     if (argc != 11)
0030     {
0031         cout << "Usage: ./TrackletAna [isdata] [evt-vtx map file] [infile] [outfile] [NevtToRun] [dRCut] [random zvtx] [random cluster set number] [clusadccutset] [clusphisizecutset]" << endl;
0032         exit(1);
0033     }
0034 
0035     // print out the input parameters
0036     for (int i = 0; i < argc; i++)
0037     {
0038         cout << "argv[" << i << "] = " << argv[i] << endl;
0039     }
0040 
0041     bool IsData = (TString(argv[1]).Atoi() == 1) ? true : false;
0042     TString EvtVtx_map_filename = TString(argv[2]);
0043     TString infilename = TString(argv[3]);
0044     TString outfilename = TString(argv[4]);
0045     int NevtToRun_ = TString(argv[5]).Atoi();
0046     float dRCut = TString(argv[6]).Atof();
0047     bool userandomzvtx = (TString(argv[7]).Atoi() == 1) ? true : false;
0048     int randclusset = TString(argv[8]).Atoi();
0049     int clusadccutset = TString(argv[9]).Atoi();
0050     int clusphisizecutset = TString(argv[10]).Atoi();
0051 
0052     TString idxstr = (IsData) ? "INTT_BCO" : "event";
0053 
0054     TRandom3 *rnd = new TRandom3(0); // Random number generator
0055 
0056     TrackletData tkldata = {};
0057 
0058     std::map<int, vector<float>> EvtVtxMap_event = EvtVtx_map_event(EvtVtx_map_filename.Data());          // use with simulation
0059     std::map<uint64_t, vector<float>> EvtVtxMap_inttbco = EvtVtx_map_inttbco(EvtVtx_map_filename.Data()); // use with data
0060 
0061     // Vertex Z reweighting
0062     TH1F *hM_vtxzweight = VtxZ_ReweiHist("/sphenix/user/hjheng/sPHENIXRepo/analysis/dNdEta_Run2023/analysis_INTT/plot/RecoPV_ana/HIJING_MDC2_ana467_20250225/VtxZ_reweight_HIJING_MDC2_ana467_20250225.root", "VtxZ_reweight_VtxZm10to10");
0063 
0064     // Random cluster
0065     int Nrandclus = RandomHit(randclusset);
0066     std::cout << "Number of random clusters: " << Nrandclus << std::endl;
0067 
0068     // cluster adc cut
0069     int clusadccut = ConstADCCut(clusadccutset);
0070     // cluster phi-size cut
0071     float clusphisizecut = ConstClusPhiCut(clusphisizecutset);
0072     // print out cut info
0073     std::cout << "Cluster ADC cut (greater than): " << clusadccut << "; Cluster Phi-size cut (smaller than): " << clusphisizecut << std::endl;
0074 
0075     TFile *f = new TFile(infilename, "READ");
0076     TTree *t = (TTree *)f->Get("EventTree");
0077     t->BuildIndex(idxstr); // Reference: https://root-forum.cern.ch/t/sort-ttree-entries/13138
0078     TTreeIndex *index = (TTreeIndex *)t->GetTreeIndex();
0079     int event, NTruthVtx, signal_process_id;
0080     vector<int> *firedTriggers = 0;
0081     uint64_t INTT_BCO;
0082     bool is_min_bias, InttBco_IsToBeRemoved;
0083     float TruthPV_trig_x, TruthPV_trig_y, TruthPV_trig_z, centrality_mbd;
0084     float mbd_south_charge_sum, mbd_north_charge_sum, mbd_charge_sum, mbd_charge_asymm, mbd_z_vtx;
0085     vector<int> *ClusLayer = 0;
0086     vector<float> *ClusX = 0, *ClusY = 0, *ClusZ = 0, *ClusPhiSize = 0, *ClusZSize = 0;
0087     vector<unsigned int> *ClusAdc = 0;
0088     vector<int> *ClusMatchedG4P_MaxE_trackID = 0, *ClusMatchedG4P_MaxClusE_ancestorTrackID = 0;
0089     vector<float> *ClusMatchedG4P_MaxE_Pt = 0, *ClusMatchedG4P_MaxE_Eta = 0, *ClusMatchedG4P_MaxE_Phi = 0;
0090     vector<int> *PrimaryG4P_PID = 0, *PrimaryG4P_trackID = 0;
0091     vector<float> *PrimaryG4P_Pt = 0, *PrimaryG4P_Eta = 0, *PrimaryG4P_Phi = 0, *PrimaryG4P_E = 0;
0092     vector<bool> *PrimaryG4P_isChargeHadron = 0;
0093     vector<int> *G4P_trackID = 0;
0094     vector<float> *G4P_Pt = 0, *G4P_Eta = 0, *G4P_Phi = 0, *G4P_E = 0;
0095     t->SetBranchAddress("event", &event);
0096     t->SetBranchAddress("INTT_BCO", &INTT_BCO);
0097     if (!IsData)
0098     {
0099         t->SetBranchAddress("NTruthVtx", &NTruthVtx);
0100         t->SetBranchAddress("signal_process_id", &signal_process_id);
0101         t->SetBranchAddress("TruthPV_trig_x", &TruthPV_trig_x);
0102         t->SetBranchAddress("TruthPV_trig_y", &TruthPV_trig_y);
0103         t->SetBranchAddress("TruthPV_trig_z", &TruthPV_trig_z);
0104         t->SetBranchAddress("PrimaryG4P_PID", &PrimaryG4P_PID);
0105         t->SetBranchAddress("PrimaryG4P_isChargeHadron", &PrimaryG4P_isChargeHadron);
0106         t->SetBranchAddress("PrimaryG4P_trackID", &PrimaryG4P_trackID);
0107         t->SetBranchAddress("PrimaryG4P_Pt", &PrimaryG4P_Pt);
0108         t->SetBranchAddress("PrimaryG4P_Eta", &PrimaryG4P_Eta);
0109         t->SetBranchAddress("PrimaryG4P_Phi", &PrimaryG4P_Phi);
0110         t->SetBranchAddress("PrimaryG4P_E", &PrimaryG4P_E);
0111         if (dotruthmatching)
0112         {
0113             t->SetBranchAddress("ClusMatchedG4P_MaxE_trackID", &ClusMatchedG4P_MaxE_trackID);
0114             t->SetBranchAddress("ClusMatchedG4P_MaxClusE_ancestorTrackID", &ClusMatchedG4P_MaxClusE_ancestorTrackID);
0115             t->SetBranchAddress("ClusMatchedG4P_MaxE_Pt", &ClusMatchedG4P_MaxE_Pt);
0116             t->SetBranchAddress("ClusMatchedG4P_MaxE_Eta", &ClusMatchedG4P_MaxE_Eta);
0117             t->SetBranchAddress("ClusMatchedG4P_MaxE_Phi", &ClusMatchedG4P_MaxE_Phi);
0118             t->SetBranchAddress("G4P_trackID", &G4P_trackID);
0119             t->SetBranchAddress("G4P_Pt", &G4P_Pt);
0120             t->SetBranchAddress("G4P_Eta", &G4P_Eta);
0121             t->SetBranchAddress("G4P_Phi", &G4P_Phi);
0122             t->SetBranchAddress("G4P_E", &G4P_E);
0123         }
0124 
0125         InttBco_IsToBeRemoved = false;
0126     }
0127     if (IsData)
0128     {
0129         t->SetBranchAddress("firedTriggers", &firedTriggers);
0130         t->SetBranchAddress("InttBco_IsToBeRemoved", &InttBco_IsToBeRemoved);
0131     }
0132     t->SetBranchAddress("is_min_bias", &is_min_bias);
0133     t->SetBranchAddress("MBD_centrality", &centrality_mbd);
0134     t->SetBranchAddress("MBD_z_vtx", &mbd_z_vtx);
0135     t->SetBranchAddress("MBD_south_charge_sum", &mbd_south_charge_sum);
0136     t->SetBranchAddress("MBD_north_charge_sum", &mbd_north_charge_sum);
0137     t->SetBranchAddress("MBD_charge_sum", &mbd_charge_sum);
0138     t->SetBranchAddress("MBD_charge_asymm", &mbd_charge_asymm);
0139     t->SetBranchAddress("ClusLayer", &ClusLayer);
0140     t->SetBranchAddress("ClusX", &ClusX);
0141     t->SetBranchAddress("ClusY", &ClusY);
0142     t->SetBranchAddress("ClusZ", &ClusZ);
0143     t->SetBranchAddress("ClusPhiSize", &ClusPhiSize);
0144     t->SetBranchAddress("ClusZSize", &ClusZSize);
0145     t->SetBranchAddress("ClusAdc", &ClusAdc);
0146 
0147     TFile *outfile = new TFile(outfilename, "RECREATE");
0148     TTree *minitree = new TTree("minitree", "Minitree of Reconstructed Tracklets");
0149     SetMinitree(minitree, tkldata, savedetail);
0150 
0151     for (int i = 0; i < NevtToRun_; i++)
0152     {
0153         Long64_t local = t->LoadTree(index->GetIndex()[i]);
0154         t->GetEntry(local);
0155 
0156         cout << "i = " << i << " event = " << event << " INTT BCO = " << INTT_BCO << " NevtToRun_ = " << NevtToRun_ << " local = " << local << " InttBco_IsToBeRemoved = " << InttBco_IsToBeRemoved << endl;
0157 
0158         // Discard events with a small mutual BCO difference between adjacent events
0159         if (InttBco_IsToBeRemoved)
0160             continue;
0161 
0162         vector<float> PV = (IsData) ? EvtVtxMap_inttbco[INTT_BCO] : EvtVtxMap_event[event];
0163         tkldata.event = event;
0164         tkldata.INTT_BCO = INTT_BCO;
0165         tkldata.PV_x = PV[0];
0166         tkldata.PV_y = PV[1];
0167         float PVz = (userandomzvtx) ? rnd->Uniform(-30, 30) : PV[2];
0168         tkldata.PV_z = PVz;
0169 
0170         if (!IsData)
0171         {
0172             tkldata.TruthPV_x = TruthPV_trig_x;
0173             tkldata.TruthPV_y = TruthPV_trig_y;
0174             tkldata.TruthPV_z = TruthPV_trig_z;
0175         }
0176 
0177         if (PV[2] < -10 || PV[2] > 10)
0178         {
0179             tkldata.vtxzwei = 1.;
0180         }
0181         else
0182         {
0183             if (IsData)
0184                 tkldata.vtxzwei = 1;
0185             else
0186                 tkldata.vtxzwei = hM_vtxzweight->GetBinContent(hM_vtxzweight->FindBin(PV[2]));
0187         }
0188 
0189         tkldata.firedTrig10_MBDSNgeq2 = (IsData) ? (find(firedTriggers->begin(), firedTriggers->end(), 10) != firedTriggers->end()) : true;
0190         tkldata.firedTrig12_vtxle10cm = (IsData) ? (find(firedTriggers->begin(), firedTriggers->end(), 12) != firedTriggers->end()) : true;
0191         tkldata.firedTrig13_vtxle30cm = (IsData) ? (find(firedTriggers->begin(), firedTriggers->end(), 13) != firedTriggers->end()) : true;
0192         tkldata.MbdNSge0 = (mbd_north_charge_sum > 0 || mbd_south_charge_sum > 0);
0193         tkldata.MbdZvtxle10cm = (mbd_z_vtx > -10 && mbd_z_vtx < 10);
0194         tkldata.validMbdVtx = (mbd_z_vtx == mbd_z_vtx);
0195         tkldata.InttBco_IsToBeRemoved = (IsData) ? InttBco_IsToBeRemoved : false;
0196 
0197         // Centrality
0198         tkldata.centrality_mbd = centrality_mbd;
0199         tkldata.mbd_z_vtx = mbd_z_vtx;
0200         tkldata.mbd_south_charge_sum = mbd_south_charge_sum;
0201         tkldata.mbd_north_charge_sum = mbd_north_charge_sum;
0202         tkldata.mbd_charge_sum = mbd_charge_sum;
0203         tkldata.mbd_charge_asymm = mbd_charge_asymm;
0204 
0205         // A selection to select events -> mimic 0-pileup environment
0206         tkldata.pu0_sel = (NTruthVtx == 1) ? true : false;
0207         tkldata.is_min_bias = is_min_bias;
0208         tkldata.process = (IsData) ? -1E9 : signal_process_id; // Exclude single diffractive events
0209 
0210         cout << "[Event info] Event " << event << " firedTrig10_MBDSNgeq2 = " << tkldata.firedTrig10_MBDSNgeq2 << " firedTrig12_vtxle10cm = " << tkldata.firedTrig12_vtxle10cm << "; firedTrig13_vtxle30cm = " << tkldata.firedTrig13_vtxle30cm << "; MbdNSge0 = " << tkldata.MbdNSge0 << "; MbdZvtxle10cm = " << tkldata.MbdZvtxle10cm << "; validMbdVtx = " << tkldata.validMbdVtx << endl;
0211 
0212         // Prepare clusters
0213         for (size_t ihit = 0; ihit < ClusLayer->size(); ihit++)
0214         {
0215             if (ClusLayer->at(ihit) < 3 || ClusLayer->at(ihit) > 6)
0216             {
0217                 cout << "[WARNING] Unknown layer: " << ClusLayer->at(ihit) << endl; // Should not happen
0218                 continue;
0219             }
0220 
0221             int layer = (ClusLayer->at(ihit) == 3 || ClusLayer->at(ihit) == 4) ? 0 : 1;
0222             Hit *hit = new Hit(ClusX->at(ihit), ClusY->at(ihit), ClusZ->at(ihit), PV[0], PV[1], PVz, layer, ClusPhiSize->at(ihit), ClusAdc->at(ihit));
0223 
0224             if (ClusAdc->at(ihit) <= clusadccut)
0225                 continue;
0226 
0227             if (ClusPhiSize->at(ihit) > clusphisizecut)
0228                 continue;
0229 
0230             if (!IsData && dotruthmatching)
0231             {
0232                 hit->SetMatchedG4P(ClusMatchedG4P_MaxE_trackID->at(ihit), ClusMatchedG4P_MaxClusE_ancestorTrackID->at(ihit), ClusMatchedG4P_MaxE_Pt->at(ihit), ClusMatchedG4P_MaxE_Eta->at(ihit), ClusMatchedG4P_MaxE_Phi->at(ihit));
0233             }
0234 
0235             tkldata.layers[layer].push_back(hit);
0236             tkldata.cluslayer.push_back(layer);
0237             tkldata.clusphi.push_back(hit->Phi());
0238             tkldata.cluseta.push_back(hit->Eta());
0239             tkldata.clusphisize.push_back(ClusPhiSize->at(ihit));
0240             tkldata.cluszsize.push_back(ClusZSize->at(ihit));
0241             tkldata.clusadc.push_back(ClusAdc->at(ihit));
0242         }
0243 
0244         std::cout << "Number of clusters before selections: " << ClusLayer->size() << ", after selections: " << tkldata.cluslayer.size() << std::endl;
0245 
0246         // For systematic: adding random clusters (after selections)
0247         // int Nrandclus = RandomHit(randclusset);
0248         if (Nrandclus > 0)
0249         {
0250             for (int irand = 0; irand < Nrandclus; irand++)
0251             {
0252                 // Hit *RandomHit(float vx, float vy, float vz, int layer)
0253                 Hit *randhit_inner = RandomHit(PV[0], PV[1], PVz, 0);
0254                 Hit *randhit_outer = RandomHit(PV[0], PV[1], PVz, 1);
0255                 tkldata.layers[0].push_back(randhit_inner);
0256                 tkldata.layers[1].push_back(randhit_outer);
0257             }
0258         }
0259 
0260         tkldata.NClusLayer1 = tkldata.layers[0].size();
0261         tkldata.NClusLayer2 = tkldata.layers[1].size();
0262 
0263         ProtoTracklets(tkldata, dRCut);
0264         RecoTracklets(tkldata);
0265 
0266         if (!IsData)
0267         {
0268             std::map<int, unsigned int> PrimaryG4PPID_count, absPrimaryG4PPID_count;
0269             int Nleptons = 0; // number of electrons and muons
0270             PrimaryG4PPID_count.clear();
0271             absPrimaryG4PPID_count.clear();
0272 
0273             // Generated charged hadrons
0274             for (size_t ihad = 0; ihad < PrimaryG4P_PID->size(); ihad++)
0275             {
0276                 TString particleclass = TString(TDatabasePDG::Instance()->GetParticle(PrimaryG4P_PID->at(ihad))->ParticleClass());
0277                 bool isStable = (TDatabasePDG::Instance()->GetParticle(PrimaryG4P_PID->at(ihad))->Stable() == 1) ? true : false;
0278                 double charge = TDatabasePDG::Instance()->GetParticle(PrimaryG4P_PID->at(ihad))->Charge();
0279                 double lifetime = TDatabasePDG::Instance()->GetParticle(PrimaryG4P_PID->at(ihad))->Lifetime(); // proper lifetime (https://root.cern.ch/doc/master/classTParticlePDG.html#a10fd025a1e867ec3b27903c1b7d2f899)
0280                 double decaylength = lifetime * TMath::C() * 1E2;                                              // decay length in cm
0281                 bool isHadron = (particleclass.Contains("Baryon") || particleclass.Contains("Meson"));
0282                 bool isChargeHadron = (isStable && (charge != 0) && isHadron);
0283                 bool isChargeHadron_alt = is_chargedHadron(PrimaryG4P_PID->at(ihad));
0284                 bool twodefsame = (isChargeHadron == isChargeHadron_alt) ? true : false;
0285 
0286                 /*
0287                 if (verbose_debug)
0288                 {
0289                     cout << std::left << std::setw(5) << "PID = " << std::setw(5) << PrimaryG4P_PID->at(ihad);
0290                     cout << std::left << std::setw(20) << " Particle class = " << std::setw(15) << particleclass;
0291                     cout << std::left << std::setw(12) << " Stable = " << std::setw(5) << isStable;
0292                     cout << std::left << std::setw(20) << " Proper lifetime = " << std::setw(15) << lifetime;
0293                     cout << std::left << std::setw(18) << " Decay length = " << std::setw(15) << decaylength;
0294                     cout << std::left << std::setw(12) << " Charge = " << std::setw(5) << charge;
0295                     cout << std::left << std::setw(15) << " Is hadron = " << std::setw(5) << isHadron;
0296                     cout << std::left << std::setw(23) << " Is charged hadron = " << std::setw(5) << isChargeHadron;
0297                     cout << std::left << std::setw(30) << " Is charged hadron (alt) = " << std::setw(5) << isChargeHadron_alt;
0298                     cout << std::left << std::setw(34) << " Two definitions are the same = " << std::setw(5) << twodefsame << endl;
0299                 }
0300                 */
0301 
0302                 if (verbose_debug)
0303                 {
0304                     PrimaryG4PPID_count[PrimaryG4P_PID->at(ihad)]++;
0305                     absPrimaryG4PPID_count[abs(PrimaryG4P_PID->at(ihad))]++;
0306                     if (abs(PrimaryG4P_PID->at(ihad)) == 11 || abs(PrimaryG4P_PID->at(ihad)) == 13)
0307                         Nleptons++;
0308                 }
0309 
0310                 tkldata.PrimaryG4P_trackID.push_back(PrimaryG4P_trackID->at(ihad));
0311                 tkldata.PrimaryG4P_PID.push_back(PrimaryG4P_PID->at(ihad));
0312                 tkldata.PrimaryG4P_Pt.push_back(PrimaryG4P_Pt->at(ihad));
0313                 tkldata.PrimaryG4P_eta.push_back(PrimaryG4P_Eta->at(ihad));
0314                 tkldata.PrimaryG4P_phi.push_back(PrimaryG4P_Phi->at(ihad));
0315                 tkldata.PrimaryG4P_IsRecotkl.push_back(false); // default
0316 
0317                 // if (is_chargedHadron(PrimaryG4P_PID->at(ihad)) == false)
0318                 //     continue;
0319 
0320                 if (is_charged(PrimaryG4P_PID->at(ihad)) == false)
0321                     continue;
0322 
0323                 // GenHadron *genhadron = new GenHadron(PrimaryG4P_Pt->at(ihad), PrimaryG4P_Eta->at(ihad), PrimaryG4P_Phi->at(ihad), PrimaryG4P_E->at(ihad));
0324                 // tkldata.GenHadrons.push_back(genhadron);
0325                 tkldata.GenHadron_PID.push_back(PrimaryG4P_PID->at(ihad));
0326                 tkldata.GenHadron_trackID.push_back(PrimaryG4P_trackID->at(ihad));
0327                 tkldata.GenHadron_Pt.push_back(PrimaryG4P_Pt->at(ihad));
0328                 tkldata.GenHadron_eta.push_back(PrimaryG4P_Eta->at(ihad));
0329                 tkldata.GenHadron_phi.push_back(PrimaryG4P_Phi->at(ihad));
0330                 tkldata.GenHadron_IsRecotkl.push_back(false); // default
0331             }
0332             tkldata.NGenHadron = tkldata.GenHadron_PID.size();
0333 
0334             if (dotruthmatching)
0335             {
0336                 for (size_t ig4p = 0; ig4p < G4P_trackID->size(); ig4p++)
0337                 {
0338                     tkldata.G4P_trackID.push_back(G4P_trackID->at(ig4p));
0339                     tkldata.G4P_Pt.push_back(G4P_Pt->at(ig4p));
0340                     tkldata.G4P_eta.push_back(G4P_Eta->at(ig4p));
0341                     tkldata.G4P_phi.push_back(G4P_Phi->at(ig4p));
0342                 }
0343             }
0344 
0345             if (verbose_debug)
0346             {
0347                 // print PrimaryG4PPID_count
0348                 // cout << "--------------------------------------------------------------------" << endl;
0349                 // for (auto it = PrimaryG4PPID_count.begin(); it != PrimaryG4PPID_count.end(); it++)
0350                 // {
0351                 //     cout << "PID = " << it->first << "; Count = " << it->second << endl;
0352                 // }
0353                 // cout << "--------------------------------------------------------------------" << endl;
0354                 // for (auto it = absPrimaryG4PPID_count.begin(); it != absPrimaryG4PPID_count.end(); it++)
0355                 // {
0356                 //     cout << "|PID| = " << it->first << "; Count = " << it->second << endl;
0357                 // }
0358                 cout << "--------------------------------------------------------------------" << endl;
0359                 cout << "Number of charged leptons (electrons and muons) = " << Nleptons << endl;
0360                 cout << "--------------------------------------------------------------------" << endl;
0361             }
0362 
0363             if (dotruthmatching)
0364                 RecoTkl_isG4P(tkldata);
0365         }
0366 
0367         cout << "Number of clusters = " << tkldata.cluslayer.size() << "; Number of reco tracklets = " << tkldata.NRecotkl_Raw << endl;
0368 
0369         // if (verbose_debug)
0370             // PrintVecSize(tkldata);
0371 
0372         minitree->Fill();
0373         ResetVec(tkldata);
0374 
0375         cout << "----------" << endl;
0376     }
0377 
0378     outfile->cd();
0379     minitree->Write("", TObject::kOverwrite);
0380     outfile->Close();
0381 
0382     f->Close();
0383 
0384     return 0;
0385 }