Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2026-07-29 09:31:53

0001 #include <cmath>
0002 #include <fstream>
0003 #include <iostream>
0004 #include <string>
0005 #include <vector>
0006 #include <algorithm>
0007 #include <utility>
0008 
0009 #include "ROOT/RDataFrame.hxx"
0010 #include <TH1D.h>
0011 #include <TProfile2D.h>
0012 #include <TFitResult.h>
0013 #include <TRandom3.h>
0014 #include <TCanvas.h>
0015 #include <TSystem.h>
0016 #include "TFile.h"
0017 #include "TChain.h"
0018 #include "TTreeReader.h"
0019 #include "TTreeReaderArray.h"
0020 #include "TLorentzVector.h"
0021 #include "TLorentzRotation.h"
0022 #include "TVector2.h"
0023 #include "TVector3.h"
0024 
0025 #include "fmt/color.h"
0026 #include "fmt/core.h"
0027 
0028 #include "edm4eic/InclusiveKinematicsData.h"
0029 #include "edm4eic/ReconstructedParticleData.h"
0030 #include "edm4hep/MCParticleData.h"
0031 
0032 #define PI            3.1415926
0033 #define MASS_ELECTRON 0.00051
0034 #define MASS_PROTON   0.93827
0035 #define MASS_PION     0.13957
0036 #define MASS_KAON     0.493667
0037 #define MASS_AU197    183.45406466643374
0038 
0039 int uchannelrho(TString rec_file="input.root", TString outputfile="output.root"){   
0040     if (gSystem->AccessPathName(rec_file.Data()) != 0) {
0041        // File does not exist
0042        cout<<Form("File %s does not exist.", rec_file.Data())<<endl;
0043        return 0;
0044     }
0045     
0046     // read our configuration   
0047     auto tree = new TChain("events");
0048     TString name_of_input = (TString) rec_file;
0049     std::cout << "Input file = " << name_of_input << endl;
0050     tree->Add(name_of_input);
0051     TTreeReader tree_reader(tree);       // !the tree reader
0052     
0053     //MC-level track attributes
0054     TTreeReaderArray<int>    mc_genStatus_array = {tree_reader, "MCParticles.generatorStatus"};
0055     TTreeReaderArray<double> mc_px_array = {tree_reader, "MCParticles.momentum.x"};
0056     TTreeReaderArray<double> mc_py_array = {tree_reader, "MCParticles.momentum.y"};
0057     TTreeReaderArray<double> mc_pz_array = {tree_reader, "MCParticles.momentum.z"};
0058     TTreeReaderArray<int>    mc_pdg_array= {tree_reader, "MCParticles.PDG"};
0059 
0060     //Reco-level track attributes
0061     TTreeReaderArray<float> reco_px_array = {tree_reader, "ReconstructedChargedParticles.momentum.x"};
0062     TTreeReaderArray<float> reco_py_array = {tree_reader, "ReconstructedChargedParticles.momentum.y"};
0063     TTreeReaderArray<float> reco_pz_array = {tree_reader, "ReconstructedChargedParticles.momentum.z"};
0064     TTreeReaderArray<float> reco_charge_array = {tree_reader, "ReconstructedChargedParticles.charge"};
0065     TTreeReaderArray<int>   reco_type     = {tree_reader,"ReconstructedChargedParticles.type"};
0066     
0067     // The association simID/recID are now stored as podio relation branches.
0068     // _ReconstructedChargedParticleAssociations_sim.index is the index into the
0069     // MCParticles collection; _..._rec.index is the index into ReconstructedChargedParticles.
0070     TTreeReaderArray<int> rec_id = {tree_reader, "_ReconstructedChargedParticleAssociations_rec.index"};
0071     TTreeReaderArray<int> sim_id = {tree_reader, "_ReconstructedChargedParticleAssociations_sim.index"};
0072     
0073     TString output_name_dir = outputfile;
0074     cout << "Output file = " << output_name_dir << endl;
0075     TFile* output = new TFile(output_name_dir,"RECREATE");
0076     
0077     //Pion reconstruction efficiency histograms
0078     TProfile2D* h_effEtaPtPi  = new TProfile2D("h_effEtaPtPi",";#eta; p_{T}(GeV/c)",50,4,6,30,0,1.6);
0079     TProfile2D* h_effEtaPtPip = new TProfile2D("h_effEtaPtPip",";#eta; p_{T}(GeV/c)",50,4,6,30,0,1.6);
0080     TProfile2D* h_effEtaPtPim = new TProfile2D("h_effEtaPtPim",";#eta; p_{T}(GeV/c)",50,4,6,30,0,1.6);
0081     TProfile2D* h_effPhiEtaPi  = new TProfile2D("h_effPhiEtaPi",";#phi (rad);#eta",50,0,6.4,50,4,6);
0082     TProfile2D* h_effPhiEtaPip = new TProfile2D("h_effPhiEtaPip",";#phi (rad);#eta",50,0,6.4,50,4,6);
0083     TProfile2D* h_effPhiEtaPim = new TProfile2D("h_effPhiEtaPim",";#phi (rad);#eta",50,0,6.4,50,4,6);
0084     
0085     //rho vector meson mass
0086     TH1D* h_VM_mass_MC = new TH1D("h_VM_mass_MC",";mass (GeV)",200,0,4);
0087     TH1D* h_VM_mass_MC_etacut = new TH1D("h_VM_mass_MC_etacut",";mass (GeV)",200,0,4);
0088     TH1D* h_VM_mass_REC = new TH1D("h_VM_mass_REC",";mass (GeV)",200,0,4);
0089     TH1D* h_VM_mass_REC_etacut = new TH1D("h_VM_mass_REC_etacut",";mass (GeV)",200,0,4);
0090     TH1D* h_VM_mass_REC_justpions = new TH1D("h_VM_mass_REC_justpions",";mass (GeV)",200,0,4);
0091     
0092     cout<<"There are "<<tree->GetEntries()<<" events!!!!!!!"<<endl;
0093     tree_reader.SetEntriesRange(0, tree->GetEntries());
0094     while (tree_reader.Next()) {
0095     
0096         TLorentzVector vmMC(0,0,0,0);
0097         TLorentzVector piplusMC(0,0,0,0);
0098         TLorentzVector piminusMC(0,0,0,0);
0099         
0100         //MC level
0101         for(unsigned int imc=0;imc<mc_px_array.GetSize();imc++){
0102             TVector3 mctrk(mc_px_array[imc],mc_py_array[imc],mc_pz_array[imc]); 
0103             if(mc_pdg_array[imc]!=11) mctrk.RotateY(0.025);//Rotate all non-electrons to hadron beam coordinate system
0104             if(mc_genStatus_array[imc]!=1) continue;
0105             if(mc_pdg_array[imc]==211 && mc_genStatus_array[imc]==1){ 
0106               piplusMC.SetVectM(mctrk,MASS_PION);  
0107             }
0108           if(mc_pdg_array[imc]==-211 && mc_genStatus_array[imc]==1){ 
0109                 piminusMC.SetVectM(mctrk,MASS_PION); 
0110             }
0111         }
0112         
0113         vmMC=piplusMC+piminusMC;
0114         if(vmMC.E()!=0){
0115             h_VM_mass_MC->Fill(vmMC.M());
0116             if(piplusMC.Theta()>0.009  && piplusMC.Theta()<0.013 && 
0117                 piminusMC.Theta()>0.009 && piminusMC.Theta()<0.013 ) h_VM_mass_MC_etacut->Fill(vmMC.M());
0118         }
0119         
0120         //4-vector for proton reconstructed as if it were a pi+
0121         TLorentzVector protonRECasifpion(0,0,0,0);
0122         //pion 4-vectors
0123         TLorentzVector piplusREC(0,0,0,0);
0124         TLorentzVector piminusREC(0,0,0,0);
0125         //rho reconstruction 4-vector
0126         TLorentzVector vmREC(0,0,0,0);
0127         //rho reconstruction from mis-identified proton 4-vector
0128         TLorentzVector vmRECfromproton(0,0,0,0);
0129         
0130         bool isPiMinusFound = false;
0131         bool isPiPlusFound = false;
0132         bool isProtonFound = false;
0133         
0134         // Map each reconstructed-track index to the MC-particle index it is
0135         // associated with. The association is a separate collection: for entry ia
0136         // rec_id[ia] is the ReconstructedChargedParticles index and sim_id[ia] the
0137         // MCParticles index. We must look up by rec_id rather than assume the
0138         // association collection is aligned with the track collection.
0139         std::vector<int> simForRec(reco_pz_array.GetSize(), -1);
0140         for(unsigned int ia=0; ia<rec_id.GetSize(); ia++){
0141             int ri = rec_id[ia];
0142             if(ri>=0 && ri<(int)simForRec.size()) simForRec[ri] = sim_id[ia];
0143         }
0144 
0145         //track loop
0146         int numpositivetracks = 0;
0147         for(unsigned int itrk=0;itrk<reco_pz_array.GetSize();itrk++){
0148             TVector3 trk(reco_px_array[itrk],reco_py_array[itrk],reco_pz_array[itrk]);
0149 
0150             //  Rotate in order to account for crossing angle
0151             //  and express coordinates in hadron beam pipe frame
0152             //  This is just a patch, not a final solution.
0153             trk.RotateY(0.025);
0154 
0155             if(reco_type[itrk] == -1){
0156                 continue;
0157             }
0158 
0159             int thisSim = simForRec[itrk];
0160 
0161           if(reco_charge_array[itrk]>0){
0162                 numpositivetracks++;
0163               if ((thisSim==4 || thisSim==5) && reco_charge_array[itrk]==1){
0164                 piplusREC.SetVectM(trk,MASS_PION);
0165                 isPiPlusFound=true;
0166               }
0167              if(thisSim==6){
0168                 protonRECasifpion.SetVectM(trk,MASS_PION);
0169                 isProtonFound=true;
0170              }
0171             }
0172           if(reco_charge_array[itrk]<0){
0173             piminusREC.SetVectM(trk,MASS_PION);
0174             if((thisSim==4 || thisSim==5) && reco_charge_array[itrk]==-1)   isPiMinusFound=true;
0175           }
0176 
0177         }
0178         
0179         //4vector of VM;
0180         if(piplusREC.E()!=0. && piminusREC.E()!=0.){
0181             vmREC=piplusREC+piminusREC;
0182         }
0183         if(protonRECasifpion.E()!=0. && piminusREC.E()!=0.){
0184           vmRECfromproton=protonRECasifpion+piminusREC;
0185         }
0186         
0187         //pion reconstruction efficiency
0188         double thispipeff = (isPiPlusFound) ? 1.0 : 0.0;
0189         double thispimeff = (isPiMinusFound) ? 1.0 : 0.0;
0190         h_effEtaPtPi ->Fill(piplusMC.Eta(), piplusMC.Pt(), thispipeff);
0191         h_effEtaPtPi ->Fill(piminusMC.Eta(),piminusMC.Pt(),thispimeff);
0192         h_effEtaPtPip->Fill(piplusMC.Eta(), piplusMC.Pt(), thispipeff);
0193         h_effEtaPtPim->Fill(piminusMC.Eta(),piminusMC.Pt(),thispimeff);
0194         //
0195         double thispipphi = piplusMC.Phi()>0  ? piplusMC.Phi()  : piplusMC.Phi()+6.2831853;
0196         double thispimphi = piminusMC.Phi()>0 ? piminusMC.Phi() : piminusMC.Phi()+6.2831853;
0197         if(piplusMC.Pt() >0.2) h_effPhiEtaPi ->Fill(thispipphi, piplusMC.Eta(), thispipeff);
0198         if(piminusMC.Pt()>0.2) h_effPhiEtaPi ->Fill(thispimphi,piminusMC.Eta(),thispimeff);
0199         if(piplusMC.Pt() >0.2) h_effPhiEtaPip->Fill(thispipphi, piplusMC.Eta(), thispipeff);
0200         if(piminusMC.Pt()>0.2) h_effPhiEtaPim->Fill(thispimphi,piminusMC.Eta(),thispimeff);
0201         //
0202         
0203         //VM rec
0204         if(vmREC.E()==0 && vmRECfromproton.E()==0) continue;
0205         double rho_mass = vmREC.M();
0206         double rho_mass_fromproton = vmRECfromproton.M();
0207         h_VM_mass_REC->Fill(rho_mass);
0208         h_VM_mass_REC->Fill(rho_mass_fromproton);
0209         if(piplusMC.Theta()>0.009  && piplusMC.Theta()<0.013 &&
0210                           piminusMC.Theta()>0.009 && piminusMC.Theta()<0.013 ) h_VM_mass_REC_etacut->Fill(vmREC.M());
0211         if(isPiMinusFound && isPiPlusFound){ 
0212           h_VM_mass_REC_justpions->Fill(rho_mass);
0213         }
0214     }
0215     output->Write();
0216     output->Close();
0217     
0218     return 0;
0219 }