File indexing completed on 2025-11-03 09:01:32
0001 
0002 
0003 #include <iostream>
0004 #include "ROOT/RDataFrame.hxx"
0005 #include "ROOT/RDF/RInterface.hxx"
0006 #include "ROOT/RVec.hxx"
0007 #include "edm4hep/MCParticleCollection.h"
0008 #include "edm4eic/ReconstructedParticleCollection.h"
0009 #include "TCanvas.h"
0010 #include "TStyle.h"
0011 
0012 void reconstructionAnalysis(TString inFile             = "/home/simong/EIC/detector_benchmarks_anl/sim_output/lowq2_reconstruction/analysis/Low-Q2_retrained_Particles_new.eicrecon.edm4hep.root",
0013                             float   beamEnergy         = 18.0,
0014                             TString outFile            = "reconstruction_results.root",
0015                             TString momentumCanvasName = "momentum_resolution.png",
0016                             TString energyThetaPhiCanvasName = "energy_theta_phi_resolution.png",
0017                             TString relationCanvasName = "relation_resolution.png",
0018                             TString resolutionGraphsCanvasName = "resolution_graphs.png",
0019                             std::string particleCollectionName = "TaggerTrackerReconstructedParticles") {
0020 
0021     
0022     gStyle->SetPadLeftMargin(0.1);  
0023     gStyle->SetPadRightMargin(0.0); 
0024     gStyle->SetPadTopMargin(0.0);   
0025     gStyle->SetPadBottomMargin(0.1);
0026     gStyle->SetTitleAlign(13);
0027     gStyle->SetTitleX(0.12);          
0028     gStyle->SetTitleY(0.985);         
0029     gStyle->SetTitleSize(0.08, "t");
0030     gStyle->SetTitleXOffset(1.0);    
0031     gStyle->SetTitleYOffset(1.0);    
0032     gStyle->SetOptStat(0);
0033 
0034     ROOT::RDataFrame d0("events",inFile, {"MCParticles",particleCollectionName});
0035 
0036     auto filterDF = d0.Define("SimParticles", "MCParticles[MCParticles.generatorStatus==1 && MCParticles.PDG==11]")
0037                       .Filter("SimParticles.size()==1")
0038                       .Filter(particleCollectionName+".size()==1");
0039 
0040     
0041     auto momentumDF = filterDF
0042         .Define("reco_momentum", particleCollectionName+"[0].momentum")
0043         .Define("mc_momentum", "SimParticles[0].momentum")
0044         .Define("px_rec", "reco_momentum.x")
0045         .Define("py_rec", "reco_momentum.y")
0046         .Define("pz_rec", "reco_momentum.z")
0047         .Define("px_mc", "mc_momentum.x")
0048         .Define("py_mc", "mc_momentum.y")
0049         .Define("pz_mc", "mc_momentum.z")
0050         
0051         .Define("theta_rec", "std::atan2(std::sqrt(px_rec*px_rec + py_rec*py_rec), pz_rec)")
0052         .Define("theta_mc", "std::atan2(std::sqrt(px_mc*px_mc + py_mc*py_mc), pz_mc)")
0053         .Define("phi_rec_rad", "std::atan2(py_rec, px_rec)")
0054         .Define("phi_mc_rad", "std::atan2(py_mc, px_mc)")
0055         .Define("phi_rec", "TMath::RadToDeg()*phi_rec_rad")
0056         .Define("phi_mc", "TMath::RadToDeg()*phi_mc_rad")
0057         
0058         .Define("px_diff", "(px_rec - px_mc)")
0059         .Define("py_diff", "(py_rec - py_mc)")
0060         .Define("pz_res", "(pz_rec - pz_mc)/pz_mc")
0061         .Define("E_rec", particleCollectionName+"[0].energy")
0062         .Define("E_mc", "std::sqrt(px_mc*px_mc + py_mc*py_mc + pz_mc*pz_mc + SimParticles[0].mass*SimParticles[0].mass)")
0063         .Define("E_res", "(E_rec - E_mc)/E_mc")
0064         .Define("theta_diff", "(theta_rec - theta_mc)")
0065         .Define("phi_diff", "TMath::RadToDeg()*ROOT::VecOps::DeltaPhi(phi_rec_rad, phi_mc_rad)");
0066 
0067     
0068     std::cout << "Original DataFrame size: " << d0.Count().GetValue() << std::endl;
0069     
0070     std::cout << "Filtered DataFrame size: " << filterDF.Count().GetValue() << std::endl;
0071 
0072     int   momentumBins      = 100;
0073     float momentumXRange[2] = {-0.1, 0.1};
0074     float momentumYRange[2] = {-0.1, 0.1};
0075     float momentumZRange[2] = {-beamEnergy, 0};
0076 
0077     int   momentumResolutionBins      = 100;
0078     float momentumDifferenceXRange[2] = {-0.05, 0.05};
0079     float momentumDifferenceYRange[2] = {-0.05, 0.05};
0080     float momentumResolutionZRange[2] = {-0.1, 0.1};
0081 
0082     int   energyBins      = 100;
0083     int   thetaBins       = 100;
0084     int   phiBins         = 100;
0085     float energyRange[2]  = {3.5, beamEnergy}; 
0086     float thetaRange[2]   = {3.134, TMath::Pi()}; 
0087     float phiRange[2]     = {-180, 180}; 
0088 
0089     int   resolutionBins           = 100;
0090     float energyResolutionRange[2] = {-0.1, 0.1};
0091     float thetaResolutionRange[2]  = {-0.003, 0.003};
0092     float phiResolutionRange[2]    = {-90, 90}; 
0093 
0094     
0095     auto px_Hist = momentumDF.Histo2D({"px_vs_px", "Reconstructed vs MC px; px reconstructed [GeV]; px MC [GeV]", momentumBins, momentumXRange[0], momentumXRange[1], momentumBins, momentumXRange[0], momentumXRange[1]}, "px_rec", "px_mc");
0096     auto py_Hist = momentumDF.Histo2D({"py_vs_py", "Reconstructed vs MC py; py reconstructed [GeV]; py MC [GeV]", momentumBins, momentumYRange[0], momentumYRange[1], momentumBins, momentumYRange[0], momentumYRange[1]}, "py_rec", "py_mc");
0097     auto pz_Hist = momentumDF.Histo2D({"pz_vs_pz", "Reconstructed vs MC pz; pz reconstructed [GeV]; pz MC [GeV]", momentumBins, momentumZRange[0], momentumZRange[1], momentumBins, momentumZRange[0], momentumZRange[1]}, "pz_rec", "pz_mc");
0098 
0099     
0100     auto px_diff_Hist = momentumDF.Histo1D({"px_diff", "px difference; px difference [GeV]; Entries", momentumResolutionBins, momentumDifferenceXRange[0], momentumDifferenceXRange[1]}, "px_diff");
0101     auto py_diff_Hist = momentumDF.Histo1D({"py_diff", "py difference; py difference [GeV]; Entries", momentumResolutionBins, momentumDifferenceYRange[0], momentumDifferenceYRange[1]}, "py_diff");
0102     auto pz_res_Hist  = momentumDF.Histo1D({"pz_res",  "pz resolution; pz resolution [GeV]; Entries", momentumResolutionBins, momentumResolutionZRange[0], momentumResolutionZRange[1]}, "pz_res");
0103 
0104     
0105     auto E_Hist     = momentumDF.Histo2D({"E_vs_E",         "Reconstructed vs MC energy; E reconstructed [GeV]; E MC [GeV]",        energyBins, energyRange[0], energyRange[1], energyBins, energyRange[0], energyRange[1]}, "E_rec", "E_mc");
0106     auto theta_Hist = momentumDF.Histo2D({"theta_vs_theta", "Reconstructed vs MC theta; theta reconstructed [rad]; theta MC [rad]", thetaBins, thetaRange[0], thetaRange[1], thetaBins, thetaRange[0], thetaRange[1]}, "theta_rec", "theta_mc");
0107     auto phi_Hist   = momentumDF.Histo2D({"phi_vs_phi",     "Reconstructed vs MC phi; phi reconstructed [deg]; phi MC [deg]",       phiBins, phiRange[0], phiRange[1], phiBins, phiRange[0], phiRange[1]}, "phi_rec", "phi_mc");
0108 
0109     auto E_res_Hist      = momentumDF.Histo1D({"E_res", "E resolution; E resolution [GeV]; Entries", resolutionBins, energyResolutionRange[0], energyResolutionRange[1]}, "E_res");
0110     auto theta_diff_Hist = momentumDF.Histo1D({"theta_diff", "theta difference; theta difference [rad]; Entries", resolutionBins, thetaResolutionRange[0], thetaResolutionRange[1]}, "theta_diff");
0111     auto phi_diff_Hist   = momentumDF.Histo1D({"phi_diff", "phi difference; phi difference [deg]; Entries", resolutionBins, phiResolutionRange[0], phiResolutionRange[1]}, "phi_diff");
0112 
0113     
0114     auto E_res_vs_E_Hist          = momentumDF.Histo2D({"E_res_vs_E", "E resolution vs E reconstructed; E reconstructed [GeV]; E resolution [GeV]", energyBins, energyRange[0], energyRange[1], resolutionBins, energyResolutionRange[0], energyResolutionRange[1]}, "E_rec", "E_res");
0115     auto E_res_vs_theta_Hist      = momentumDF.Histo2D({"E_res_vs_theta", "E resolution vs theta reconstructed; theta reconstructed [rad]; E resolution [GeV]", thetaBins, thetaRange[0], thetaRange[1], resolutionBins, energyResolutionRange[0], energyResolutionRange[1]}, "theta_rec", "E_res");
0116     auto E_res_vs_phi_Hist        = momentumDF.Histo2D({"E_res_vs_phi", "E resolution vs phi reconstructed; phi reconstructed [deg]; E resolution [GeV]", phiBins, phiRange[0], phiRange[1], resolutionBins, energyResolutionRange[0], energyResolutionRange[1]}, "phi_rec", "E_res");
0117     auto theta_diff_vs_E_Hist     = momentumDF.Histo2D({"theta_diff_vs_E", "theta difference vs E reconstructed; E reconstructed [GeV]; theta difference [rad]", energyBins, energyRange[0], energyRange[1], resolutionBins, thetaResolutionRange[0], thetaResolutionRange[1]}, "E_rec", "theta_diff");
0118     auto theta_diff_vs_theta_Hist = momentumDF.Histo2D({"theta_diff_vs_theta", "theta difference vs theta reconstructed; theta reconstructed [rad]; theta difference [rad]", thetaBins, thetaRange[0], thetaRange[1], resolutionBins, thetaResolutionRange[0], thetaResolutionRange[1]}, "theta_rec", "theta_diff");
0119     auto theta_diff_vs_phi_Hist   = momentumDF.Histo2D({"theta_diff_vs_phi", "theta difference vs phi reconstructed; phi reconstructed [deg]; theta difference [rad]", phiBins, phiRange[0], phiRange[1], resolutionBins, thetaResolutionRange[0], thetaResolutionRange[1]}, "phi_rec", "theta_diff");
0120     auto phi_diff_vs_E_Hist       = momentumDF.Histo2D({"phi_diff_vs_E", "phi difference vs E reconstructed; E reconstructed [GeV]; phi difference [rad]", energyBins, energyRange[0], energyRange[1], resolutionBins, phiResolutionRange[0], phiResolutionRange[1]}, "E_rec", "phi_diff");
0121     auto phi_diff_vs_theta_Hist   = momentumDF.Histo2D({"phi_diff_vs_theta", "phi difference vs theta reconstructed; theta reconstructed [rad]; phi difference [deg]", thetaBins, thetaRange[0], thetaRange[1], resolutionBins, phiResolutionRange[0], phiResolutionRange[1]}, "theta_rec", "phi_diff");
0122     auto phi_diff_vs_phi_Hist     = momentumDF.Histo2D({"phi_diff_vs_phi", "phi difference vs phi reconstructed; phi reconstructed [deg]; phi difference [deg]", phiBins, phiRange[0], phiRange[1], resolutionBins, phiResolutionRange[0], phiResolutionRange[1]}, "phi_rec", "phi_diff");
0123 
0124     
0125     TCanvas *cMomentum = new TCanvas("momentum_canvas", "Momentum Resolution", 3000, 1600);
0126     cMomentum->Divide(3, 2);
0127     cMomentum->cd(1);
0128     px_Hist->Draw("colz");
0129     gPad->SetLogz();
0130     cMomentum->cd(2);
0131     py_Hist->Draw("colz");
0132     gPad->SetLogz();
0133     cMomentum->cd(3);
0134     pz_Hist->Draw("colz");
0135     gPad->SetLogz();
0136     cMomentum->cd(4);
0137     px_diff_Hist->Draw();
0138     cMomentum->cd(5);
0139     py_diff_Hist->Draw();
0140     cMomentum->cd(6);
0141     pz_res_Hist->Draw();
0142     cMomentum->SetGrid();
0143     cMomentum->Update();
0144     
0145     cMomentum->SaveAs(momentumCanvasName);
0146 
0147     
0148     TCanvas *cEnergyThetaPhi = new TCanvas("energy_theta_phi_canvas", "Energy, Theta and Phi Resolution", 3000, 1600);
0149     cEnergyThetaPhi->Divide(3, 2);
0150     cEnergyThetaPhi->cd(1);
0151     E_Hist->Draw("colz");
0152     gPad->SetLogz();
0153     cEnergyThetaPhi->cd(2);
0154     theta_Hist->Draw("colz");
0155     gPad->SetLogz();
0156     cEnergyThetaPhi->cd(3);
0157     phi_Hist->Draw("colz");
0158     gPad->SetLogz();
0159     cEnergyThetaPhi->cd(4);
0160     E_res_Hist->Draw();
0161     cEnergyThetaPhi->cd(5);
0162     theta_diff_Hist->Draw();
0163     cEnergyThetaPhi->cd(6);
0164     phi_diff_Hist->Draw();
0165     cEnergyThetaPhi->SetGrid();
0166     cEnergyThetaPhi->Update();
0167     
0168     cEnergyThetaPhi->SaveAs(energyThetaPhiCanvasName);
0169     
0170     
0171     TCanvas *cResolutionVsMC = new TCanvas("resolution_vs_mc_canvas", "Resolution vs MC Values", 3000, 1600);
0172     cResolutionVsMC->Divide(3, 3);
0173     cResolutionVsMC->cd(1);
0174     E_res_vs_E_Hist->Draw("colz");
0175     gPad->SetLogz();
0176     cResolutionVsMC->cd(2);
0177     E_res_vs_theta_Hist->Draw("colz");
0178     gPad->SetLogz();
0179     cResolutionVsMC->cd(3);
0180     E_res_vs_phi_Hist->Draw("colz");
0181     gPad->SetLogz();
0182     cResolutionVsMC->cd(4);
0183     theta_diff_vs_E_Hist->Draw("colz");
0184     gPad->SetLogz();
0185     cResolutionVsMC->cd(5);
0186     theta_diff_vs_theta_Hist->Draw("colz");
0187     gPad->SetLogz();
0188     cResolutionVsMC->cd(6);
0189     theta_diff_vs_phi_Hist->Draw("colz");
0190     gPad->SetLogz();
0191     cResolutionVsMC->cd(7);
0192     phi_diff_vs_E_Hist->Draw("colz");
0193     gPad->SetLogz();
0194     cResolutionVsMC->cd(8);
0195     phi_diff_vs_theta_Hist->Draw("colz");
0196     gPad->SetLogz();
0197     cResolutionVsMC->cd(9);
0198     phi_diff_vs_phi_Hist->Draw("colz");
0199     gPad->SetLogz();
0200     cResolutionVsMC->SetGrid();
0201     cResolutionVsMC->Update();
0202     
0203     cResolutionVsMC->SaveAs(relationCanvasName);
0204 
0205     
0206     TObjArray* fitE_vs_E = new TObjArray();
0207     TObjArray* fitTheta_vs_E = new TObjArray();
0208     TObjArray* fitPhi_vs_theta = new TObjArray();
0209     E_res_vs_E_Hist->FitSlicesY(nullptr, 1, -1, 0, "Q", fitE_vs_E);
0210     theta_diff_vs_E_Hist->FitSlicesY(nullptr, 1, -1, 0, "Q", fitTheta_vs_E);
0211     phi_diff_vs_theta_Hist->FitSlicesY(nullptr, 1, -1, 0, "Q", fitPhi_vs_theta);
0212 
0213     
0214     TH1* hE_vs_E_mean         = (TH1*)fitE_vs_E->At(1);     
0215     TH1* hE_vs_E_stddev       = (TH1*)fitE_vs_E->At(2);     
0216     TH1* hTheta_vs_E_mean     = (TH1*)fitTheta_vs_E->At(1); 
0217     TH1* hTheta_vs_E_stddev   = (TH1*)fitTheta_vs_E->At(2); 
0218     TH1* hPhi_vs_theta_mean   = (TH1*)fitPhi_vs_theta->At(1);   
0219     TH1* hPhi_vs_theta_stddev = (TH1*)fitPhi_vs_theta->At(2);   
0220 
0221     
0222     TCanvas *cResolutionGraphs = new TCanvas("resolution_graphs_canvas", "Resolution Graphs", 1200, 800);
0223     cResolutionGraphs->Divide(3, 2);
0224     cResolutionGraphs->cd(1);
0225     hE_vs_E_mean->SetTitle("Mean Energy Offset vs E MC; Energy MC [GeV]; Mean Energy Offset [GeV]");
0226     hE_vs_E_mean->SetMarkerStyle(20);
0227     hE_vs_E_mean->SetMarkerColor(kBlue);
0228     hE_vs_E_mean->SetMaximum(0.01); 
0229     hE_vs_E_mean->SetMinimum(-0.01); 
0230     hE_vs_E_mean->Draw();
0231     cResolutionGraphs->cd(4);
0232     hE_vs_E_stddev->SetTitle("Energy Resolution vs E MC; Energy MC [GeV]; Energy Resolution [GeV]");
0233     hE_vs_E_stddev->SetMarkerStyle(20);
0234     hE_vs_E_stddev->SetMarkerColor(kRed);
0235     hE_vs_E_stddev->SetMaximum(0.04); 
0236     hE_vs_E_stddev->SetMinimum(0); 
0237     hE_vs_E_stddev->Draw();
0238     cResolutionGraphs->cd(2);
0239     hTheta_vs_E_mean->SetTitle("Mean Theta Offset vs E MC; Energy MC [GeV]; Mean Theta Offset [rad]");
0240     hTheta_vs_E_mean->SetMarkerStyle(20);
0241     hTheta_vs_E_mean->SetMarkerColor(kBlue);
0242     hTheta_vs_E_mean->SetMaximum(0.0003); 
0243     hTheta_vs_E_mean->SetMinimum(-0.0003); 
0244     hTheta_vs_E_mean->Draw();
0245     cResolutionGraphs->cd(5);
0246     hTheta_vs_E_stddev->SetTitle("Std Dev Theta Offset vs E MC; Energy MC [GeV]; Std Dev Theta Offset [rad]");
0247     hTheta_vs_E_stddev->SetMarkerStyle(20);
0248     hTheta_vs_E_stddev->SetMarkerColor(kRed);
0249     hTheta_vs_E_stddev->SetMaximum(0.0005); 
0250     hTheta_vs_E_stddev->SetMinimum(0); 
0251     hTheta_vs_E_stddev->Draw();
0252     cResolutionGraphs->cd(3);
0253     hPhi_vs_theta_mean->SetTitle("Mean Phi Offset vs theta MC; theta MC [rad]; Mean Phi Offset [rad]");
0254     hPhi_vs_theta_mean->SetMarkerStyle(20);
0255     hPhi_vs_theta_mean->SetMarkerColor(kBlue);
0256     hPhi_vs_theta_mean->SetMaximum(5); 
0257     hPhi_vs_theta_mean->SetMinimum(-5); 
0258     hPhi_vs_theta_mean->Draw();
0259     cResolutionGraphs->cd(6);
0260     hPhi_vs_theta_stddev->SetTitle("Std Dev Phi Offset vs theta MC; theta MC [rad]; Std Dev Phi Offset [rad]");
0261     hPhi_vs_theta_stddev->SetMarkerStyle(20);
0262     hPhi_vs_theta_stddev->SetMarkerColor(kRed);
0263     hPhi_vs_theta_stddev->SetMaximum(60); 
0264     hPhi_vs_theta_stddev->SetMinimum(0); 
0265     hPhi_vs_theta_stddev->Draw();
0266     cResolutionGraphs->SetGrid();
0267     cResolutionGraphs->Update();
0268     
0269     cResolutionGraphs->SaveAs(resolutionGraphsCanvasName);
0270 
0271     TFile *f = new TFile(outFile,"RECREATE");
0272     cMomentum->Write();
0273     cEnergyThetaPhi->Write();
0274     cResolutionVsMC->Write();
0275     cResolutionGraphs->Write();
0276     px_Hist->Write();
0277     py_Hist->Write();
0278     pz_Hist->Write();
0279     px_diff_Hist->Write();
0280     py_diff_Hist->Write();
0281     pz_res_Hist->Write();
0282     E_Hist->Write();
0283     theta_Hist->Write();
0284     phi_Hist->Write();
0285     E_res_Hist->Write();
0286     theta_diff_Hist->Write();
0287     phi_diff_Hist->Write();
0288     E_res_vs_E_Hist->Write();
0289     E_res_vs_theta_Hist->Write();
0290     E_res_vs_phi_Hist->Write();
0291     theta_diff_vs_E_Hist->Write();
0292     theta_diff_vs_theta_Hist->Write();
0293     theta_diff_vs_phi_Hist->Write();
0294     phi_diff_vs_E_Hist->Write();
0295     phi_diff_vs_theta_Hist->Write();
0296     phi_diff_vs_phi_Hist->Write();
0297 
0298     hE_vs_E_mean->Write();
0299     hE_vs_E_stddev->Write();
0300     hTheta_vs_E_mean->Write();
0301     hTheta_vs_E_stddev->Write();
0302     hPhi_vs_theta_mean->Write();
0303     hPhi_vs_theta_stddev->Write();
0304 
0305     f->Close();
0306 
0307 
0308     
0309     
0310     
0311     
0312     
0313     
0314     
0315     
0316 
0317     
0318     
0319     
0320     
0321     
0322     
0323     
0324 
0325     
0326     
0327     
0328     
0329     
0330     
0331     
0332 
0333     
0334 
0335     
0336     
0337     
0338     
0339     
0340     
0341     
0342     
0343     
0344     
0345     
0346     
0347     
0348 
0349     
0350     
0351     
0352     
0353 
0354     
0355     
0356     
0357     
0358     
0359     
0360     
0361     
0362     
0363     
0364     
0365     
0366     
0367 
0368 }
0369