Weights (#21)

* provide an example of Integral-1D plot * enable weights in filling * provide an example of histo with weights * change default AT version to v2.3.4
HeavyIonAnalysis · Feb 13, 2024 · a6c214e · a6c214e
1 parent 4fd3941
commit a6c214e
Show file tree

Hide file tree

Showing 7 changed files with 84 additions and 35 deletions.
diff --git a/.gitignore b/.gitignore
@@ -39,6 +39,7 @@
 .idea/misc.xml
 .idea/modules.xml
 .idea/vcs.xml
+.idea/
 
 *build*/
-install/
+install/
diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -13,7 +13,7 @@ set(AnalysisTreeQA_BUILD_TASKS OFF CACHE BOOL "Build user' AnalysisTreeQA tasks
 set(AnalysisTreeQA_BUILD_EXAMPLES ON CACHE BOOL "Build AnalysisTreeQA examples (examples/)")
 set(AnalysisTreeQA_BUILD_TESTS OFF CACHE BOOL "Build tests for AnalysisTreeQA")
 set(AnalysisTreeQA_BUNDLED_AT ON CACHE BOOL "Get and build AnalysisTree")
-set(AnalysisTreeQA_BUNDLED_AT_VERSION "v2.2.8" CACHE STRING "Bundled AnalysisTree version")
+set(AnalysisTreeQA_BUNDLED_AT_VERSION "v2.3.4" CACHE STRING "Bundled AnalysisTree version")
 set(AnalysisTreeQA_BUNDLED_AT_GIT_SHALLOW ON CACHE BOOL "Use CMake GIT_SHALLOW option")
 set(AnalysisTreeQA_BUNDLED_CUTS ON CACHE BOOL "Get and build AnalysisTreeCuts")
 set(AnalysisTreeQA_BUNDLED_CUTS_VERSION "v0.0.1" CACHE STRING "Bundled AnalysisTreeCuts version")

diff --git a/examples/example.cpp b/examples/example.cpp
@@ -20,6 +20,9 @@ void example(const std::string& filelist){
   Cuts* pT_cut =  new Cuts("pT_cut", {RangeCut("VtxTracks.pT", 1, 1.5)});
   task->AddH1({"p_{T}, GeV/c", Variable::FromString("VtxTracks.pT"), {100, 0, 3}}, pT_cut);
 
+  // 1D histo with weight
+  task->AddH1({"PSD module id", Variable::FromString("PsdModules.id"), {51, -1, 50}}, nullptr, Variable::FromString("PsdModules.signal"));
+
   // AnalysisTree::Variable in case of more complicated plot
   Variable chi2_over_ndf("#chi^{2}/NDF", {{"VtxTracks", "chi2"}, {"VtxTracks", "ndf"}}, []( std::vector<double>& var ) { return var.at(0)/var.at(1); });
   task->AddH1({"#chi^{2}/NDF", chi2_over_ndf, {100, 0, 10}});
@@ -44,9 +47,18 @@ void example(const std::string& filelist){
   task->AddProfile({"#it{y}", Variable::FromString("SimParticles.rapidity"), {20, 0.5, 2.5}}, {"MC-protons v_{1}", v1, {}}, mc_protons);
   task->AddProfile({"#it{y}", Variable::FromString("SimParticles.rapidity"), {20, 0.5, 2.5}}, {"MC-protons v_{1}", v1, {}}, mc_pions);
 
+  //Integral plot 1D
+  SimpleCut vtx_chi2_track_cut = RangeCut({"VtxTracks.vtx_chi2"}, 0, 3);
+  SimpleCut nhits_cut = RangeCut({"VtxTracks.nhits"}, 4, 100);
+  SimpleCut chi2_cut({"VtxTracks.chi2", "VtxTracks.ndf"}, [](std::vector<double> par){ return par[0]/par[1] < 10; });
+  SimpleCut eta_cut = RangeCut({"VtxTracks.eta"}, 0, 3);
+  Cuts* vertex_tracks_cuts = new Cuts("GoodCentralityTracks", {vtx_chi2_track_cut, nhits_cut, chi2_cut, eta_cut});
+  task->AddIntegral({"Multiplicity", Variable::FromString("VtxTracks.ones"), {1000, 0, 1000}}, vertex_tracks_cuts);
+
   man->AddTask(task);
 
   man->Init({filelist}, {"rTree"});
+  man->SetVerbosityPeriod(100);
   man->Run(-1);
   man->Finish();
 }
@@ -62,4 +74,4 @@ int main(int argc, char* argv[]){
   example(filename);
 
   return 0;
-}
+}
diff --git a/src/EntryConfig.cpp b/src/EntryConfig.cpp
@@ -8,25 +8,34 @@
 namespace AnalysisTree {
 namespace QA {
 
-struct fill_struct : public Utils::Visitor<void> {
-  fill_struct(double val1, double val2) : val1_(val1), val2_(val2) {}
-  void operator()(TH1*) const { throw std::runtime_error("Cannot apply Fill(va1, val2) to TH1"); }
+struct fill2_struct : public Utils::Visitor<void> {
+  fill2_struct(double val1, double val2) : val1_(val1), val2_(val2) {}
+  void operator()(TH1* h1) const { h1->Fill(val1_, val2_); }
   void operator()(TH2* h2) const { h2->Fill(val1_, val2_); }
   void operator()(TProfile* p) const { p->Fill(val1_, val2_); }
   double val1_, val2_;
 };
 
+struct fill3_struct : public Utils::Visitor<void> {
+  fill3_struct(double val1, double val2, double val3) : val1_(val1), val2_(val2), val3_(val3) {}
+  void operator()(TH1*) const { throw std::runtime_error("EntryConfig, fill3_struct: cannot fill TH1 with 3 arguments"); }
+  void operator()(TH2* h2) const { h2->Fill(val1_, val2_, val3_); }
+  void operator()(TProfile* p) const { p->Fill(val1_, val2_, val3_); }
+  double val1_, val2_, val3_;
+};
+
 struct write_struct : public Utils::Visitor<void> {
   explicit write_struct(std::string n) : name_(std::move(n)) {}
   template<class PlotType>
   void operator()(PlotType* p) const { p->Write(name_.c_str()); }
   std::string name_;
 };
 
-EntryConfig::EntryConfig(const Axis& axis, Cuts* cuts, bool is_integral)
+EntryConfig::EntryConfig(const Axis& axis, Variable& weight, Cuts* cuts, bool is_integral)
   : name_(axis.GetName()),
     type_(is_integral ? PlotType::kIntegral1D : PlotType::kHisto1D),
     axes_({axis}),
+    var4weight_(weight),
     entry_cuts_(cuts) {
   if (cuts)
     name_ += "_" + cuts->GetName();
@@ -36,8 +45,9 @@ EntryConfig::EntryConfig(const Axis& axis, Cuts* cuts, bool is_integral)
   InitPlot();
 }
 
-EntryConfig::EntryConfig(const Axis& x, const Axis& y, Cuts* cuts, bool is_profile) : type_(is_profile ? PlotType::kProfile : PlotType::kHisto2D),
+EntryConfig::EntryConfig(const Axis& x, const Axis& y, Variable& weight, Cuts* cuts, bool is_profile) : type_(is_profile ? PlotType::kProfile : PlotType::kHisto2D),
                                                                                       axes_({x, y}),
+                                                                                      var4weight_(weight),
                                                                                       entry_cuts_(cuts) {
   Set2DName();
   InitPlot();
@@ -127,7 +137,11 @@ void EntryConfig::Set2DName() {
 }
 
 void EntryConfig::Fill(double value1, double value2) {
-  ANALYSISTREE_UTILS_VISIT(fill_struct(value1, value2), plot_);
+  ANALYSISTREE_UTILS_VISIT(fill2_struct(value1, value2), plot_);
+}
+
+void EntryConfig::Fill(double value1, double value2, double value3) {
+  ANALYSISTREE_UTILS_VISIT(fill3_struct(value1, value2, value3), plot_);
 }
 
 void EntryConfig::Write() const {

diff --git a/src/EntryConfig.hpp b/src/EntryConfig.hpp
@@ -22,8 +22,6 @@ class Axis : public Variable, public TAxis {
   Axis(const std::string& title, const Variable& var, const TAxis& a) : Variable(var), TAxis(a) {
     this->SetTitle(title.c_str());
     if(this->GetFields().size() == 1 && this->GetFields().at(0).GetName() == "ones"){
-//      fields_[0] = Field(fields_[0].GetBranchName(), "ones");
-//      fields_.clear();
       this->lambda_ = [](const std::vector<double>& ){ return 1; };
       this->name_ = "Ones";
     }
@@ -48,8 +46,8 @@ class EntryConfig {
   };
 
   EntryConfig() = default;
-  explicit EntryConfig(const Axis& axis, Cuts* cuts = nullptr, bool is_integral = false);
-  EntryConfig(const Axis& x, const Axis& y, Cuts* cuts = nullptr, bool is_profile = false);
+  explicit EntryConfig(const Axis& axis, [[maybe_unused]] Variable& weight, Cuts* cuts = nullptr, bool is_integral = false);
+  EntryConfig(const Axis& x, const Axis& y, Variable& weight, Cuts* cuts = nullptr, bool is_profile = false);
   EntryConfig(const Axis& x, Cuts* cuts_x, const Axis& y, Cuts* cuts_y);
 
   EntryConfig(const EntryConfig&) = default;
@@ -60,6 +58,7 @@ class EntryConfig {
 
   void Fill(double value);
   void Fill(double value1, double value2);
+  void Fill(double value1, double value2, double value3);
   void Write() const;
 
   ANALYSISTREE_ATTR_NODISCARD const std::vector<Axis>& GetAxes() const { return axes_; }
@@ -80,6 +79,10 @@ class EntryConfig {
     return vars;
   }
 
+  ANALYSISTREE_ATTR_NODISCARD Variable GetVariableForWeight() const {
+    return var4weight_;
+  }
+
   ANALYSISTREE_ATTR_NODISCARD std::string GetDirectoryName() const;
 
   void SetOutDir(TDirectory* out_dir) { out_dir_ = out_dir; }
@@ -99,6 +102,7 @@ class EntryConfig {
   PlotType type_{PlotType(-1)};
 
   std::vector<Axis> axes_{};
+  Variable var4weight_{};
   Cuts* entry_cuts_{nullptr};
   std::vector<std::pair<int, int>> vars_id_{};
 

diff --git a/src/Task.cpp b/src/Task.cpp
@@ -36,15 +36,21 @@ void Task::Exec() {
       continue;
     }
     auto var_ids = plot.GetVariablesId();
+    auto weights = this->GetWeights(var_ids.at(0).first);
+    int ivw{-1};
     for (auto var : this->GetValues(var_ids.at(0).first)) {
-
+      ++ivw;
+      auto weight = weights.at(ivw);
+      if(std::fabs(weight) < 1e-6) continue;
       switch (plot.GetNdimentions()) {
         case 1: {
-          plot.Fill(var[var_ids.at(0).second]);
+          if(std::fabs(weight - 1) < 1e-4) plot.Fill(var[var_ids.at(0).second]);
+          else                                plot.Fill(var[var_ids.at(0).second], weight);
           break;
         }
         case 2: {
-          plot.Fill(var[var_ids.at(0).second], var[var_ids.at(1).second]);
+          if(std::fabs(weight - 1) < 1e-4) plot.Fill(var[var_ids.at(0).second], var[var_ids.at(1).second]);
+          else                                plot.Fill(var[var_ids.at(0).second], var[var_ids.at(1).second], weight);
           break;
         }
       }
@@ -79,12 +85,12 @@ void Task::Init() {
   }
 }
 
-size_t Task::AddAnalysisEntry(const Axis& a, Cuts* cuts, bool is_integral){
-  entries_.emplace_back(EntryConfig(a, cuts, is_integral));
-  auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
-  entries_.back().SetVariablesId({{var_id.first, var_id.second.at(0)}});
-  return entries_.size() - 1;
-}
+//size_t Task::AddAnalysisEntry(const Axis& a, Cuts* cuts, bool is_integral){
+//  entries_.emplace_back(EntryConfig(a, cuts, is_integral));
+//  auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
+//  entries_.back().SetVariablesId({{var_id.first, var_id.second.at(0)}});
+//  return entries_.size() - 1;
+//}
 
 
 }// namespace QA

diff --git a/src/Task.hpp b/src/Task.hpp
@@ -20,30 +20,42 @@ class Task : public AnalysisTask {
   void Exec() override;
   void Finish() override;
 
-  size_t AddH1(const Axis& x, Cuts* cuts = nullptr) {
-    entries_.emplace_back(EntryConfig(x, cuts, false));
-    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
+  size_t AddH1(const Axis& x, Cuts* cuts = nullptr, Variable weight = Variable{}) {
+    if(weight.GetNumberOfBranches() == 0) {
+      weight = Variable::FromString(x.GetBranchName() + ".ones");
+    }
+    entries_.emplace_back(EntryConfig(x, weight, cuts, false));
+    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts(), entries_.back().GetVariableForWeight()));
     entries_.back().SetVariablesId({{var_id.first, var_id.second.at(0)}});
     return entries_.size() - 1;
   }
 
-  size_t AddH2(const Axis& x, const Axis& y, Cuts* cuts = nullptr) {
-    entries_.emplace_back(EntryConfig(x, y, cuts));
-    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
+  size_t AddH2(const Axis& x, const Axis& y, Cuts* cuts = nullptr, Variable weight = Variable{}) {
+    if(weight.GetNumberOfBranches() == 0) {
+      weight = Variable::FromString(x.GetBranchName() + ".ones");
+    }
+    entries_.emplace_back(EntryConfig(x, y, weight, cuts));
+    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts(), entries_.back().GetVariableForWeight()));
     entries_.back().SetVariablesId({ {var_id.first, var_id.second.at(0)}, {var_id.first, var_id.second.at(1)} });
     return entries_.size() - 1;
   }
 
-  size_t AddProfile(const Axis& x, const Axis& y, Cuts* cuts = nullptr) {
-    entries_.emplace_back(EntryConfig(x, y, cuts, true));
-    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
+  size_t AddProfile(const Axis& x, const Axis& y, Cuts* cuts = nullptr, Variable weight = Variable{}) {
+    if(weight.GetNumberOfBranches() == 0) {
+      weight = Variable::FromString(x.GetBranchName() + ".ones");
+    }
+    entries_.emplace_back(EntryConfig(x, y, weight, cuts, true));
+    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts(), entries_.back().GetVariableForWeight()));
     entries_.back().SetVariablesId({ {var_id.first, var_id.second.at(0)}, {var_id.first, var_id.second.at(1)} });
     return entries_.size() - 1;
   }
 
-  size_t AddIntegral(const Axis& x, Cuts* cuts = nullptr) {
-    entries_.emplace_back(EntryConfig(x, cuts, true));
-    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts()));
+  size_t AddIntegral(const Axis& x, Cuts* cuts = nullptr, Variable weight = Variable{}) {
+    if(weight.GetNumberOfBranches() == 0) {
+      weight = Variable::FromString(x.GetBranchName() + ".ones");
+    }
+    entries_.emplace_back(EntryConfig(x, weight, cuts, true));
+    auto var_id = AddEntry(AnalysisEntry(entries_.back().GetVariables(), entries_.back().GetEntryCuts(), entries_.back().GetVariableForWeight()));
     entries_.back().SetVariablesId({{var_id.first, var_id.second.at(0)}});
     return entries_.size() - 1;
   }
@@ -61,7 +73,7 @@ class Task : public AnalysisTask {
 
  private:
   void FillIntegral(EntryConfig& plot);
-  size_t AddAnalysisEntry(const Axis& a, Cuts* cuts, bool is_integral);
+//  size_t AddAnalysisEntry(const Axis& a, Cuts* cuts, bool is_integral);
 
   std::vector<EntryConfig> entries_{};
   std::map<std::string, TDirectory*> dir_map_{};