File indexing completed on 2024-09-27 07:02:38
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
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
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
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
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
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
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
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 }