Back to home page

EIC code displayed by LXR

 
 

    


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

0001 ////////////////////////////////////////
0002 // Read reconstruction ROOT output file
0003 // Plot variables
0004 ////////////////////////////////////////
0005 
0006 #include "ROOT/RDataFrame.hxx"
0007 #include <iostream>
0008 #include <fstream>
0009 #include <fmt/core.h>
0010 
0011 #include "edm4hep/MCParticleCollection.h"
0012 #include "edm4hep/SimCalorimeterHitCollection.h"
0013 
0014 #include "TCanvas.h"
0015 #include "TStyle.h"
0016 #include "TMath.h"
0017 #include "TH1.h"
0018 #include "TF1.h"
0019 #include "TH1D.h"
0020 #include <nlohmann/json.hpp>
0021 
0022 #include "emcal_barrel_common_functions.h"
0023 
0024 using ROOT::RDataFrame;
0025 using namespace ROOT::VecOps;
0026 using json = nlohmann::json;
0027 
0028 void save_canvas(TCanvas* c, std::string label)
0029 {
0030   c->SaveAs(fmt::format("results/{}.png",label).c_str());
0031   c->SaveAs(fmt::format("results/{}.pdf",label).c_str());
0032 }
0033 
0034 void save_canvas(TCanvas* c, std::string label, std::string particle_label)
0035 {
0036   std::string label_with_E = fmt::format("emcal_barrel_{}_{}", particle_label, label); 
0037   save_canvas(c, label_with_E);
0038 }
0039 
0040 void emcal_barrel_particles_analysis(std::string particle_name = "electron", bool save_calib = false)
0041 {
0042   // Setting for graphs
0043   gROOT->SetStyle("Plain");
0044   gStyle->SetOptFit(1);
0045   gStyle->SetLineWidth(2);
0046   gStyle->SetPadTickX(1);
0047   gStyle->SetPadTickY(1);
0048   gStyle->SetPadGridX(1);
0049   gStyle->SetPadGridY(1);
0050   gStyle->SetPadLeftMargin(0.14);
0051   gStyle->SetPadRightMargin(0.14);
0052 
0053   json j;
0054   // variables that will be saved in the JSON file
0055   double Ethr_mean;
0056   double fSam_mean;
0057   double fSam_img_mean;
0058   double fSam_scfi_mean;
0059   double fSam_mean_err;
0060   double fSam_img_mean_err;
0061   double fSam_scfi_mean_err;
0062 
0063   ROOT::EnableImplicitMT();
0064   std::string input_fname = fmt::format("sim_output/sim_emcal_barrel_{}.edm4hep.root", particle_name);
0065   ROOT::RDataFrame d0("events", input_fname);
0066 
0067   // Environment Variables
0068   std::string detector_path = "";
0069   std::string detector_name = "athena";
0070   if(std::getenv("DETECTOR_PATH")) {
0071     detector_path = std::getenv("DETECTOR_PATH");
0072   }
0073   if(std::getenv("DETECTOR_CONFIG")) {
0074     detector_name = std::getenv("DETECTOR_CONFIG");
0075   }
0076 
0077   // Thrown Energy [GeV]
0078   auto Ethr = [](std::vector<edm4hep::MCParticleData> const& input) {
0079     auto p = input[2];
0080     auto energy = TMath::Sqrt(p.momentum.x * p.momentum.x + p.momentum.y * p.momentum.y + p.momentum.z * p.momentum.z + p.mass * p.mass);
0081     return energy;
0082   };
0083 
0084   // Number of hits
0085   auto nhits = [] (const std::vector<edm4hep::SimCalorimeterHitData>& evt) {
0086     return (int) evt.size();
0087   };
0088 
0089   // Energy deposition [GeV]
0090   auto Esim = [](const std::vector<edm4hep::SimCalorimeterHitData>& evt) {
0091     auto total_edep = 0.0;
0092     for (const auto& i: evt)
0093       total_edep += i.energy;
0094     return total_edep;
0095   };
0096 
0097   // Sampling fraction = Esampling / Ethrown
0098   auto fsam = [](const double sampled, const double thrown) {
0099     return sampled / thrown;
0100   };
0101 
0102   // Define variables
0103   auto d1 = ROOT::RDF::RNode(
0104     d0.Define("Ethr", Ethr, {"MCParticles"})
0105   );
0106 
0107   // Define optional variables
0108   auto fsam_estimate = 1.0;
0109   if (d1.HasColumn("EcalBarrelScFiHits")) {
0110     d1 = d1.Define("nhits", nhits, {"EcalBarrelImagingHits"})
0111            .Define("EsimImg", Esim, {"EcalBarrelImagingHits"})
0112            .Define("EsimScFi", Esim, {"EcalBarrelScFiHits"})
0113            .Define("Esim", "EsimImg+EsimScFi")
0114            .Define("fsamImg", fsam, {"EsimImg", "Ethr"})
0115            .Define("fsamScFi", fsam, {"EsimScFi", "Ethr"})
0116            .Define("fsam", fsam, {"Esim", "Ethr"});
0117     fsam_estimate = 0.1;
0118   } else {
0119     d1 = d1.Define("nhits", nhits, {"EcalBarrelSciGlassHits"})
0120            .Define("Esim", Esim, {"EcalBarrelSciGlassHits"})
0121            .Define("fsam", fsam, {"Esim", "Ethr"});
0122     fsam_estimate = 1.0;
0123   }
0124 
0125 
0126   // Define Histograms
0127   auto Ethr_max = 25.0;
0128   auto hEthr = d1.Histo1D(
0129       {"hEthr", "Thrown Energy; Thrown Energy [GeV]; Events", 100, 0.0, Ethr_max},
0130       "Ethr");
0131   auto hNhits =
0132       d1.Histo1D({"hNhits", "Number of hits per events; Number of hits; Events", 100, 0.0, 2000.0},
0133       "nhits");
0134   auto hEsim = d1.Histo1D(
0135       {"hEsim", "Energy Deposit; Energy Deposit [GeV]; Events", 500, 0.0, fsam_estimate * Ethr_max},
0136       "Esim");
0137   auto hfsam = d1.Histo1D(
0138       {"hfsam", "Sampling Fraction; Sampling Fraction; Events", 400, 0.0, 2.0 * fsam_estimate},
0139       "fsam");
0140 
0141   // Define optional histograms
0142   ROOT::RDF::RResultPtr<::TH1D> hEsimImg, hEsimScFi, hfsamImg, hfsamScFi;
0143   if (d0.HasColumn("EcalBarrelScFiHits")) {
0144     hEsimImg = d1.Histo1D(
0145       {"hEsimImg", "Energy Deposit; Energy Deposit [GeV]; Events", 500, 0.0, fsam_estimate * Ethr_max},
0146       "EsimImg");
0147     hEsimScFi = d1.Histo1D(
0148       {"hEsimScFi", "Energy Deposit; Energy Deposit [GeV]; Events", 500, 0.0, fsam_estimate * Ethr_max},
0149       "EsimScFi");
0150     hfsamImg = d1.Histo1D(
0151       {"hfsamImg", "Sampling Fraction; Sampling Fraction; Events", 400, 0.0, 2.0 * fsam_estimate},
0152       "fsamImg");
0153     hfsamScFi = d1.Histo1D(
0154       {"hfsamScFi", "Sampling Fraction; Sampling Fraction; Events", 400, 0.0, 2.0 * fsam_estimate},
0155       "fsamScFi");
0156   }
0157 
0158   addDetectorName(detector_name, hEthr.GetPtr());
0159   addDetectorName(detector_name, hEsim.GetPtr());
0160   addDetectorName(detector_name, hfsam.GetPtr());
0161 
0162   // Event Counts
0163   auto nevents_thrown = d1.Count();
0164   std::cout << "Number of Thrown Events: " << (*nevents_thrown) << "\n";
0165 
0166   // Draw Histograms
0167   if (hEthr.IsReady()) {
0168     TCanvas* c1 = new TCanvas("c1", "c1", 700, 500);
0169     c1->SetLogy(1);
0170     auto h = hEthr->DrawCopy();
0171     Ethr_mean = h->GetMean();
0172     h->SetLineWidth(2);
0173     h->SetLineColor(kBlue);
0174     save_canvas(c1,"Ethr",particle_name);
0175   }
0176 
0177   if (hNhits.IsReady()) {
0178     TCanvas* c2 = new TCanvas("c2", "c2", 700, 500);
0179     c2->SetLogy(1);
0180     auto h = hNhits->DrawCopy();
0181     h->SetLineWidth(2);
0182     h->SetLineColor(kBlue);
0183     save_canvas(c2,"nhits",particle_name);
0184   }
0185 
0186   if (hEsim.IsReady()) {
0187     TCanvas* c3 = new TCanvas("c3", "c3", 700, 500);
0188     c3->SetLogy(1);
0189     auto h = hEsim->DrawCopy();
0190     h->SetLineWidth(2);
0191     h->SetLineColor(kBlue);
0192     double up_fit = h->GetMean() + 5*h->GetStdDev();
0193     double down_fit = h->GetMean() - 5*h->GetStdDev();
0194     h->GetXaxis()->SetRangeUser(0.,up_fit);
0195     save_canvas(c3,"Esim",particle_name);
0196   }
0197 
0198   if (hfsam.IsReady()) {
0199     TCanvas* c4 = new TCanvas("c4", "c4", 700, 500);
0200     auto h = hfsam->DrawCopy();
0201     h->SetLineWidth(2);
0202     h->SetLineColor(kBlue);
0203     double up_fit = h->GetMean() + 5*h->GetStdDev();
0204     double down_fit = h->GetMean() - 5*h->GetStdDev();
0205     h->Fit("gaus", "", "", down_fit, up_fit);
0206     h->GetXaxis()->SetRangeUser(down_fit,up_fit);
0207     TF1 *gaus = h->GetFunction("gaus");
0208     fSam_mean = gaus->GetParameter(1);
0209     fSam_mean_err = gaus->GetParError(1);
0210     gaus->SetLineWidth(2);
0211     gaus->SetLineColor(kRed); 
0212     save_canvas(c4,"fsam",particle_name);
0213   }
0214 
0215   if (hfsamImg.IsReady()) {
0216     TCanvas* c5 = new TCanvas("c5", "c5", 700, 500);
0217     auto h = hfsamImg->DrawCopy();
0218     h->SetLineWidth(2);
0219     h->SetLineColor(kBlue);
0220     double up_fit = h->GetMean() + 5*h->GetStdDev();
0221     double down_fit = h->GetMean() - 5*h->GetStdDev();
0222     h->Fit("gaus", "", "", down_fit, up_fit);
0223     h->GetXaxis()->SetRangeUser(down_fit,up_fit);
0224     TF1 *gaus = h->GetFunction("gaus");
0225     fSam_img_mean = gaus->GetParameter(1);
0226     fSam_img_mean_err = gaus->GetParError(1);
0227     gaus->SetLineWidth(2);
0228     gaus->SetLineColor(kRed); 
0229     save_canvas(c5,"fsamImg",particle_name);
0230   }
0231 
0232   if (hfsamScFi.IsReady()) {
0233     TCanvas* c6 = new TCanvas("c6", "c6", 700, 500);
0234     auto h = hfsamScFi->DrawCopy();
0235     h->SetLineWidth(2);
0236     h->SetLineColor(kBlue);
0237     double up_fit = h->GetMean() + 5*h->GetStdDev();
0238     double down_fit = h->GetMean() - 5*h->GetStdDev();
0239     h->Fit("gaus", "", "", down_fit, up_fit);
0240     h->GetXaxis()->SetRangeUser(down_fit,up_fit);
0241     TF1 *gaus = h->GetFunction("gaus");
0242     fSam_scfi_mean = gaus->GetParameter(1);
0243     fSam_scfi_mean_err = gaus->GetParError(1);
0244     gaus->SetLineWidth(2);
0245     gaus->SetLineColor(kRed); 
0246     save_canvas(c6,"fsamScFi",particle_name);
0247   }
0248 
0249   j[particle_name] = {
0250     {"particle_name", particle_name},
0251     {"thrown_energy", Ethr_mean},
0252     {"sampling_fraction", fSam_mean},
0253     {"sampling_fraction_error", fSam_mean_err},
0254     {"sampling_fraction_img", fSam_img_mean},
0255     {"sampling_fraction_error_img", fSam_img_mean_err},
0256     {"sampling_fraction_scfi", fSam_scfi_mean},
0257     {"sampling_fraction_error_scfi", fSam_scfi_mean_err}      
0258   };
0259   if (save_calib) {
0260     std::string calib_output_path = fmt::format("results/emcal_barrel_{}_calibration.json", particle_name);
0261     std::cout << "Saving calibration results to " << calib_output_path << std::endl;
0262     std::ofstream o(calib_output_path);
0263     o << std::setw(4) << j << std::endl;
0264   }
0265 }
0266