dump_good_sim.cpp 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141
  1. #include <iostream>
  2. #include <vector>
  3. #include <cmath>
  4. #include <numeric>
  5. #include <TSystem.h>
  6. #include <Math/VectorUtil.h>
  7. #include "filval.hpp"
  8. #include "root_filval.hpp"
  9. #include "common.hpp"
  10. #include "analysis/TrackingNtupleObjs.hpp"
  11. typedef ROOT::Math::LorentzVector<ROOT::Math::PxPyPzE4D<float>> Vec4;
  12. SeedCollection seeds;
  13. SimTrackCollection sim_tracks;
  14. SimVertexCollection sim_vertices;
  15. TrackCollection gsf_tracks;
  16. SuperClusterCollection scls;
  17. void register_objects(TrackingDataSet& tds){
  18. /* seeds.init(tds); */
  19. sim_tracks.init(tds);
  20. sim_vertices.init(tds);
  21. /* gsf_tracks.init(tds); */
  22. /* scls.init(tds); */
  23. }
  24. struct {
  25. Value<float>* x;
  26. Value<float>* y;
  27. Value<float>* z;
  28. Value<float>* sigma_x;
  29. Value<float>* sigma_y;
  30. Value<float>* sigma_z;
  31. } bs;
  32. template <typename T>
  33. struct KinColl {
  34. T* pt;
  35. T* eta;
  36. T* phi;
  37. };
  38. bool in_lum_region(const SimVertex& vertex) {
  39. /* Here approximate the luminous region as a cylinder with length 5*\sigma_z and radius
  40. * sqrt(\simga_x^2 + \sigma_y^2) centered as beamspot position.
  41. */
  42. float dx = vertex.x() - bs.x->get();
  43. float dy = vertex.y() - bs.y->get();
  44. float dz = vertex.z() - bs.z->get();
  45. auto radius = static_cast<float>(5 * sqrt(pow(bs.sigma_x->get(), 2) + pow(bs.sigma_y->get(), 2)));
  46. float half_len = 5*bs.sigma_z->get();
  47. return sqrt(dx*dx + dy*dy) < radius and abs(dz) < half_len;
  48. };
  49. bool is_good_sim(const SimTrack& sim_track) {
  50. const auto& vertex = sim_vertices[sim_track.parentVtxIdx()];
  51. return abs(sim_track.pdgId()) == 11 and in_lum_region(vertex);
  52. };
  53. template <typename T>
  54. std::ostream& operator<< (std::ostream& out, const std::vector<T>& v) {
  55. if ( !v.empty() ) {
  56. out << '[';
  57. std::copy (v.begin(), v.end(), std::ostream_iterator<T>(out, ", "));
  58. out << "\b\b]";
  59. }
  60. return out;
  61. }
  62. void run(){
  63. using namespace std;
  64. using namespace fv;
  65. using namespace fv_root;
  66. using namespace fv_util;
  67. auto file_list = the_config->get_source_files();
  68. string output_filename = the_config->get_output_filename();
  69. TrackingDataSet tds(output_filename, file_list, "trackingNtuple/tree");
  70. register_objects(tds);
  71. // Declare TTree
  72. bs = {
  73. tds.track_branch<float>("bsp_x"),
  74. tds.track_branch<float>("bsp_y"),
  75. tds.track_branch<float>("bsp_z"),
  76. tds.track_branch<float>("bsp_sigmax"),
  77. tds.track_branch<float>("bsp_sigmay"),
  78. tds.track_branch<float>("bsp_sigmaz")
  79. };
  80. while (tds.next()) {
  81. cout << "Event: " << tds.get_current_event() << endl;
  82. for (const auto& sim_track : sim_tracks) {
  83. if (!is_good_sim(sim_track)) continue;
  84. /* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; */
  85. /* SimVertex parentvtx = sim_vertices[sim_track.parentVtxIdx()]; */
  86. cout << "\t" << sim_track.pdgId() << " "
  87. << sim_track.pt() << " "
  88. << sim_track.eta() << " "
  89. << sim_track.phi() << " "
  90. << sim_track.parentVtxIdx() << " "
  91. << sim_track.decayVtxIdx() << " ";
  92. if (sim_track.decayVtxIdx().size() > 0) {
  93. cout << sim_tracks[sim_vertices[sim_track.decayVtxIdx()[0]].daughterSimIdx()[0]].pdgId();
  94. }
  95. cout << endl;
  96. // Log Sim Track
  97. }
  98. }
  99. tds.save_all();
  100. }
  101. int main(int argc, char * argv[]){
  102. using namespace fv_util;
  103. ArgParser args(argc, argv);
  104. if(args.cmd_option_exists("-c")) {
  105. init_config(args.get_cmd_option("-c"));
  106. args.update_config();
  107. if (args.cmd_option_exists("-s")) {
  108. the_config->update_key("silent", true);
  109. }
  110. if (args.cmd_option_exists("-b")) {
  111. the_config->update_key("batch", true);
  112. }
  113. init_log(LogPriority::kLogInfo);
  114. // gSystem->Load("libfilval.so");
  115. run();
  116. } else {
  117. cout << "Usage: ./" << argv[0] << " (-s) (-b) -c config_file.yaml" << endl;
  118. }
  119. return 0;
  120. }