Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-01-09 10:34:30

0001 #ifdef __CINT__
0002 
0003 #pragma link off all globals;
0004 #pragma link off all classes;
0005 #pragma link off all functions;
0006 
0007 #pragma link C++ class PlotFile;
0008 #endif
0009 
0010 #ifndef __CINT__
0011 #include <stdio.h>
0012 #include <stdlib.h>
0013 #include <fstream>
0014 #include <iostream>
0015 #include <iomanip>
0016 #include <string>
0017 #include <sys/types.h>
0018 #include <sys/stat.h>
0019 #include <dirent.h>
0020 
0021 #include "math.h"
0022 #include "string.h"
0023 
0024 #include "TROOT.h"
0025 #include "TFile.h"
0026 #include "TChain.h"
0027 #include "TH1D.h"
0028 #include "TH2D.h"
0029 #include "TH3D.h"
0030 #include "TStyle.h"
0031 #include "TCanvas.h"
0032 #include "TProfile.h"
0033 #include "TTree.h"
0034 #include "TNtuple.h"
0035 #include "TRandom3.h"
0036 #include "TMath.h"
0037 #include "TSystem.h"
0038 #include "TUnixSystem.h"
0039 #include "TVector2.h"
0040 #include "TVector3.h"
0041 #include "TLorentzVector.h"
0042 #include "TTreeReader.h"
0043 #include "TTreeReaderValue.h"
0044 #include "TTreeReaderArray.h"
0045 #include "TLatex.h"
0046 #endif
0047 
0048 
0049 #include "StPhysicalHelix.h"
0050 #include "SystemOfUnits.h"
0051 #include "PhysicalConstants.h"
0052 
0053 using namespace std;
0054 
0055 const double gPionMass = 0.13957;
0056 const double gKaonMass = 0.493677;
0057 
0058 const double twoPi = 2.*3.1415927;
0059 const double eMass = 0.000511;
0060 
0061 const double bField = -1.7; // Tesla
0062 
0063 TVector3 getDcaToVtx(TVector3 pos, TVector3 mom, int charge, TVector3 vtx);
0064 
0065 TLorentzVector getPairParent(TVector3 pos1, TVector3 mom1, int charge1, double mass1,
0066                  TVector3 pos2, TVector3 mom2, int charge2, double mass2,
0067                  TVector3 vtx,
0068                  float &dcaDaughters, float &cosTheta, float &decayLength, float &V0DcaToVtx);
0069 
0070 int main(int argc, char **argv)
0071 {
0072   if(argc!=3 && argc!=1) return 0;
0073 
0074   TString listname;
0075   TString outname;
0076 
0077   if(argc==1)
0078     {
0079       listname  = "test.list";
0080       outname = "test.root";
0081     }
0082 
0083   if(argc==3)
0084     {
0085       listname = argv[1];
0086       outname = argv[2];
0087     }  
0088 
0089   // ===== read in root files =====
0090   TChain *chain = new TChain("events");
0091 
0092   int nfiles = 0;
0093   char filename[512];
0094   ifstream *inputstream = new ifstream;
0095   inputstream->open(listname.Data());
0096   if(!inputstream)
0097     {
0098       printf("[e] Cannot open file list: %s\n", listname.Data());
0099     }
0100   while(inputstream->good())
0101     {
0102       inputstream->getline(filename, 512);
0103       if(inputstream->good())
0104     {
0105       TFile *ftmp = TFile::Open(filename, "read");
0106       if(!ftmp||!(ftmp->IsOpen())||!(ftmp->GetNkeys())) 
0107         {
0108           printf("[e] Could you open file: %s\n", filename);
0109         } 
0110       else
0111         {
0112           cout<<"[i] Add "<<nfiles<<"th file: "<<filename<<endl;
0113           chain->Add(filename);
0114           nfiles++;
0115         }
0116     }
0117     }
0118   inputstream->close();
0119   printf("[i] Read in %d files with %lld events in total\n", nfiles, chain->GetEntries());
0120 
0121   // ===== define histograms =====
0122   TH1F *hEventStat = new TH1F("hEventStat", "Event statistics", 5, 0, 5);
0123   hEventStat->GetXaxis()->SetBinLabel(1, "MC events");
0124   hEventStat->GetXaxis()->SetBinLabel(2, "D0");
0125   hEventStat->GetXaxis()->SetBinLabel(3, "D0 -> pi+K");
0126   hEventStat->GetXaxis()->SetBinLabel(4, "Reco D0");
0127 
0128   TH1F *hMcVtxX = new TH1F("hMcVtxX", "x position of MC vertex;x (mm)", 100, -5.05, 4.95);
0129   TH1F *hMcVtxY = new TH1F("hMcVtxY", "y position of MC vertex;y (mm)", 500, -5.01, 4.99);
0130   TH1F *hMcVtxZ = new TH1F("hMcVtxZ", "z position of MC vertex;z (mm)", 400, -200, 200);
0131 
0132   TH2F *hMCD0PtRap = new TH2F("hMCD0PtRap", "MC D^{0};y;p_{T} (GeV/c)", 20, -5, 5, 100, 0, 10);
0133   
0134   TH3F *h3InvMass = new TH3F("h3InvMass", "Invariant mass of unlike-sign #piK pairs;p_{T} (GeV/c);y;M_{#piK} (GeV/c^{2})", 100, 0, 10, 20, -5, 5, 100, 1.6, 2.0);
0135 
0136 
0137   TTreeReader treereader(chain);
0138   // MC  
0139   TTreeReaderArray<int> mcPartGenStatus = {treereader, "MCParticles.generatorStatus"};
0140   TTreeReaderArray<int> mcPartPdg = {treereader, "MCParticles.PDG"};
0141   TTreeReaderArray<float> mcPartCharge = {treereader, "MCParticles.charge"};
0142   TTreeReaderArray<unsigned int> mcPartParent_begin = {treereader, "MCParticles.parents_begin"};
0143   TTreeReaderArray<unsigned int> mcPartParent_end = {treereader, "MCParticles.parents_end"};
0144   TTreeReaderArray<int> mcPartParent_index = {treereader, "_MCParticles_parents.index"};
0145   TTreeReaderArray<unsigned int> mcPartDaughter_begin = {treereader, "MCParticles.daughters_begin"};
0146   TTreeReaderArray<unsigned int> mcPartDaughter_end = {treereader, "MCParticles.daughters_end"};
0147   TTreeReaderArray<int> mcPartDaughter_index = {treereader, "_MCParticles_daughters.index"};
0148   TTreeReaderArray<double> mcPartMass = {treereader, "MCParticles.mass"};
0149   TTreeReaderArray<double> mcPartVx = {treereader, "MCParticles.vertex.x"};
0150   TTreeReaderArray<double> mcPartVy = {treereader, "MCParticles.vertex.y"};
0151   TTreeReaderArray<double> mcPartVz = {treereader, "MCParticles.vertex.z"};
0152   TTreeReaderArray<double> mcMomPx = {treereader, "MCParticles.momentum.x"};
0153   TTreeReaderArray<double> mcMomPy = {treereader, "MCParticles.momentum.y"};
0154   TTreeReaderArray<double> mcMomPz = {treereader, "MCParticles.momentum.z"};
0155   TTreeReaderArray<double> mcEndPointX = {treereader, "MCParticles.endpoint.x"};
0156   TTreeReaderArray<double> mcEndPointY = {treereader, "MCParticles.endpoint.y"};
0157   TTreeReaderArray<double> mcEndPointZ = {treereader, "MCParticles.endpoint.z"};
0158 
0159   TTreeReaderArray<unsigned int> assocChSimID = {treereader, "ReconstructedChargedParticleAssociations.simID"};
0160   TTreeReaderArray<unsigned int> assocChRecID = {treereader, "ReconstructedChargedParticleAssociations.recID"};
0161   TTreeReaderArray<float> assocWeight = {treereader, "ReconstructedChargedParticleAssociations.weight"};
0162 
0163   TTreeReaderArray<float> rcMomPx = {treereader, "ReconstructedChargedParticles.momentum.x"};
0164   TTreeReaderArray<float> rcMomPy = {treereader, "ReconstructedChargedParticles.momentum.y"};
0165   TTreeReaderArray<float> rcMomPz = {treereader, "ReconstructedChargedParticles.momentum.z"};
0166   TTreeReaderArray<float> rcPosx = {treereader, "ReconstructedChargedParticles.referencePoint.x"};
0167   TTreeReaderArray<float> rcPosy = {treereader, "ReconstructedChargedParticles.referencePoint.y"};
0168   TTreeReaderArray<float> rcPosz = {treereader, "ReconstructedChargedParticles.referencePoint.z"};
0169   TTreeReaderArray<float> rcCharge = {treereader, "ReconstructedChargedParticles.charge"};
0170   TTreeReaderArray<int>   rcPdg = {treereader, "ReconstructedChargedParticles.PDG"};
0171 
0172   // needed for accessing number of hits used for track reconstruction
0173   TTreeReaderArray<unsigned int> partAssocTrk_begin = {treereader, "ReconstructedChargedParticles.tracks_begin"};
0174   TTreeReaderArray<int> partAssocTrk_index = {treereader, "_ReconstructedChargedParticles_tracks.index"};
0175   TTreeReaderArray<int> trkAssocTraj_index = {treereader, "_CentralCKFTracks_trajectory.index"};
0176   TTreeReaderArray<unsigned int> nMeasurements = {treereader, "CentralCKFTrajectories.nMeasurements"};
0177 
0178   TTreeReaderArray<float> rcTrkLoca = {treereader, "CentralCKFTrackParameters.loc.a"};
0179   TTreeReaderArray<float> rcTrkLocb = {treereader, "CentralCKFTrackParameters.loc.b"};
0180   TTreeReaderArray<float> rcTrkqOverP = {treereader, "CentralCKFTrackParameters.qOverP"};
0181   TTreeReaderArray<float> rcTrkTheta = {treereader, "CentralCKFTrackParameters.theta"};
0182   TTreeReaderArray<float> rcTrkPhi = {treereader, "CentralCKFTrackParameters.phi"};
0183 
0184   TTreeReaderArray<float> CTVx = {treereader, "CentralTrackVertices.position.x"};
0185   TTreeReaderArray<float> CTVy = {treereader, "CentralTrackVertices.position.y"};
0186   TTreeReaderArray<float> CTVz = {treereader, "CentralTrackVertices.position.z"};
0187   TTreeReaderArray<int> CTVndf = {treereader, "CentralTrackVertices.ndf"};
0188   TTreeReaderArray<float> CTVchi2 = {treereader, "CentralTrackVertices.chi2"};
0189   TTreeReaderArray<float> CTVerr_xx = {treereader, "CentralTrackVertices.positionError.xx"};
0190   TTreeReaderArray<float> CTVerr_yy = {treereader, "CentralTrackVertices.positionError.yy"};
0191   TTreeReaderArray<float> CTVerr_zz = {treereader, "CentralTrackVertices.positionError.zz"};
0192 
0193   TTreeReaderArray<int> prim_vtx_index = {treereader, "PrimaryVertices_objIdx.index"};
0194 
0195   TTreeReaderArray<unsigned int> vtxAssocPart_begin = {treereader, "CentralTrackVertices.associatedParticles_begin"};
0196   TTreeReaderArray<unsigned int> vtxAssocPart_end = {treereader, "CentralTrackVertices.associatedParticles_end"};
0197   TTreeReaderArray<int> vtxAssocPart_index = {treereader, "_CentralTrackVertices_associatedParticles.index"};
0198   
0199   int nevents = 0;
0200   while(treereader.Next())
0201     {
0202       if(nevents%1000==0) printf("\n[i] New event %d\n",nevents);
0203 
0204       // ===== find MC primary vertex
0205       int nMCPart = mcPartMass.GetSize();
0206       TVector3 vertex_mc(-999., -999., -999.);
0207       for(int imc=0; imc<nMCPart; imc++)
0208     {
0209       if(mcPartGenStatus[imc] == 4 && mcPartPdg[imc] == 11)
0210         {
0211           vertex_mc.SetXYZ(mcEndPointX[imc], mcEndPointY[imc], mcEndPointZ[imc]);
0212           break;
0213         }
0214     }
0215       hEventStat->Fill(0.5);
0216       hMcVtxX->Fill(vertex_mc.x());
0217       hMcVtxY->Fill(vertex_mc.y());
0218       hMcVtxZ->Fill(vertex_mc.z());
0219 
0220       // ===== get RC primary vertex
0221       TVector3 vertex_rc(-999., -999., -999.);
0222       if(prim_vtx_index.GetSize()>0)
0223     {
0224       int rc_vtx_index = prim_vtx_index[0];
0225       vertex_rc.SetXYZ(CTVx[rc_vtx_index], CTVy[rc_vtx_index], CTVz[rc_vtx_index]);
0226     }
0227 
0228       // ===== map MC and RC particles
0229       int nAssoc = assocChRecID.GetSize();
0230       map<int, int> assoc_map_to_rc;
0231       map<int, int> assoc_map_to_mc;
0232       for(unsigned int rc_index=0; rc_index<rcMomPx.GetSize(); rc_index++)
0233     {
0234       // loop over the association to find the matched MC particle
0235       // with largest weight
0236       double max_weight = 0;
0237       int matched_mc_index = -1;
0238       for(int j=0; j<nAssoc; j++)
0239         {
0240           if(assocChRecID[j] != rc_index) continue;
0241           if(assocWeight[j] > max_weight)
0242         {
0243           max_weight = assocWeight[j];
0244           matched_mc_index = assocChSimID[j];
0245         }
0246         }
0247 
0248       // build the map
0249       assoc_map_to_rc[matched_mc_index] = rc_index;
0250       assoc_map_to_mc[rc_index] = matched_mc_index;
0251     }
0252 
0253       // ===== look for D0 -> pi+K
0254       for(int imc=0; imc<nMCPart; imc++)
0255     {
0256       if(fabs(mcPartPdg[imc]) != 421) continue;
0257       hEventStat->Fill(1.5);
0258       
0259       int nDuaghters = mcPartDaughter_end[imc]-mcPartDaughter_begin[imc];
0260       if(nDuaghters!=2) continue;
0261 
0262       // find D0 that decay into pi+K
0263       bool is_pik_decay = false;      
0264       int daug_index_1 = mcPartDaughter_index[mcPartDaughter_begin[imc]];
0265       int daug_index_2 = mcPartDaughter_index[mcPartDaughter_begin[imc]+1];
0266       int daug_pdg_1 = mcPartPdg[daug_index_1];
0267       int daug_pdg_2 = mcPartPdg[daug_index_2];
0268       if( (fabs(daug_pdg_1)==321 && fabs(daug_pdg_2)==211) || (fabs(daug_pdg_1)==211 && fabs(daug_pdg_2)==321) )
0269         {
0270           is_pik_decay = true;
0271         }
0272       if(!is_pik_decay) continue;
0273       hEventStat->Fill(2.5);
0274 
0275       // D0 kinematics
0276       TLorentzVector mc_mom_vec;
0277       mc_mom_vec.SetXYZM(mcMomPx[imc], mcMomPy[imc], mcMomPz[imc], mcPartMass[imc]);
0278       double mcRap = mc_mom_vec.Rapidity();
0279       double mcPt = mc_mom_vec.Pt();
0280       hMCD0PtRap->Fill(mcRap, mcPt);
0281     }
0282 
0283 
0284       // ===== Get reconstructed pions and kaons
0285       const int pid_mode = 1; // 0 - truth; 1 - realistic
0286       vector<unsigned int> pi_index;
0287       vector<unsigned int> k_index;
0288       pi_index.clear();
0289       k_index.clear();
0290       for(unsigned int rc_index=0; rc_index<rcMomPx.GetSize(); rc_index++)
0291     {
0292 
0293       // cut on nhits
0294 
0295       // this assumes 1-to-1-to-1 correspondance between particle-track_trajectory
0296       int nhits = nMeasurements[rc_index];
0297       
0298       // the following implemenation is more robust
0299       // int begin_index = partAssocTrk_begin[rc_index]; // take the first track; currently there is 1-to-1 correspondance between track and particle
0300       // int track_index = partAssocTrk_index[begin_index];
0301       // int traj_index = trkAssocTraj_index[track_index]; // 1-to-1 correspondance between track and trajectory
0302       // int nhits = nMeasurements[traj_index];
0303       // printf("rc_index = %d, begin_index = %d, track_index = %d, traj_index = %d, nMeasurements = %d = %d\n", rc_index, begin_index, track_index, traj_index, nhits, nMeasurements[rc_index]);
0304 
0305       if(nhits<4) continue;
0306       
0307       if(pid_mode==0)
0308         {
0309           // Use truth PID information
0310           int iSimPartID = -1;
0311           if(assoc_map_to_mc.find(rc_index) != assoc_map_to_mc.end()) iSimPartID = assoc_map_to_mc[rc_index];
0312           if(iSimPartID>=0)
0313         {
0314           if(fabs(mcPartPdg[iSimPartID]) == 211) pi_index.push_back(rc_index);
0315           if(fabs(mcPartPdg[iSimPartID]) == 321) k_index.push_back(rc_index);
0316         }
0317         }
0318       else if(pid_mode==1)
0319         {
0320           // Use reconstructed PID
0321           if(fabs(rcPdg[rc_index]) == 211) pi_index.push_back(rc_index);
0322           if(fabs(rcPdg[rc_index]) == 321) k_index.push_back(rc_index);
0323         }
0324     }
0325 
0326       // ===== pair pion and kaon
0327       for(unsigned int i=0; i<pi_index.size(); i++)
0328     {
0329       // convert from local to global coordinate
0330       // Helix code uses centimeter as the default unit, while ePIC software uses millimeter
0331       TVector3 pos1(rcTrkLoca[pi_index[i]] * sin(rcTrkPhi[pi_index[i]]) * -1 * millimeter, rcTrkLoca[pi_index[i]] * cos(rcTrkPhi[pi_index[i]]) * millimeter, rcTrkLocb[pi_index[i]] * millimeter);
0332       TVector3 mom1(rcMomPx[pi_index[i]], rcMomPy[pi_index[i]], rcMomPz[pi_index[i]]);
0333       int charge1 = rcCharge[pi_index[i]];
0334       
0335       TVector3 dcaToVtx = getDcaToVtx(pos1, mom1, charge1, vertex_rc);
0336       for(unsigned int j=0; j<k_index.size(); j++)
0337         {
0338           TVector3 pos2(rcTrkLoca[k_index[j]] * sin(rcTrkPhi[k_index[j]]) * -1 * millimeter, rcTrkLoca[k_index[j]] * cos(rcTrkPhi[k_index[j]]) * millimeter, rcTrkLocb[k_index[j]] * millimeter);
0339           TVector3 mom2(rcMomPx[k_index[j]], rcMomPy[k_index[j]], rcMomPz[k_index[j]]);
0340           int charge2 = rcCharge[k_index[j]];
0341       
0342           TVector3 dcaToVtx2 = getDcaToVtx(pos2, mom2, charge2, vertex_rc);
0343           if(charge1*charge2<0)
0344         {
0345           // -- only look at unlike-sign pi+k pair
0346 
0347           float dcaDaughters, cosTheta, decayLength, V0DcaToVtx;
0348           TLorentzVector parent = getPairParent(pos1, mom1, charge1, gPionMass,
0349                             pos2, mom2, charge2, gKaonMass,
0350                             vertex_rc,
0351                             dcaDaughters, cosTheta, decayLength, V0DcaToVtx);
0352           h3InvMass->Fill(parent.Pt(), parent.Rapidity(), parent.M());
0353         }
0354         }
0355     }         
0356       nevents++;
0357     }
0358 
0359   TFile *outfile = new TFile(outname.Data(), "recreate");
0360 
0361   hEventStat->Write();
0362   hMcVtxX->Write();
0363   hMcVtxY->Write();
0364   hMcVtxZ->Write();
0365   hMCD0PtRap->Write();
0366   h3InvMass->Write();
0367   outfile->Close();
0368 
0369 }
0370 
0371 //======================================
0372 TVector3 getDcaToVtx(TVector3 pos, TVector3 mom, int charge, TVector3 vtx)
0373 {
0374   StPhysicalHelix pHelix(mom, pos, bField * tesla, charge);
0375 
0376   TVector3 vtx_tmp;
0377   vtx_tmp.SetXYZ(vtx.x()*millimeter, vtx.y()*millimeter, vtx.z()*millimeter);
0378   
0379   pHelix.moveOrigin(pHelix.pathLength(vtx_tmp));
0380   TVector3 dcaToVtx = pHelix.origin() - vtx_tmp;
0381 
0382   dcaToVtx.SetXYZ(dcaToVtx.x()/millimeter, dcaToVtx.y()/millimeter, dcaToVtx.z()/millimeter);
0383   
0384   return dcaToVtx;
0385 }
0386 
0387 //======================================
0388 TLorentzVector getPairParent(TVector3 pos1, TVector3 mom1, int charge1, double mass1,
0389                  TVector3 pos2, TVector3 mom2, int charge2, double mass2,
0390                  TVector3 vtx,
0391                  float &dcaDaughters, float &cosTheta, float &decayLength, float &V0DcaToVtx)
0392 {
0393   // -- get helix  
0394   StPhysicalHelix p1Helix(mom1, pos1, bField * tesla, charge1);
0395   StPhysicalHelix p2Helix(mom2, pos2, bField * tesla, charge2);
0396 
0397   TVector3 vtx_tmp;
0398   vtx_tmp.SetXYZ(vtx.x()*millimeter, vtx.y()*millimeter, vtx.z()*millimeter);
0399 
0400   // -- find two-track DCA point 
0401   pair<double, double> const ss = p1Helix.pathLengths(p2Helix);
0402   TVector3 const p1AtDcaToP2 = p1Helix.at(ss.first);
0403   TVector3 const p2AtDcaToP1 = p2Helix.at(ss.second);
0404   
0405   // -- calculate DCA of particle1 to particle2 at their DCA
0406   dcaDaughters = (p1AtDcaToP2 - p2AtDcaToP1).Mag()/millimeter;
0407     
0408   // -- calculate Lorentz vector of particle1-particle2 pair
0409   TVector3 const p1MomAtDca = p1Helix.momentumAt(ss.first,  bField * tesla);
0410   TVector3 const p2MomAtDca = p2Helix.momentumAt(ss.second, bField * tesla);
0411   
0412   TLorentzVector p1FourMom(p1MomAtDca, sqrt(p1MomAtDca.Mag2()+mass1*mass2));
0413   TLorentzVector p2FourMom(p2MomAtDca, sqrt(p2MomAtDca.Mag2()+mass2*mass2));
0414   
0415   TLorentzVector parent = p1FourMom + p2FourMom;
0416     
0417   // -- calculate decay vertex (secondary or tertiary)
0418   TVector3 decayVertex = (p1AtDcaToP2 + p2AtDcaToP1) * 0.5 ;
0419     
0420   // -- calculate pointing angle and decay length with respect to primary vertex
0421   TVector3 vtxToV0 = decayVertex - vtx_tmp;
0422   float pointingAngle = vtxToV0.Angle(parent.Vect());
0423   cosTheta = std::cos(pointingAngle);
0424   decayLength = vtxToV0.Mag()/millimeter;
0425 
0426   // -- calculate V0 DCA to primary vertex
0427   V0DcaToVtx = decayLength * std::sin(pointingAngle);
0428 
0429   return parent;
0430 }
0431