Resolve "Pion rejection YR benchmark"

Closes #34 (closed)

benchmarks/barrel_ecal/scripts/emcal_barrel_pion_rejection_analysis.cxx

+////////////////////////////////////////
+// Read reconstruction ROOT output file
+// Plot variables
+////////////////////////////////////////
+
+#include "ROOT/RDataFrame.hxx"
+#include <iostream>
+#include <algorithm>
+
+#include "dd4pod/Geant4ParticleCollection.h"
+#include "dd4pod/CalorimeterHitCollection.h"
+
+#include "benchmark.h"
+#include "mt.h"
+#include "util.h"
+
+#include "TCanvas.h"
+#include "TStyle.h"
+#include "TMath.h"
+#include "TH1.h"
+#include "TF1.h"
+#include "TH1D.h"
+#include "TFitResult.h"
+#include "TLegend.h"
+#include "TString.h"
+
+R__LOAD_LIBRARY(libfmt.so)
+#include "fmt/core.h"
+#include "DD4hep/Detector.h"
+#include "DDG4/Geant4Data.h"
+#include "DDRec/CellIDPositionConverter.h"
+#include "emcal_barrel_common_functions.h"
+
+using ROOT::RDataFrame;
+using namespace ROOT::VecOps;
+
+void emcal_barrel_pion_rejection_analysis(
+                                          const char* input_fname1 = "sim_output/sim_emcal_barrel_uniform_electrons.root", 
+                                          const char* input_fname2 = "sim_output/sim_emcal_barrel_piminus.root"
+                                          /*
+                                          const char* input_fname1 = "../sim_output/sim_emcal_barrel_uniform_electron_E5.root", 
+                                          const char* input_fname2 = "../sim_output/sim_emcal_barrel_uniform_electron_E10.root",
+                                          const char* input_fname3 = "../sim_output/sim_emcal_barrel_uniform_electron_E18.root",
+                                          const char* input_fname4 = "../sim_output/sim_emcal_barrel_uniform_piminus_E5.root", 
+                                          const char* input_fname5 = "../sim_output/sim_emcal_barrel_uniform_piminus_E10.root",
+                                          const char* input_fname6 = "../sim_output/sim_emcal_barrel_uniform_piminus_E18.root"
+                                          /**/
+                                          )
+{
+  // Setting for graphs
+  gROOT->SetStyle("Plain");
+  gStyle->SetOptFit(1);
+  gStyle->SetLineWidth(2);
+  gStyle->SetPadTickX(1);
+  gStyle->SetPadTickY(1);
+  gStyle->SetPadGridX(1);
+  gStyle->SetPadGridY(1);
+  gStyle->SetPadLeftMargin(0.14);
+  gStyle->SetPadRightMargin(0.14);
+
+  ROOT::EnableImplicitMT();
+  ROOT::RDataFrame d0("events", {input_fname1, input_fname2});
+  //ROOT::RDataFrame d0("events", {input_fname1, input_fname2, input_fname3, input_fname4, input_fname5, input_fname6});
+
+  // Environment Variables
+  std::string detector_path = "";
+  std::string detector_name = "topside";
+  if(std::getenv("DETECTOR_PATH")) {
+    detector_path = std::getenv("DETECTOR_PATH");
+  }
+  if(std::getenv("JUGGLER_DETECTOR")) {
+    detector_name = std::getenv("JUGGLER_DETECTOR");
+  }
+/*
+  // Using the detector layers 
+  dd4hep::Detector& detector = dd4hep::Detector::getInstance();
+  detector.fromCompact(fmt::format("{}/{}.xml", detector_path,detector_name));
+  //dd4hep::rec::CellIDPositionConverter cellid_converter(detector);
+
+  auto decoder = detector.readout("EcalBarrelHits").idSpec().decoder();
+  fmt::print("{}\n", decoder->fieldDescription());
+  auto layer_index = decoder->index("layer");
+  fmt::print(" layer index is {}.\n", layer_index);
+*/
+  // Rejection Value [GeV] based upon Energy deposit in the first 4 layers.
+  // Currently constructed from looking at tthe pion and electron Energy deposit plots
+  // should eventually grab a value from a json file
+  // Suggested cut at 0.005
+  double rejectCut = 0.0;
+
+  // Thrown Energy [GeV]
+  auto Ethr = [](std::vector<dd4pod::Geant4ParticleData> const& input) {
+    return TMath::Sqrt(input[2].psx*input[2].psx + input[2].psy*input[2].psy + input[2].psz*input[2].psz + input[2].mass*input[2].mass);
+  };
+
+  // Thrown Momentum [GeV]
+  auto Pthr = [](std::vector<dd4pod::Geant4ParticleData> const& input) {
+    return TMath::Sqrt(input[2].psx*input[2].psx + input[2].psy*input[2].psy + input[2].psz*input[2].psz);
+  };
+
+  // Thrown Eta 
+  auto Eta = [](std::vector<dd4pod::Geant4ParticleData> const& input) {
+    double E  = TMath::Sqrt(input[2].psx*input[2].psx + input[2].psy*input[2].psy + input[2].psz*input[2].psz + input[2].mass*input[2].mass);
+    return 0.5*TMath::Log((E + input[2].psz) / (E - input[2].psz));
+  };
+
+  // Number of hits
+  auto nhits = [] (const std::vector<dd4pod::CalorimeterHitData>& evt) {return (int) evt.size(); };
+
+  // Energy deposition [GeV]
+  auto Esim = [](const std::vector<dd4pod::CalorimeterHitData>& evt) {
+    auto total_edep = 0.0;
+    for (const auto& i: evt)
+      total_edep += i.energyDeposit;
+  return total_edep;
+  };
+/*
+  // Energy deposititon [GeV] in the first 4 layers
+  auto Esim_front = [=](const std::vector<dd4pod::CalorimeterHitData>& evt) {
+    auto total_edep = 0.0;
+    for (const auto& i: evt) {
+      //fmt::print("cell id {}, layer {}\n",i.cellID, decoder->get(i.cellID, layer_index));
+      if( decoder->get(i.cellID, layer_index) < 4 ){
+        total_edep += i.energyDeposit;
+      }
+    }
+    return total_edep;
+  };
+*/
+  // Sampling fraction = Esampling / Ethrown
+  auto fsam = [](const double& sampled, const double& thrown) {
+    return sampled / thrown;
+  };
+
+  // E_front / p
+  auto fEp = [](const double& E_front, const double& mom) {
+    return E_front / mom;
+  };
+
+  // Returns the pdgID of the particle
+  auto getpid = [](std::vector<dd4pod::Geant4ParticleData> const& input) {
+    return input[2].pdgID;
+  };
+
+  // Returns number of particle daughters
+  auto getdau = [](std::vector<dd4pod::Geant4ParticleData> const& input){
+    return input[2].daughters_begin;
+  };
+
+  // Filter function to get electrons
+  auto is_electron = [](std::vector<dd4pod::Geant4ParticleData> const& input){
+    if (input[2].pdgID == 11){return true;}
+    else {return false;}
+  };
+
+  // Filter function to get just negative pions
+  auto is_piMinus = [](std::vector<dd4pod::Geant4ParticleData> const& input){
+    if (input[2].pdgID == -211){return true;}
+    else {return false;}
+  };
+
+
+  // Define variables
+  auto d1 = d0.Define("Ethr",       Ethr,       {"mcparticles"})
+              .Define("Pthr",       Pthr,       {"mcparticles"})
+              .Define("nhits",      nhits,      {"EcalBarrelHits"})
+              .Define("Esim",       Esim,       {"EcalBarrelHits"})
+              .Define("fsam",       fsam,       {"Esim","Ethr"})
+              .Define("pid",        getpid,     {"mcparticles"})
+              //.Define("Esim_front", Esim_front, {"EcalBarrelHits"})
+              //.Define("EOverP",     fEp,        {"Esim_front", "Pthr"})
+              .Define("Eta",        Eta,        {"mcparticles"})
+              ;
+  
+  // Energy Deposit Filters
+  //auto d_ele_rej = d_ele.Filter("Esim_front > " + to_string(rejectCut));
+  //auto d_pim_rej = d_pim.Filter("Esim_front > " + to_string(rejectCut));
+  
+  // Particle Filters
+  auto d_ele = d1.Filter(is_electron, {"mcparticles"});
+  auto d_pim = d1.Filter(is_piMinus, {"mcparticles"});
+
+
+  // Gathering benchmarks and plotting distrubutions
+  std::vector<double> E                 = {5, 10, 18};
+  std::vector<double> ledges5           = {2.8, 0.4, 0.3, 0.5};
+  std::vector<double> ledges10          = {1.4, 0.5, 0.6, 1.0};
+  std::vector<double> ledges18          = {0.9, 0.9, 1.0, 1.8};
+  std::vector<double> maxRate5          = {0.1, 100, 500, 1000};
+  std::vector<double> maxRate10         = {10, 400, 800, 1000};
+  std::vector<double> maxRate18         = {200, 800, 1000, 100};
+  std::vector<vector<double>> maxRate   = {maxRate5, maxRate10, maxRate18};
+  std::vector<vector<double>> lowEdges  = {ledges5, ledges10, ledges18};
+  double suppression                    = 1e-4;
+  std::vector<vector<double>> rejRatios = lowEdges;
+  std::vector<std::string> etaBin       = {"Eta >= -3.5 && Eta < -2.0", "Eta >= -2.0 && Eta < -1.0", "Eta >= -1.0 && Eta < 0", "Eta >= 0 && Eta < 1.0"};
+  std::vector<std::string> etaTitle     = {"-3.5 < #eta < -2.0", "-2.0 < #eta < -1.0", "-1.0 < #eta < 0", "0 < #eta < 1.0"};
+
+  // Barrel eta cuts
+  // The eta range for the barrel is -1 < eta < 1
+  // Threfore the first bins are empty and will not be iterated over
+
+  for (int i = 0; i < 3; i++){   // E loop
+    for (int j = 2; j < 4; j++){ // Eta Looop
+      // Apply eta cuts/binning
+      auto e_eta = d_ele.Filter(etaBin[j]);
+      auto p_eta = d_pim.Filter(etaBin[j]);
+      
+      // Print out the momentum distributions for the electron and pi-
+      std::string title = "e^{-} (E = " + std::to_string((int)E[i]) + " GeV) : " + etaTitle[j] + "; p [GeV]; Events";
+      auto he = e_eta.Histo1D({"he", title.c_str(), 100, lowEdges[i][j], E[i]}, "Pthr");
+      he->GetYaxis()->SetTitleOffset(1.4);
+      he->SetLineWidth(2);
+      he->SetLineColor(kBlue);
+      
+      title = "#pi^{-} (E = " + std::to_string((int)E[i]) + " GeV) : " + etaTitle[j] + "; p [GeV]; Events";
+      auto hp = p_eta.Histo1D({"hp", title.c_str(), 100, lowEdges[i][j], E[i]}, "Pthr");
+      hp->GetYaxis()->SetTitleOffset(1.4);
+      hp->SetLineWidth(2);
+      hp->SetLineColor(kBlue);
+
+      auto c = new TCanvas("c", "c", 700, 500);
+      c->SetLogy(1);
+      he->DrawClone();
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_mom_ele_E{}_eta{}.png", (int)E[i], j)).c_str());
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_mom_ele_E{}_eta{}.pdf", (int)E[i], j)).c_str());
+      c->Clear();
+
+      hp->DrawClone();
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_mom_pim_E{}_eta{}.png", (int)E[i], j)).c_str());
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_mom_pim_E{}_eta{}.pdf", (int)E[i], j)).c_str());
+      c->Clear();
+
+      // Gather ratio of pi/e for each Energy and eta bin
+      // Then plot the distributions
+      rejRatios[i][j] = (double)hp->Integral() / (double)he->Integral();
+      std::cout << hp->Integral() << " " << he->Integral()  << " " << rejRatios[i][j] << std::endl;
+
+      hp->Divide(he.GetPtr());
+      title = "#pi^{-}/e^{-} (E = " + std::to_string((int)E[i]) + " GeV) : " + etaTitle[j];
+      hp->SetTitle(title.c_str());
+      hp->DrawClone();
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_ratio_pim_E{}_eta{}.png", (int)E[i], j)).c_str());
+      c->SaveAs((fmt::format("results/emcal_barrel_pion_rej_ratio_pim_E{}_eta{}.pdf", (int)E[i], j)).c_str());
+      c->Clear();
+    }
+  }
+
+  // Writing out benchmarks
+  //Tests
+  std::string test_tag = "Barrel_emcal_pion_rejection";
+  //TODO: Change test_tag to something else
+  std:string detector = "Barrel_emcal";
+  
+  for (int i = 0; i < etaTitle.size(); i++){
+    etaTitle[i].erase(std::remove(etaTitle[i].begin(), etaTitle[i].end(), '#'), etaTitle[i].end());
+    std::replace(etaTitle[i].begin(), etaTitle[i].end(), 'e', 'E');    
+  }
+  
+  // E, Eta = 18, 2
+  eic::util::Test pion_rejection_E18_Eta2{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[0], 2)},
+     {"title", "Pion Rejection1"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[0], etaTitle[2])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[0][2])},
+     {"value", std::to_string(rejRatios[0][2])}
+    }
+  };
+  pion_rejection_E18_Eta2.pass(0.1);
+
+  // E, Eta = 18, 3
+  eic::util::Test pion_rejection_E18_Eta3{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[0], 3)},
+     {"title", "Pion Rejection2"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[0], etaTitle[3])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[0][3])},
+     {"value", std::to_string(rejRatios[0][3])}
+   }
+  };
+  pion_rejection_E18_Eta3.pass(0.1);
+
+  // E, Eta = 10, 2
+  eic::util::Test pion_rejection_E10_Eta2{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[1], 2)},
+     {"title", "Pion Rejection"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[1], etaTitle[2])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[1][2])},
+     {"value", std::to_string(rejRatios[1][2])}
+    }
+  };
+  pion_rejection_E10_Eta2.pass(0.1);
+
+  // E, Eta = 10, 3
+  eic::util::Test pion_rejection_E10_Eta3{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[1], 3)},
+     {"title", "Pion Rejection"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[1], etaTitle[3])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[1][3])},
+     {"value", std::to_string(rejRatios[1][3])}
+    }
+  };
+  pion_rejection_E10_Eta3.pass(0.1);
+
+  // E, Eta = 5, 2
+  eic::util::Test pion_rejection_E5_Eta2{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[2], 2)},
+     {"title", "Pion Rejection"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[2], etaTitle[2])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[2][2])},
+     {"value", std::to_string(rejRatios[2][2])}
+    }
+  };
+  pion_rejection_E5_Eta2.pass(0.1);
+
+  // E, Eta = 5, 3
+  eic::util::Test pion_rejection_E5_Eta3{
+    {{"name", fmt::format("{}_E{}_EtaBin{}", test_tag, (int)E[2], 3)},
+     {"title", "Pion Rejection"},
+     {"description", fmt::format("Pion rejection with E = {}, and {}", (int)E[2], etaTitle[3])},
+     {"quantity", "pi-/e-"},
+     {"target", std::to_string(suppression * maxRate[2][3])},
+     {"value", std::to_string(rejRatios[2][3])}
+    }
+  };
+  pion_rejection_E5_Eta3.pass(0.1);
+  
+  // Writing out all tests
+  eic::util::write_test({pion_rejection_E18_Eta2, 
+                         pion_rejection_E18_Eta3, 
+                         pion_rejection_E10_Eta2, 
+                         pion_rejection_E10_Eta3, 
+                         pion_rejection_E5_Eta2, 
+                         pion_rejection_E5_Eta3}, 
+                         fmt::format("results/{}_pion_rej.json", detector));
+
+}