Add RNTuple LHCB benchmark

Author: eguiraud

root/tree/tree/CMakeLists.txt

@@ -32,4 +32,10 @@ if(ROOT_root7_FOUND AND rootbench-datafiles)
     SETUP "hadd -f404 ${RB_DATASETDIR}/h1dst-lz4.root ${RB_DATASETDIR}/h1dst-lzma.root"
     SETUP "hadd -f506 ${RB_DATASETDIR}/h1dst-zstd.root ${RB_DATASETDIR}/h1dst-lzma.root"
     SETUP "hadd -f0 ${RB_DATASETDIR}/h1dst-none.root ${RB_DATASETDIR}/h1dst-lzma.root")
+
+  RB_ADD_GBENCHMARK(RNTupleLHCBBenchmarks
+   RNTupleLHCBBenchmarks.cxx
+   LABEL short
+   LIBRARIES Core Hist MathCore RIO Tree ROOTNTuple
+   DOWNLOAD_DATAFILES B2HHH~none.ntuple)
 endif()

root/tree/tree/RNTupleLHCBBenchmarks.cxx

@@ -0,0 +1,94 @@
+#include <ROOT/RNTuple.hxx>
+#include <TH1D.h>
+#include <benchmark/benchmark.h>
+#include <rootbench/RBConfig.h>
+
+namespace {
+constexpr double kKaonMassMeV = 493.677;
+
+double GetP2(double px, double py, double pz)
+{
+   return px * px + py * py + pz * pz;
+}
+
+double GetKE(double px, double py, double pz)
+{
+   double p2 = GetP2(px, py, pz);
+   return sqrt(p2 + kKaonMassMeV * kKaonMassMeV);
+}
+} // namespace
+
+static void BM_RNTuple_LHCB(benchmark::State &state)
+{
+   using RNTupleReader = ROOT::Experimental::RNTupleReader;
+   using RNTupleModel = ROOT::Experimental::RNTupleModel;
+
+   auto model = RNTupleModel::Create();
+   auto ntuple = RNTupleReader::Open(std::move(model), "DecayTree", RB::GetDataDir() + "/B2HHH~none.ntuple");
+   auto viewH1IsMuon = ntuple->GetView<int>("H1_isMuon");
+   auto viewH2IsMuon = ntuple->GetView<int>("H2_isMuon");
+   auto viewH3IsMuon = ntuple->GetView<int>("H3_isMuon");
+
+   auto viewH1PX = ntuple->GetView<double>("H1_PX");
+   auto viewH1PY = ntuple->GetView<double>("H1_PY");
+   auto viewH1PZ = ntuple->GetView<double>("H1_PZ");
+   auto viewH1ProbK = ntuple->GetView<double>("H1_ProbK");
+   auto viewH1ProbPi = ntuple->GetView<double>("H1_ProbPi");
+
+   auto viewH2PX = ntuple->GetView<double>("H2_PX");
+   auto viewH2PY = ntuple->GetView<double>("H2_PY");
+   auto viewH2PZ = ntuple->GetView<double>("H2_PZ");
+   auto viewH2ProbK = ntuple->GetView<double>("H2_ProbK");
+   auto viewH2ProbPi = ntuple->GetView<double>("H2_ProbPi");
+
+   auto viewH3PX = ntuple->GetView<double>("H3_PX");
+   auto viewH3PY = ntuple->GetView<double>("H3_PY");
+   auto viewH3PZ = ntuple->GetView<double>("H3_PZ");
+   auto viewH3ProbK = ntuple->GetView<double>("H3_ProbK");
+   auto viewH3ProbPi = ntuple->GetView<double>("H3_ProbPi");
+
+   auto hMass = new TH1D("B_mass", "", 500, 5050, 5500);
+
+   for (auto _ : state) {
+      hMass->Reset();
+      for (auto i : ntuple->GetEntryRange()) {
+         if (viewH1IsMuon(i) || viewH2IsMuon(i) || viewH3IsMuon(i)) {
+            continue;
+         }
+
+         constexpr double prob_k_cut = 0.5;
+         if (viewH1ProbK(i) < prob_k_cut)
+            continue;
+         if (viewH2ProbK(i) < prob_k_cut)
+            continue;
+         if (viewH3ProbK(i) < prob_k_cut)
+            continue;
+
+         constexpr double prob_pi_cut = 0.5;
+         if (viewH1ProbPi(i) > prob_pi_cut)
+            continue;
+         if (viewH2ProbPi(i) > prob_pi_cut)
+            continue;
+         if (viewH3ProbPi(i) > prob_pi_cut)
+            continue;
+
+         double b_px = viewH1PX(i) + viewH2PX(i) + viewH3PX(i);
+         double b_py = viewH1PY(i) + viewH2PY(i) + viewH3PY(i);
+         double b_pz = viewH1PZ(i) + viewH2PZ(i) + viewH3PZ(i);
+         double b_p2 = GetP2(b_px, b_py, b_pz);
+         double k1_E = GetKE(viewH1PX(i), viewH1PY(i), viewH1PZ(i));
+         double k2_E = GetKE(viewH2PX(i), viewH2PY(i), viewH2PZ(i));
+         double k3_E = GetKE(viewH3PX(i), viewH3PY(i), viewH3PZ(i));
+         double b_E = k1_E + k2_E + k3_E;
+         double b_mass = sqrt(b_E * b_E - b_p2);
+         hMass->Fill(b_mass);
+      }
+   }
+   RB::Ensure(int(hMass->GetMean()) == 5262);
+   RB::Ensure(int(hMass->GetEntries()) == 23895);
+   delete hMass;
+}
+BENCHMARK(BM_RNTuple_LHCB)->Unit(benchmark::kMicrosecond)->Iterations(3);
+
+BENCHMARK_MAIN();
+