obj_types.hpp 8.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213
  1. /**
  2. * @file
  3. * @author Caleb Fangmeier <caleb@fangmeier.tech>
  4. * @version 0.1
  5. *
  6. * @section LICENSE
  7. *
  8. *
  9. * MIT License
  10. *
  11. * Copyright (c) 2017 Caleb Fangmeier
  12. *
  13. * Permission is hereby granted, free of charge, to any person obtaining a copy
  14. * of this software and associated documentation files (the "Software"), to deal
  15. * in the Software without restriction, including without limitation the rights
  16. * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  17. * copies of the Software, and to permit persons to whom the Software is
  18. * furnished to do so, subject to the following conditions:
  19. *
  20. * The above copyright notice and this permission notice shall be included in all
  21. * copies or substantial portions of the Software.
  22. *
  23. * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  24. * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  25. * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  26. * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  27. * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  28. * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  29. * SOFTWARE.
  30. *
  31. * @section DESCRIPTION
  32. * Define some common data structures used throughout the analysis and methods
  33. * to make them available.
  34. */
  35. #ifndef OBJ_TYPES_HPP
  36. #define OBJ_TYPES_HPP
  37. #include "TLorentzVector.h"
  38. #include "filval/filval.hpp"
  39. #include "analysis/MiniTreeDataSet.hpp"
  40. struct Jet{
  41. float b_cmva;
  42. int mcMatchFlav;
  43. };
  44. struct Lepton {
  45. int pdg_id;
  46. int mcMatchPdgId;
  47. int charge;
  48. };
  49. struct GenPart{
  50. int pdgId;
  51. int motherIndex;
  52. int motherId;
  53. int status;
  54. };
  55. struct Particle{
  56. enum{JET,
  57. LEPTON,
  58. GENPART,
  59. } tag;
  60. int idx;
  61. TLorentzVector v;
  62. union
  63. {
  64. Jet jet;
  65. Lepton lepton;
  66. GenPart genpart;
  67. };
  68. static Particle
  69. Jet(int idx, TLorentzVector v, Jet&& jet){
  70. Particle p = {Particle::JET, idx, v, {}};
  71. p.jet = jet;
  72. return p;
  73. }
  74. static Particle
  75. Lepton(int idx, TLorentzVector v, Lepton&& lepton){
  76. Particle p = {Particle::LEPTON, idx, v, {}};
  77. p.lepton = lepton;
  78. return p;
  79. }
  80. static Particle
  81. GenPart(int idx, TLorentzVector v, GenPart&& genpart){
  82. Particle p = {Particle::GENPART, idx, v, {}};
  83. p.genpart = genpart;
  84. return p;
  85. }
  86. };
  87. decltype(auto)
  88. construct_jets_value(MiniTreeDataSet& mt){
  89. mt.track_branch<int>("nJet");
  90. mt.track_branch_vec<float>("nJet", "Jet_pt");
  91. mt.track_branch_vec<float>("nJet", "Jet_eta");
  92. mt.track_branch_vec<float>("nJet", "Jet_phi");
  93. mt.track_branch_vec<float>("nJet", "Jet_mass");
  94. auto Jet_4v = lorentz_vectors("Jet_pt", "Jet_eta", "Jet_phi", "Jet_mass", "Jet_4v");
  95. auto jets = apply(GenFunction::register_function<std::vector<Particle>(std::vector<TLorentzVector>,
  96. std::vector<float>,
  97. std::vector< int >
  98. )>("build_reco_jets",
  99. FUNC(([](const std::vector<TLorentzVector>& vs,
  100. const std::vector<float>& b_cmvas,
  101. const std::vector< int >& mcMatchFlavs
  102. ){
  103. std::vector<Particle> jets;
  104. for(int i=0; i<vs.size(); i++){
  105. Particle p = Particle::Jet(i, vs[i], {b_cmvas[i],
  106. mcMatchFlavs[i]
  107. });
  108. jets.push_back(p);
  109. }
  110. return jets;
  111. }))), fv::tuple(Jet_4v,
  112. mt.track_branch_vec<float>("nJet", "Jet_btagCMVA"),
  113. mt.track_branch_vec< int >("nJet", "Jet_mcMatchFlav")
  114. ), "jets");
  115. return jets;
  116. }
  117. decltype(auto)
  118. construct_leptons_value(MiniTreeDataSet& mt){
  119. mt.track_branch<int>("nLepGood");
  120. mt.track_branch_vec<float>("nLepGood", "LepGood_pt");
  121. mt.track_branch_vec<float>("nLepGood", "LepGood_eta");
  122. mt.track_branch_vec<float>("nLepGood", "LepGood_phi");
  123. mt.track_branch_vec<float>("nLepGood", "LepGood_mass");
  124. auto LepGood_4v = lorentz_vectors("LepGood_pt", "LepGood_eta", "LepGood_phi", "LepGood_mass", "LepGood_4v");
  125. auto leptons = apply(GenFunction::register_function<std::vector<Particle>(std::vector<TLorentzVector>,
  126. std::vector<int>,
  127. std::vector<int>,
  128. std::vector<int>)>("build_reco_leptons",
  129. FUNC(([](const std::vector<TLorentzVector>& vs,
  130. const std::vector<int>& pdgIds,
  131. const std::vector<int>& mcMatchPdgIds,
  132. const std::vector<int>& charges
  133. ){
  134. std::vector<Particle> leptons;
  135. for(int i=0; i<vs.size(); i++){
  136. Particle p = Particle::Lepton(i, vs[i], {pdgIds[i],
  137. mcMatchPdgIds[i],
  138. charges[i]});
  139. leptons.push_back(p);
  140. }
  141. return leptons;
  142. }))), fv::tuple(LepGood_4v,
  143. mt.track_branch_vec< int >("nLepGood", "LepGood_pdgId"),
  144. mt.track_branch_vec< int >("nLepGood", "LepGood_mcMatchPdgId"),
  145. mt.track_branch_vec< int >("nLepGood", "LepGood_charge")
  146. ), "leptons");
  147. return leptons;
  148. }
  149. decltype(auto)
  150. construct_mc_jets_value(MiniTreeDataSet& mt){
  151. mt.track_branch<int>("nGenPart");
  152. mt.track_branch_vec<float>("nGenPart", "GenPart_pt");
  153. mt.track_branch_vec<float>("nGenPart", "GenPart_eta");
  154. mt.track_branch_vec<float>("nGenPart", "GenPart_phi");
  155. mt.track_branch_vec<float>("nGenPart", "GenPart_mass");
  156. auto Jet_4v = lorentz_vectors("GenPart_pt", "GenPart_eta", "GenPart_phi", "GenPart_mass", "GenPart_4v");
  157. energies(Jet_4v, "GenPart_energy");
  158. auto mc_jets = apply(GenFunction::register_function<std::vector<Particle>(std::vector<TLorentzVector>,
  159. std::vector<int>,
  160. std::vector<int>,
  161. std::vector<int>,
  162. std::vector<int>)>("build_mc_jets",
  163. FUNC(([](const std::vector<TLorentzVector>& vs,
  164. const std::vector<int>& pdgIds,
  165. const std::vector<int>& motherIndices,
  166. const std::vector<int>& motherIds,
  167. const std::vector<int>& statuses){
  168. std::vector<Particle> mc_jets;
  169. for(int i=0; i<vs.size(); i++){
  170. Particle p = Particle::GenPart(i, vs[i], {pdgIds[i],
  171. motherIndices[i],
  172. motherIds[i],
  173. statuses[i]});
  174. mc_jets.push_back(p);
  175. }
  176. return mc_jets;
  177. }))), fv::tuple(Jet_4v,
  178. mt.track_branch_vec<int>("nGenPart", "GenPart_pdgId"),
  179. mt.track_branch_vec<int>("nGenPart", "GenPart_motherIndex"),
  180. mt.track_branch_vec<int>("nGenPart", "GenPart_motherId"),
  181. mt.track_branch_vec<int>("nGenPart", "GenPart_status")
  182. ),"mc_jets");
  183. return mc_jets;
  184. }
  185. void create_all_common_values(MiniTreeDataSet& mt){
  186. construct_jets_value(mt);
  187. construct_mc_jets_value(mt);
  188. construct_leptons_value(mt);
  189. }
  190. #endif