MVA_Creation.cpp 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156
  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. * MVA Creation script. This file demonstrates how to use filval_root's TMVA integration.
  33. */
  34. #include <iostream>
  35. #include <vector>
  36. #include <utility>
  37. #include <numeric>
  38. #include <limits>
  39. #include "filval/filval.hpp"
  40. #include "filval/root/root_filval.hpp"
  41. #include "MiniTreeDataSet.hpp"
  42. #include <TSystem.h>
  43. #define PI 3.14159
  44. #define W_MASS 80.385 // GeV/c^2
  45. #define Z_MASS 91.188 // GeV/c^2
  46. #define T_MASS 172.44 // GeV/c^2
  47. using namespace std;
  48. using namespace fv;
  49. using namespace fv::root;
  50. void enable_branches(MiniTreeDataSet& mt){
  51. mt.track_branch<int>("nLepGood");
  52. mt.track_branch<int>("nJet");
  53. mt.track_branch<int>("nGenPart");
  54. mt.track_branch<int>("nBJetLoose40");
  55. mt.track_branch<int>("nBJetMedium40");
  56. mt.track_branch<int>("nBJetTight40");
  57. }
  58. #define MVA_DTYPES int, int, int, int, int
  59. void declare_values(MiniTreeDataSet& mt){
  60. auto event_number = mt.get_current_event_number();
  61. auto is_training = fv::apply(fv::GenFunction::reg_func<bool(int)>("is_odd",
  62. FUNC(([](int n){
  63. return (n%2) == 1;
  64. }))), fv::tuple(event_number));
  65. auto is_signal = fv::bound(fv::GenFunction::reg_func<bool()>("is_signal",
  66. FUNC(([mt=&mt](){
  67. const std::string& label = mt->get_current_event_label();
  68. return label == "signal";
  69. }))), "is_signal");
  70. auto weight = fv::constant<double>("1", 1);
  71. auto mva_data = fv::root::mva_data<MVA_DTYPES>(is_training, is_signal, weight,
  72. {"nJet", lookup<int>("nJet")},
  73. {"nBJetLoose40", lookup<int>("nBJetLoose40")},
  74. {"nBJetMedium40", lookup<int>("nBJetMedium40")},
  75. {"nBJetTight40", lookup<int>("nBJetTight40")},
  76. {"nLepGood", lookup<int>("nLepGood")}
  77. );
  78. GenValue::alias("mva_data", mva_data);
  79. }
  80. void declare_containers(MiniTreeDataSet& mt){
  81. auto mva_data = (MVAData<MVA_DTYPES>*)lookup<MVAData<MVA_DTYPES>::type>("mva_data");
  82. auto mva = mt.register_container<MVA<MVA_DTYPES>>("my_mva", mva_data);
  83. mva->add_method("KNN", "H:nkNN=20:ScaleFrac=0.8:SigmaFact=1.0:Kernel=Gaus:UseKernel=F:UseWeight=T:!Trim");
  84. }
  85. void create_mva(const std::string& output_filename, const std::vector<std::string>& sig_files, const std::vector<std::string>& bg_files, bool silent){
  86. gSystem->Load("libfilval.so");
  87. auto replace_suffix = [](const std::string& input, const std::string& new_suffix){
  88. return input.substr(0, input.find_last_of(".")) + new_suffix;
  89. };
  90. string log_filename = replace_suffix(output_filename, ".log");
  91. fv::util::Log::init_logger(log_filename, fv::util::LogPriority::kLogDebug);
  92. std::map<std::string, std::string> filenames_with_labels;
  93. for (const std::string& fname : sig_files){
  94. filenames_with_labels[fname] = "signal";
  95. }
  96. for (const std::string& fname : bg_files){
  97. filenames_with_labels[fname] = "background";
  98. }
  99. MiniTreeDataSet mt(output_filename, filenames_with_labels);
  100. enable_branches(mt);
  101. declare_values(mt);
  102. declare_containers(mt);
  103. mt.process(silent);
  104. mt.save_all();
  105. }
  106. int main(int argc, char * argv[])
  107. {
  108. fv::util::ArgParser args(argc, argv);
  109. if(args.cmdOptionExists("-h")) {
  110. cout << "Usage: ./mva (-s) -out outfile.root -sig [signal_minitree.root]+ -bg [background_minitree.root]+" << endl;
  111. return 0;
  112. }
  113. bool silent = args.cmdOptionExists("-s");
  114. string output_filename = args.getCmdOption("-out");
  115. std::vector<std::string> sig_files;
  116. std::vector<std::string> bg_files;
  117. std::vector<std::string>* cur_flist = nullptr;
  118. for(int i=1; i<argc; i++){
  119. if (!strncmp(argv[i], "-sig", 4)){
  120. cur_flist = &sig_files;
  121. }
  122. else if (!strncmp(argv[i], "-bg", 3)){
  123. cur_flist = &bg_files;
  124. continue;
  125. }
  126. else if (!strncmp(argv[i], "-out", 4)){
  127. cur_flist = nullptr;
  128. }
  129. else if (cur_flist != nullptr){
  130. cur_flist->push_back(argv[i]);
  131. }
  132. }
  133. create_mva(output_filename, sig_files, bg_files, silent);
  134. }