Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-09-27 07:02:39

0001 #include <fstream>
0002 #include <iostream>
0003 #include <string>
0004 
0005 #include <TChain.h>
0006 #include <TFile.h>
0007 #include <TTreeReader.h>
0008 #include <TTreeReaderArray.h>
0009 #include <TH1D.h>
0010 #include <TH2D.h>
0011 #include <Math/Vector3D.h>
0012 #include <Math/Vector4D.h>
0013 #include <Math/RotationY.h>
0014 #include <TMath.h>
0015 #include <TVector3.h>
0016 #include <TLorentzVector.h>
0017 
0018 #include "fmt/color.h"
0019 #include "fmt/core.h"
0020 
0021 #include "nlohmann/json.hpp"
0022 
0023 void trk_dis_analysis(const std::string& config_name)
0024 {
0025     // Read our configuration

0026     std::ifstream  config_file{config_name};
0027     nlohmann::json config;
0028     config_file >> config;
0029   
0030     const std::string rec_file      = config["rec_file"];
0031     const std::string detector      = config["detector"];
0032     const std::string output_prefix = config["output_prefix"];
0033     const int         ebeam         = config["ebeam"];
0034     const int         pbeam         = config["pbeam"];
0035     const int         Q2_min        = config["Min_Q2"];
0036     
0037     fmt::print(fmt::emphasis::bold | fg(fmt::color::forest_green),
0038                 "Running DIS tracking analysis...\n");
0039     fmt::print(" - Detector package: {}\n", detector);
0040     fmt::print(" - input file: {}\n", rec_file);
0041     fmt::print(" - output prefix for histograms: {}\n", output_prefix);
0042     fmt::print(" - ebeam: {}\n", ebeam);
0043     fmt::print(" - pbeam: {}\n", pbeam);
0044     fmt::print(" - Minimum Q2: {}\n", Q2_min);
0045     
0046     //--------------------------------------------------------------------------------------------------------------------------------------------

0047 
0048     // Set output file for the histograms

0049     std::string output_name_hists = fmt::format("{}.root", output_prefix);
0050     cout << "Output file for histograms = " << output_name_hists << endl;
0051     TFile* ofile = new TFile(output_name_hists.c_str(), "RECREATE");
0052 
0053     //--------------------------------------------------------------------------------------------------------------------------------------------

0054  
0055     // Set up input file chain

0056     TChain *mychain = new TChain("events");
0057     mychain->Add(rec_file.c_str());
0058 
0059     //--------------------------------------------------------------------------------------------------------------------------------------------

0060 
0061     // Initialize reader

0062     TTreeReader tr(mychain);
0063 
0064     // Generated Particle Information

0065     TTreeReaderArray<int>   gen_status(tr, "MCParticles.generatorStatus");
0066     TTreeReaderArray<int>   gen_pid(tr, "MCParticles.PDG");
0067     TTreeReaderArray<float> gen_px(tr, "MCParticles.momentum.x");
0068     TTreeReaderArray<float> gen_py(tr, "MCParticles.momentum.y");
0069     TTreeReaderArray<float> gen_pz(tr, "MCParticles.momentum.z");
0070     TTreeReaderArray<double> gen_mass(tr, "MCParticles.mass"); //Not important here

0071     TTreeReaderArray<float> gen_charge(tr, "MCParticles.charge");
0072     TTreeReaderArray<double> gen_vx(tr, "MCParticles.vertex.x");
0073     TTreeReaderArray<double> gen_vy(tr, "MCParticles.vertex.y");
0074     TTreeReaderArray<double> gen_vz(tr, "MCParticles.vertex.z");
0075 
0076     // Reconstructed real-seeded tracks (charged particles)

0077     TTreeReaderArray<float> rec_px(tr, "ReconstructedChargedParticles.momentum.x");
0078     TTreeReaderArray<float> rec_py(tr, "ReconstructedChargedParticles.momentum.y");
0079     TTreeReaderArray<float> rec_pz(tr, "ReconstructedChargedParticles.momentum.z");
0080     TTreeReaderArray<float> rec_mass(tr, "ReconstructedChargedParticles.mass");
0081     TTreeReaderArray<int> rec_type(tr, "ReconstructedChargedParticles.type"); //Type 0: successful eta/phi match to generated particle

0082     TTreeReaderArray<int> rec_pdg(tr, "ReconstructedChargedParticles.PDG"); //Uses PID lookup table information

0083 
0084     // Reconstructed truth-seeded tracks (charged particles)

0085     TTreeReaderArray<float> rec_ts_px(tr, "ReconstructedTruthSeededChargedParticles.momentum.x");
0086     TTreeReaderArray<float> rec_ts_py(tr, "ReconstructedTruthSeededChargedParticles.momentum.y");
0087     TTreeReaderArray<float> rec_ts_pz(tr, "ReconstructedTruthSeededChargedParticles.momentum.z");
0088     TTreeReaderArray<float> rec_ts_mass(tr, "ReconstructedTruthSeededChargedParticles.mass");
0089     TTreeReaderArray<int> rec_ts_type(tr, "ReconstructedTruthSeededChargedParticles.type"); //Type 0: successful eta/phi match to generated particle

0090     TTreeReaderArray<int> rec_ts_pdg(tr, "ReconstructedTruthSeededChargedParticles.PDG"); //Uses PID lookup table information

0091 
0092     // Hit-based track to MC Particle association weight

0093     TTreeReaderArray<float> hit_assoc_weight(tr,"CentralCKFTrackAssociations.weight"); //Real-seeded tracking

0094     TTreeReaderArray<float> hit_assoc_weight_ts(tr,"CentralCKFTruthSeededTrackAssociations.weight"); //Truth-seeded tracking

0095 
0096     //-------------------------------------------------------------------------------------------------------------------------------------------- 

0097     // Define Histograms

0098 
0099     //Eta distribution of generated charged particles

0100     TH1 *h1a = new TH1D("h1a","Generated Charged Particles",100,-4,4);
0101     h1a->GetXaxis()->SetTitle("#eta_{gen.}");h1a->GetXaxis()->CenterTitle();
0102     h1a->SetLineColor(kTeal);h1a->SetLineWidth(2);
0103 
0104     TH1 *h1a1 = new TH1D("h1a1","Generated Charged Particles",100,-4,4);  //Minimum momentum cut of Pt > 200 MeV/c

0105     h1a1->GetXaxis()->SetTitle("#eta_{gen.}");h1a1->GetXaxis()->CenterTitle();
0106     h1a1->SetLineColor(kRed);h1a1->SetLineWidth(2);
0107 
0108     TH1 *h1a2 = new TH1D("h1a2","Generated Charged Particles",100,-4,4);  //Minimum momentum cut of Pt > 500 MeV/c

0109     h1a2->GetXaxis()->SetTitle("#eta_{gen.}");h1a2->GetXaxis()->CenterTitle();
0110     h1a2->SetLineColor(kBlack);h1a2->SetLineWidth(2);
0111     h1a2->SetFillColor(kBlack);h1a2->SetFillStyle(3244);
0112 
0113     //Eta distribution of reconstructed real-seeded tracks (charged particles)

0114     TH1 *h1b = new TH1D("h1b","Reconstructed Real-seeded tracks",100,-4,4);
0115     h1b->GetXaxis()->SetTitle("#eta_{rec.}");h1b->GetXaxis()->CenterTitle();
0116     h1b->SetLineColor(kGreen);h1b->SetLineWidth(2);
0117 
0118     TH1 *h1b1 = new TH1D("h1b1","Reconstructed Real-seeded tracks",100,-4,4); //Minimum momentum cut of Pt > 200 MeV/c

0119     h1b1->GetXaxis()->SetTitle("#eta_{rec.}");h1b1->GetXaxis()->CenterTitle();
0120     h1b1->SetLineColor(kBlue);h1b1->SetLineWidth(2);
0121 
0122     TH1 *h1b2 = new TH1D("h1b2","Reconstructed Real-seeded tracks",100,-4,4); //Minimum momentum cut of Pt > 500 MeV/c

0123     h1b2->GetXaxis()->SetTitle("#eta_{rec.}");h1b2->GetXaxis()->CenterTitle();
0124     h1b2->SetLineColor(kRed);h1b2->SetLineWidth(2);
0125     h1b2->SetMarkerColor(kRed);h1b2->SetMarkerStyle(kFullCrossX);
0126 
0127     //Eta distribution of reconstructed truth-seeded tracks (charged particles)

0128     TH1 *h1c = new TH1D("h1c","Reconstructed Truth-seeded tracks",100,-4,4);
0129     h1c->GetXaxis()->SetTitle("#eta_{rec.}");h1c->GetXaxis()->CenterTitle();
0130     h1c->SetLineColor(kRed);h1c->SetLineWidth(2);
0131 
0132     TH1 *h1c1 = new TH1D("h1c1","Reconstructed Truth-seeded tracks",100,-4,4); //Minimum momentum cut of Pt > 200 MeV/c

0133     h1c1->GetXaxis()->SetTitle("#eta_{rec.}");h1c1->GetXaxis()->CenterTitle();
0134     h1c1->SetLineColor(kOrange);h1c1->SetLineWidth(2);
0135 
0136     TH1 *h1c2 = new TH1D("h1c2","Reconstructed Truth-seeded tracks",100,-4,4); //Minimum momentum cut of Pt > 500 MeV/c

0137     h1c2->GetXaxis()->SetTitle("#eta_{rec.}");h1c2->GetXaxis()->CenterTitle();
0138     h1c2->SetLineColor(kMagenta);h1c2->SetLineWidth(2);
0139     h1c2->SetMarkerColor(kMagenta);h1c2->SetMarkerStyle(kFullCrossX);
0140 
0141     //Track purity

0142     TH1 *h2a = new TH1D("h2a","Real-seeded tracks purity",100,0,1.1);
0143     h2a->GetXaxis()->SetTitle("Fraction of track measurements from a given MC Particle");h2a->GetXaxis()->CenterTitle();
0144     h2a->GetYaxis()->SetTitle("Number of Tracks");h2a->GetYaxis()->CenterTitle();
0145     h2a->SetLineColor(kBlack);h2a->SetLineWidth(2);
0146 
0147     TH1 *h2b = new TH1D("h2b","Truth-seeded tracks purity",100,0,1.1);
0148     h2b->GetXaxis()->SetTitle("Fraction of track measurements from a given MC Particle");h2b->GetXaxis()->CenterTitle();
0149     h2b->GetYaxis()->SetTitle("Number of Tracks");h2b->GetYaxis()->CenterTitle();
0150     h2b->SetLineColor(kBlack);h2b->SetLineWidth(2);
0151 
0152     //Define additional variables

0153     TLorentzVector gen_vec;
0154     TVector3 gen_vertex;
0155 
0156     TLorentzVector rec_vec;
0157     TVector3 track_vec; //Reconstructed track momentum vector

0158     int counter(0);
0159 
0160     //Loop over events

0161     std::cout<<"Analyzing "<<mychain->GetEntries()<<" events!"<<std::endl;
0162     while (tr.Next()) {
0163 
0164         if(counter%100==0) std::cout<<"Analyzing event "<<counter<<std::endl;
0165         counter++;
0166 
0167         //Loop over generated particles

0168         for(size_t igen=0;igen<gen_status.GetSize();igen++){
0169             
0170             auto charge = gen_charge[igen];
0171             auto status = gen_status[igen];
0172 
0173             //Require final-state, charged particle (no secondaries)

0174             if(status==1 && fabs(charge) > 0.01 ){
0175                 gen_vec.SetXYZM(gen_px[igen],gen_py[igen],gen_pz[igen],gen_mass[igen]);
0176                 gen_vertex.SetXYZ(gen_vx[igen],gen_vy[igen],gen_vz[igen]);
0177                 
0178                 //Fill eta histogram

0179                 h1a->Fill(gen_vec.Eta());
0180                 if( gen_vec.Pt()>0.2 ) h1a1->Fill(gen_vec.Eta());
0181                 if( gen_vec.Pt()>0.5 ) h1a2->Fill(gen_vec.Eta());
0182             }
0183         } //End loop over generated particles

0184         
0185         //Loop over reconstructed real-seeded charged particles (copy of tracks with PID info)

0186         size_t rec_mult = rec_type.GetSize();
0187 
0188         for(size_t irec=0;irec<rec_mult;irec++){
0189 
0190             rec_vec.SetXYZM(rec_px[irec],rec_py[irec],rec_pz[irec],rec_mass[irec]);
0191             
0192             //Fill histograms

0193             h1b->Fill(rec_vec.Eta());
0194             if( rec_vec.Pt() > 0.2 ) h1b1->Fill(rec_vec.Eta());
0195             if( rec_vec.Pt() > 0.5 ) h1b2->Fill(rec_vec.Eta());
0196 
0197         } //End loop over reconstructed particles

0198 
0199         //Loop over reconstructed truth-seeded charged particles (copy of tracks with PID info)

0200         size_t rec_ts_mult = rec_ts_type.GetSize();
0201 
0202         for(size_t irec=0;irec<rec_ts_mult;irec++){
0203 
0204             rec_vec.SetXYZM(rec_ts_px[irec],rec_ts_py[irec],rec_ts_pz[irec],rec_ts_mass[irec]);
0205             
0206             //Fill histograms

0207             h1c->Fill(rec_vec.Eta());
0208             if( rec_vec.Pt() > 0.2 ) h1c1->Fill(rec_vec.Eta());
0209             if( rec_vec.Pt() > 0.5 ) h1c2->Fill(rec_vec.Eta());
0210 
0211         } //End loop over reconstructed particles

0212 
0213         // Loop over truth-seeded hit-based associations

0214         for(size_t iassoc=0;iassoc<hit_assoc_weight.GetSize();iassoc++){
0215             
0216             auto assoc_weight = hit_assoc_weight[iassoc];
0217             h2a->Fill(assoc_weight);
0218         
0219         } //End loop over hit associations

0220 
0221         // Loop over truth-seeded hit-based associations

0222         for(size_t iassoc=0;iassoc<hit_assoc_weight_ts.GetSize();iassoc++){
0223             
0224             auto assoc_weight = hit_assoc_weight_ts[iassoc];
0225             h2b->Fill(assoc_weight);
0226         
0227         } //End loop over hit associations

0228 
0229     } //End loop over events

0230 
0231     //--------------------------------------------------------------------------------------------------------------------------------------------

0232 
0233     ofile->Write(); // Write histograms to file

0234     ofile->Close(); // Close output file

0235 
0236 }