Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-02-22 11:02:47

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<float> mcMomPx = {treereader, "MCParticles.momentum.x"};
0153   TTreeReaderArray<float> mcMomPy = {treereader, "MCParticles.momentum.y"};
0154   TTreeReaderArray<float> 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   
0162   TTreeReaderArray<float> rcMomPx = {treereader, "ReconstructedChargedParticles.momentum.x"};
0163   TTreeReaderArray<float> rcMomPy = {treereader, "ReconstructedChargedParticles.momentum.y"};
0164   TTreeReaderArray<float> rcMomPz = {treereader, "ReconstructedChargedParticles.momentum.z"};
0165   TTreeReaderArray<float> rcPosx = {treereader, "ReconstructedChargedParticles.referencePoint.x"};
0166   TTreeReaderArray<float> rcPosy = {treereader, "ReconstructedChargedParticles.referencePoint.y"};
0167   TTreeReaderArray<float> rcPosz = {treereader, "ReconstructedChargedParticles.referencePoint.z"};
0168   TTreeReaderArray<float> rcCharge = {treereader, "ReconstructedChargedParticles.charge"};
0169   TTreeReaderArray<int>   rcPdg = {treereader, "ReconstructedChargedParticles.PDG"};
0170 
0171   TTreeReaderArray<float> rcTrkLoca = {treereader, "CentralCKFTrackParameters.loc.a"};
0172   TTreeReaderArray<float> rcTrkLocb = {treereader, "CentralCKFTrackParameters.loc.b"};
0173   TTreeReaderArray<float> rcTrkqOverP = {treereader, "CentralCKFTrackParameters.qOverP"};
0174   TTreeReaderArray<float> rcTrkTheta = {treereader, "CentralCKFTrackParameters.theta"};
0175   TTreeReaderArray<float> rcTrkPhi = {treereader, "CentralCKFTrackParameters.phi"};
0176 
0177   TTreeReaderArray<float> CTVx = {treereader, "CentralTrackVertices.position.x"};
0178   TTreeReaderArray<float> CTVy = {treereader, "CentralTrackVertices.position.y"};
0179   TTreeReaderArray<float> CTVz = {treereader, "CentralTrackVertices.position.z"};
0180   TTreeReaderArray<int> CTVndf = {treereader, "CentralTrackVertices.ndf"};
0181   TTreeReaderArray<float> CTVchi2 = {treereader, "CentralTrackVertices.chi2"};
0182   TTreeReaderArray<float> CTVerr_xx = {treereader, "CentralTrackVertices.positionError.xx"};
0183   TTreeReaderArray<float> CTVerr_yy = {treereader, "CentralTrackVertices.positionError.yy"};
0184   TTreeReaderArray<float> CTVerr_zz = {treereader, "CentralTrackVertices.positionError.zz"};
0185 
0186   TTreeReaderArray<int> prim_vtx_index = {treereader, "PrimaryVertices_objIdx.index"};
0187 
0188   TTreeReaderArray<unsigned int> vtxAssocPart_begin = {treereader, "CentralTrackVertices.associatedParticles_begin"};
0189   TTreeReaderArray<unsigned int> vtxAssocPart_end = {treereader, "CentralTrackVertices.associatedParticles_end"};
0190   TTreeReaderArray<int> vtxAssocPart_index = {treereader, "_CentralTrackVertices_associatedParticles.index"};
0191   
0192   int nevents = 0;
0193   while(treereader.Next())
0194     {
0195       if(nevents%1000==0) printf("\n[i] New event %d\n",nevents);
0196 
0197       // ===== find MC primary vertex
0198       int nMCPart = mcPartMass.GetSize();
0199       TVector3 vertex_mc(-999., -999., -999.);
0200       for(int imc=0; imc<nMCPart; imc++)
0201     {
0202       if(mcPartGenStatus[imc] == 4 && mcPartPdg[imc] == 11)
0203         {
0204           vertex_mc.SetXYZ(mcEndPointX[imc], mcEndPointY[imc], mcEndPointZ[imc]);
0205           break;
0206         }
0207     }
0208       hEventStat->Fill(0.5);
0209       hMcVtxX->Fill(vertex_mc.x());
0210       hMcVtxY->Fill(vertex_mc.y());
0211       hMcVtxZ->Fill(vertex_mc.z());
0212 
0213       // ===== get RC primary vertex
0214       TVector3 vertex_rc(-999., -999., -999.);
0215       if(prim_vtx_index.GetSize()>0)
0216     {
0217       int rc_vtx_index = prim_vtx_index[0];
0218       vertex_rc.SetXYZ(CTVx[rc_vtx_index], CTVy[rc_vtx_index], CTVz[rc_vtx_index]);
0219     }
0220 
0221       // ===== map MC and RC particles
0222       int nAssoc = assocChRecID.GetSize();
0223       map<int, int> assoc_map_to_rc;
0224       map<int, int> assoc_map_to_mc;
0225       for(int j=0; j<nAssoc; j++)
0226     {
0227       assoc_map_to_rc[assocChSimID[j]] = assocChRecID[j];
0228       assoc_map_to_mc[assocChRecID[j]] = assocChSimID[j];
0229     }
0230 
0231       // ===== look for D0 -> pi+K
0232       for(int imc=0; imc<nMCPart; imc++)
0233     {
0234       if(fabs(mcPartPdg[imc]) != 421) continue;
0235       hEventStat->Fill(1.5);
0236       
0237       int nDuaghters = mcPartDaughter_end[imc]-mcPartDaughter_begin[imc];
0238       if(nDuaghters!=2) continue;
0239 
0240       // find D0 that decay into pi+K
0241       bool is_pik_decay = false;      
0242       int daug_index_1 = mcPartDaughter_index[mcPartDaughter_begin[imc]];
0243       int daug_index_2 = mcPartDaughter_index[mcPartDaughter_begin[imc]+1];
0244       int daug_pdg_1 = mcPartPdg[daug_index_1];
0245       int daug_pdg_2 = mcPartPdg[daug_index_2];
0246       if( (fabs(daug_pdg_1)==321 && fabs(daug_pdg_2)==211) || (fabs(daug_pdg_1)==211 && fabs(daug_pdg_2)==321) )
0247         {
0248           is_pik_decay = true;
0249         }
0250       if(!is_pik_decay) continue;
0251       hEventStat->Fill(2.5);
0252 
0253       // D0 kinematics
0254       TLorentzVector mc_mom_vec;
0255       mc_mom_vec.SetXYZM(mcMomPx[imc], mcMomPy[imc], mcMomPz[imc], mcPartMass[imc]);
0256       double mcRap = mc_mom_vec.Rapidity();
0257       double mcPt = mc_mom_vec.Pt();
0258       hMCD0PtRap->Fill(mcRap, mcPt);
0259     }
0260 
0261 
0262       // ===== Get reconstructed pions and kaons
0263       const int pid_mode = 1; // 0 - truth; 1 - realistic
0264       vector<unsigned int> pi_index;
0265       vector<unsigned int> k_index;
0266       pi_index.clear();
0267       k_index.clear();
0268       for(unsigned int rc_index=0; rc_index<rcMomPx.GetSize(); rc_index++)
0269     {     
0270       if(pid_mode==0)
0271         {
0272           // Use truth PID information
0273           int iSimPartID = -1;
0274           if(assoc_map_to_mc.find(rc_index) != assoc_map_to_mc.end()) iSimPartID = assoc_map_to_mc[rc_index];
0275           if(iSimPartID>=0)
0276         {
0277           if(fabs(mcPartPdg[iSimPartID]) == 211) pi_index.push_back(rc_index);
0278           if(fabs(mcPartPdg[iSimPartID]) == 321) k_index.push_back(rc_index);
0279         }
0280         }
0281       else if(pid_mode==1)
0282         {
0283           // Use reconstructed PID
0284           if(fabs(rcPdg[rc_index]) == 211) pi_index.push_back(rc_index);
0285           if(fabs(rcPdg[rc_index]) == 321) k_index.push_back(rc_index);
0286         }
0287     }
0288 
0289       // ===== pair pion and kaon
0290       for(unsigned int i=0; i<pi_index.size(); i++)
0291     {
0292       // convert from local to global coordinate
0293       // Helix code uses centimeter as the default unit, while ePIC software uses millimeter
0294       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);
0295       TVector3 mom1(rcMomPx[pi_index[i]], rcMomPy[pi_index[i]], rcMomPz[pi_index[i]]);
0296       int charge1 = rcCharge[pi_index[i]];
0297       
0298       TVector3 dcaToVtx = getDcaToVtx(pos1, mom1, charge1, vertex_rc);
0299       for(unsigned int j=0; j<k_index.size(); j++)
0300         {
0301           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);
0302           TVector3 mom2(rcMomPx[k_index[j]], rcMomPy[k_index[j]], rcMomPz[k_index[j]]);
0303           int charge2 = rcCharge[k_index[j]];
0304       
0305           TVector3 dcaToVtx2 = getDcaToVtx(pos2, mom2, charge2, vertex_rc);
0306           if(charge1*charge2<0)
0307         {
0308           // -- only look at unlike-sign pi+k pair
0309 
0310           float dcaDaughters, cosTheta, decayLength, V0DcaToVtx;
0311           TLorentzVector parent = getPairParent(pos1, mom1, charge1, gPionMass,
0312                             pos2, mom2, charge2, gKaonMass,
0313                             vertex_rc,
0314                             dcaDaughters, cosTheta, decayLength, V0DcaToVtx);
0315           h3InvMass->Fill(parent.Pt(), parent.Rapidity(), parent.M());
0316         }
0317         }
0318     }         
0319       nevents++;
0320     }
0321 
0322   TFile *outfile = new TFile(outname.Data(), "recreate");
0323 
0324   hEventStat->Write();
0325   hMcVtxX->Write();
0326   hMcVtxY->Write();
0327   hMcVtxZ->Write();
0328   hMCD0PtRap->Write();
0329   h3InvMass->Write();
0330   outfile->Close();
0331 
0332 }
0333 
0334 //======================================
0335 TVector3 getDcaToVtx(TVector3 pos, TVector3 mom, int charge, TVector3 vtx)
0336 {
0337   StPhysicalHelix pHelix(mom, pos, bField * tesla, charge);
0338 
0339   TVector3 vtx_tmp;
0340   vtx_tmp.SetXYZ(vtx.x()*millimeter, vtx.y()*millimeter, vtx.z()*millimeter);
0341   
0342   pHelix.moveOrigin(pHelix.pathLength(vtx_tmp));
0343   TVector3 dcaToVtx = pHelix.origin() - vtx_tmp;
0344 
0345   dcaToVtx.SetXYZ(dcaToVtx.x()/millimeter, dcaToVtx.y()/millimeter, dcaToVtx.z()/millimeter);
0346   
0347   return dcaToVtx;
0348 }
0349 
0350 //======================================
0351 TLorentzVector getPairParent(TVector3 pos1, TVector3 mom1, int charge1, double mass1,
0352                  TVector3 pos2, TVector3 mom2, int charge2, double mass2,
0353                  TVector3 vtx,
0354                  float &dcaDaughters, float &cosTheta, float &decayLength, float &V0DcaToVtx)
0355 {
0356   // -- get helix  
0357   StPhysicalHelix p1Helix(mom1, pos1, bField * tesla, charge1);
0358   StPhysicalHelix p2Helix(mom2, pos2, bField * tesla, charge2);
0359 
0360   TVector3 vtx_tmp;
0361   vtx_tmp.SetXYZ(vtx.x()*millimeter, vtx.y()*millimeter, vtx.z()*millimeter);
0362 
0363   // -- find two-track DCA point 
0364   pair<double, double> const ss = p1Helix.pathLengths(p2Helix);
0365   TVector3 const p1AtDcaToP2 = p1Helix.at(ss.first);
0366   TVector3 const p2AtDcaToP1 = p2Helix.at(ss.second);
0367   
0368   // -- calculate DCA of particle1 to particle2 at their DCA
0369   dcaDaughters = (p1AtDcaToP2 - p2AtDcaToP1).Mag()/millimeter;
0370     
0371   // -- calculate Lorentz vector of particle1-particle2 pair
0372   TVector3 const p1MomAtDca = p1Helix.momentumAt(ss.first,  bField * tesla);
0373   TVector3 const p2MomAtDca = p2Helix.momentumAt(ss.second, bField * tesla);
0374   
0375   TLorentzVector p1FourMom(p1MomAtDca, sqrt(p1MomAtDca.Mag2()+mass1*mass2));
0376   TLorentzVector p2FourMom(p2MomAtDca, sqrt(p2MomAtDca.Mag2()+mass2*mass2));
0377   
0378   TLorentzVector parent = p1FourMom + p2FourMom;
0379     
0380   // -- calculate decay vertex (secondary or tertiary)
0381   TVector3 decayVertex = (p1AtDcaToP2 + p2AtDcaToP1) * 0.5 ;
0382     
0383   // -- calculate pointing angle and decay length with respect to primary vertex
0384   TVector3 vtxToV0 = decayVertex - vtx_tmp;
0385   float pointingAngle = vtxToV0.Angle(parent.Vect());
0386   cosTheta = std::cos(pointingAngle);
0387   decayLength = vtxToV0.Mag()/millimeter;
0388 
0389   // -- calculate V0 DCA to primary vertex
0390   V0DcaToVtx = decayLength * std::sin(pointingAngle);
0391 
0392   return parent;
0393 }
0394