tracking_detectors/analysis/sim_track_hits.cxx

0001 #include "ROOT/RDataFrame.hxx"
0002 #include "TCanvas.h"
0003 #include "TH1D.h"
0004 #include "TLegend.h"
0005 #include "TProfile.h"
0006 #include "THStack.h"
0007 #include "Math/Vector4D.h"
0008
0009 #include <cstdlib>
0010 #include <iostream>
0011
0012 R__LOAD_LIBRARY(libedm4eic.so)
0013
0014 #include <fmt/format.h>
0015
0016 #include "edm4hep/MCParticleCollection.h"
0017 #include "edm4hep/SimTrackerHitCollection.h"
0018 #include "edm4eic/ClusterCollection.h"
0019 #include "edm4eic/TrackParametersCollection.h"
0020 #include "edm4eic/TrackerHitCollection.h"
0021
0022 using ROOT::RDataFrame;
0023 using namespace ROOT::VecOps;
0024
0025 auto p_track = [](std::vector<edm4eic::TrackParametersData> const& in) {
0026   std::vector<double> result;
0027   for (size_t i = 0; i < in.size(); ++i) {
0028     result.push_back(std::abs(1.0 / (in[i].qOverP)));
0029   }
0030   return result;
0031 };
0032
0033 std::vector<float> pt(std::vector<edm4hep::MCParticleData> const& in) {
0034   std::vector<float> result;
0035   for (size_t i = 0; i < in.size(); ++i) {
0036     result.push_back(std::sqrt(in[i].momentum.x * in[i].momentum.x + in[i].momentum.y * in[i].momentum.y));
0037   }
0038   return result;
0039 }
0040
0041 auto momentum = [](std::vector<ROOT::Math::PxPyPzMVector> const& in) {
0042   std::vector<double> result;
0043   for (size_t i = 0; i < in.size(); ++i) {
0044     result.push_back(in[i].E());
0045   }
0046   return result;
0047 };
0048 auto theta = [](std::vector<ROOT::Math::PxPyPzMVector> const& in) {
0049   std::vector<double> result;
0050   for (size_t i = 0; i < in.size(); ++i) {
0051     result.push_back(in[i].Theta() * 180 / M_PI);
0052   }
0053   return result;
0054 };
0055 auto fourvec = [](ROOT::VecOps::RVec<edm4hep::MCParticleData> const& in) {
0056   std::vector<ROOT::Math::PxPyPzMVector> result;
0057   ROOT::Math::PxPyPzMVector lv;
0058   for (size_t i = 0; i < in.size(); ++i) {
0059     lv.SetCoordinates(in[i].momentum.x, in[i].momentum.y, in[i].momentum.z, in[i].mass);
0060     result.push_back(lv);
0061   }
0062   return result;
0063 };
0064
0065 auto delta_p = [](const std::vector<double>& tracks, const std::vector<double>& thrown) {
0066   std::vector<double> res;
0067   for (const auto& p1 : thrown) {
0068     for (const auto& p2 : tracks) {
0069       res.push_back(p1 - p2);
0070     }
0071   }
0072   return res;
0073 };
0074
0075 auto delta_p_over_p = [](const std::vector<double>& tracks, const std::vector<double>& thrown) {
0076   std::vector<double> res;
0077   for (const auto& p1 : thrown) {
0078     for (const auto& p2 : tracks) {
0079       res.push_back((p1 - p2) / p1);
0080     }
0081   }
0082   return res;
0083 };
0084
0085 // Create dataframe with extra definitions for hit multiplicities
0086 ROOT::RDF::RNode add_subsystems(ROOT::RDF::RNode df, std::vector<std::pair<std::string, std::string>> hitcols) {
0087   if (hitcols.empty()) {
0088     return df;
0089   }
0090   const auto [name, collection] = hitcols.back();
0091   std::cout << " --> registering subsystem " << collection << "\n";
0092   auto df2 = df.Define("N_" + name, [](std::vector<edm4hep::SimTrackerHitData> hits) { return hits.size(); }, {collection});
0093   hitcols.pop_back();
0094   return add_subsystems(df2, hitcols);
0095 };
0096
0097 int sim_track_hits(const char* fname = "sim_track_hits.edm4hep.root") {
0098
0099   ROOT::EnableImplicitMT();
0100   ROOT::RDataFrame df("events", fname);
0101
0102   // detect detector setup
0103   std::string detector = "default";
0104   if (const char* env_detector = std::getenv("DETECTOR_VERSION")) {
0105     if (detector.find("acadia") != std::string::npos) {
0106       detector = "acadia";
0107     } else if (detector.find("canyonlands") != std::string::npos) {
0108       detector = "canyonlands";
0109     }
0110   }
0111   std::cout << "sim_track_hits: detector set to " << detector << std::endl;
0112
0113   // minimal hit collection setup
0114   std::vector<std::pair<std::string, std::string>> hitCollections{{"vtx_barrel", "VertexBarrelHits"},
0115                                                                   {"si_barrel", "SiBarrelHits"},
0116                                                                   {"trk_endcap", "TrackerEndcapHits"},
0117                                                                   {"mpgd_barrel", "MPGDBarrelHits"},
0118   };
0119
0120   auto df0 = df.Define("isThrown", "MCParticles.generatorStatus == 1")
0121                  .Define("thrownParticles", "MCParticles[isThrown]")
0122                  .Define("thrownP", fourvec, {"thrownParticles"})
0123                  .Define("p_thrown", momentum, {"thrownP"})
0124                  .Define("theta_thrown", theta, {"thrownP"})
0125                  .Define("theta0", "theta_thrown[0]");
0126   auto df1 = add_subsystems(df0, hitCollections);
0127
0128   // get 1D and 2D histogram definitions
0129   std::vector<decltype(df1.Histo2D({}, "", ""))> h2D;
0130   std::vector<decltype(df1.Histo1D(""))> hnth;
0131   std::vector<decltype(df1.Histo1D(""))> hn;
0132
0133   for (const auto& [name, col] : hitCollections) {
0134     hnth.push_back(
0135         df1.Histo1D({(name + "_nhits_vs_theta").c_str(), "; #theta [deg.]; ", 20, 0, 180}, "theta0", "N_" + name));
0136     hn.push_back(df1.Histo1D({(name + "_nhits").c_str(), "; Nhits; #", 20, 0, 20}, "N_" + name));
0137     h2D.push_back(df1.Histo2D({(name + "_x_vs_y").c_str(), "; x ; y ", 100, -600, 600, 100, -600, 600},
0138                               col + ".position.x", col + ".position.y"));
0139     h2D.push_back(df1.Histo2D({(name + "_x_vs_z").c_str(), "; z ; x ", 100, -1200, 1200, 100, -900, 900},
0140                               col + ".position.z", col + ".position.x"));
0141   }
0142   auto hth = df1.Histo1D({"theta0", "; #theta [deg.]", 20, 0, 180}, "theta0");
0143
0144   // Draw multiplicity versus theta
0145   {
0146     TCanvas c;
0147     THStack hs{"n_hits", "; #theta [deg.] "};
0148     int idx = 1;
0149     for (auto& h : hnth) {
0150       h->Divide(&*hth);
0151       h->SetLineColor(idx++);
0152       hs.Add((TH1D*)h->Clone());
0153     }
0154     hs.Draw("nostack, hist");
0155     c.BuildLegend();
0156     c.SaveAs("results/tracking_detectors/sim_track_hits_n_hits_vs_theta.png");
0157     c.SaveAs("results/tracking_detectors/sim_track_hits_n_hits_vs_theta.pdf");
0158   }
0159   // Draw nhits
0160   {
0161     TCanvas c;
0162     THStack hs{"n_hits", "; NHits "};
0163     int idx = 1;
0164     for (auto& h : hn) {
0165       h->SetLineColor(idx++);
0166       hs.Add((TH1D*)h->Clone());
0167     }
0168     hs.Draw("nostack, hist");
0169     c.BuildLegend();
0170     c.SaveAs("results/tracking_detectors/sim_track_hits_n_hits.png");
0171     c.SaveAs("results/tracking_detectors/sim_track_hits_n_hits.pdf");
0172   }
0173
0174   // Draw total multiplicity versus theta
0175   {
0176     TCanvas c;
0177     hth->DrawClone();
0178     c.BuildLegend();
0179     c.SaveAs("results/tracking_detectors/sim_track_hits_theta.png");
0180     c.SaveAs("results/tracking_detectors/sim_track_hits_theta.pdf");
0181   }
0182
0183   for (auto& h : h2D) {
0184     TCanvas c;
0185     h->DrawClone("colz");
0186     c.SaveAs(fmt::format("results/tracking_detectors/sim_track_hits_{}.png", h->GetName()).c_str());
0187     c.SaveAs(fmt::format("results/tracking_detectors/sim_track_hits_{}.pdf", h->GetName()).c_str());
0188   }
0189
0190   return 0;
0191 }