Add RNTuple benchmark for ATLAS file

root/tree/tree/CMakeLists.txt

@@ -16,7 +16,11 @@ if(ROOT_root7_FOUND AND rootbench-datafiles AND ROOT_dataframe_FOUND)
             DEPENDS libh1event
   )
 
-  RB_ADD_GBENCHMARK(RNTupleAnalysisBenchmarks
+  RB_ADD_TOOL(gen_atlas ${CMAKE_CURRENT_SOURCE_DIR}/gen_atlas.cxx
+            LIBRARIES Core RIO Tree ROOTNTuple ROOTNTupleUtil
+  )
+
+  RB_ADD_GBENCHMARK(RNTupleH1Benchmarks
     RNTupleH1Benchmarks.cxx
     LABEL short
     LIBRARIES Core Hist MathCore RIO Tree ROOTNTuple
@@ -30,6 +34,15 @@ if(ROOT_root7_FOUND AND rootbench-datafiles AND ROOT_dataframe_FOUND)
     SETUP "${RB_HADD_EXECUTABLE} -f404 ${RB_DATASETDIR}/h1dst-lz4.root ${RB_DATASETDIR}/h1dst-lzma.root"
     SETUP "${RB_HADD_EXECUTABLE} -f506 ${RB_DATASETDIR}/h1dst-zstd.root ${RB_DATASETDIR}/h1dst-lzma.root"
     SETUP "${RB_HADD_EXECUTABLE} -f0 ${RB_DATASETDIR}/h1dst-none.root ${RB_DATASETDIR}/h1dst-lzma.root")
+
+  RB_ADD_GBENCHMARK(RNTupleATLASBenchmarks
+    RNTupleATLASBenchmarks.cxx
+    LABEL short
+    LIBRARIES Core Hist MathCore RIO Tree ROOTNTuple
+    DOWNLOAD_DATAFILES gg_data-zstd.root
+    SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c lzma"
+    SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c zstd"
+    SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c none")
 endif(ROOT_root7_FOUND AND rootbench-datafiles AND ROOT_dataframe_FOUND)
 
 if(ROOT_root7_FOUND AND rootbench-datafiles)

root/tree/tree/RNTupleATLASBenchmarks.cxx

@@ -0,0 +1,118 @@
+#include <benchmark/benchmark.h>
+
+#include <cstdint>
+#include <string>
+#include <vector>
+
+#include <ROOT/RNTupleReader.hxx>
+#include <TH1D.h>
+
+#include <Math/Vector4D.h>
+
+#include <rootbench/RBConfig.h>
+
+
+static float ComputeInvariantMass(float pt0, float pt1, float eta0, float eta1, float phi0, float phi1, float e0, float e1)
+{
+    ROOT::Math::PtEtaPhiEVector p1(pt0, eta0, phi0, e0);
+    ROOT::Math::PtEtaPhiEVector p2(pt1, eta1, phi1, e1);
+    return (p1 + p2).mass() / 1000.0;
+}
+
+
+static void ProcessNTuple(ROOT::RNTupleReader *ntuple, TH1D *hMass, bool isMC)
+{
+
+    auto viewTrigP           = ntuple->GetView<bool>("trigP");
+    auto viewPhotonN         = ntuple->GetView<std::uint32_t>("photon_n");
+    auto viewPhotonIsTightId = ntuple->GetView<std::vector<bool>>("photon_isTightID");
+    auto viewPhotonPt        = ntuple->GetView<std::vector<float>>("photon_pt");
+    auto viewPhotonEta       = ntuple->GetView<std::vector<float>>("photon_eta");
+    auto viewPhotonPhi       = ntuple->GetView<std::vector<float>>("photon_phi");
+    auto viewPhotonE         = ntuple->GetView<std::vector<float>>("photon_E");
+    auto viewPhotonPtCone30  = ntuple->GetView<std::vector<float>>("photon_ptcone30");
+    auto viewPhotonEtCone20  = ntuple->GetView<std::vector<float>>("photon_etcone20");
+
+    auto viewScaleFactorPhoton        = ntuple->GetView<float>("scaleFactor_PHOTON");
+    auto viewScaleFactorPhotonTrigger = ntuple->GetView<float>("scaleFactor_PhotonTRIGGER");
+    auto viewScaleFactorPileUp        = ntuple->GetView<float>("scaleFactor_PILEUP");
+    auto viewMcWeight                 = ntuple->GetView<float>("mcWeight");
+
+    unsigned nevents = 0;
+    for (auto e : ntuple->GetEntryRange()) {
+        nevents++;
+
+        if (!viewTrigP(e)) continue;
+
+        std::vector<size_t> idxGood;
+        auto isTightId = viewPhotonIsTightId(e);
+        auto pt = viewPhotonPt(e);
+        auto eta = viewPhotonEta(e);
+
+        for (size_t i = 0; i < viewPhotonN(e); ++i) {
+            if (!isTightId[i]) continue;
+            if (pt[i] <= 25000.) continue;
+            if (abs(eta[i]) >= 2.37) continue;
+            if (abs(eta[i]) >= 1.37 && abs(eta[i]) <= 1.52) continue;
+            idxGood.push_back(i);
+        }
+        if (idxGood.size() != 2) continue;
+
+        auto ptCone30 = viewPhotonPtCone30(e);
+        auto etCone20 = viewPhotonEtCone20(e);
+
+        bool isIsolatedPhotons = true;
+        for (int i = 0; i < 2; ++i) {
+            if ((ptCone30[idxGood[i]] / pt[idxGood[i]] >= 0.065) ||
+                (etCone20[idxGood[i]] / pt[idxGood[i]] >= 0.065))
+            {
+            isIsolatedPhotons = false;
+            break;
+            }
+        }
+        if (!isIsolatedPhotons) continue;
+
+        auto phi = viewPhotonPhi(e);
+        auto E = viewPhotonE(e);
+
+        float myy = ComputeInvariantMass(
+            pt[idxGood[0]], pt[idxGood[1]],
+            eta[idxGood[0]], eta[idxGood[1]],
+            phi[idxGood[0]], phi[idxGood[1]],
+            E[idxGood[0]], E[idxGood[1]]);
+
+        if (pt[idxGood[0]] / 1000. / myy <= 0.35) continue;
+        if (pt[idxGood[1]] / 1000. / myy <= 0.25) continue;
+        if (myy <= 105) continue;
+        if (myy >= 160) continue;
+
+        if (isMC) {
+            auto weight = viewScaleFactorPhoton(e) * viewScaleFactorPhotonTrigger(e) *
+                        viewScaleFactorPileUp(e) * viewMcWeight(e);
+            hMass->Fill(myy, weight);
+        } else {
+            hMass->Fill(myy);
+        }
+    }
+}
+
+
+static void BM_RNTuple_ATLAS(benchmark::State &state, const std::string &comprAlgorithm)
+{
+    using ROOT::RNTupleReader;
+    std::string path = RB::GetDataDir() + "/atlas-" + comprAlgorithm + ".ntuple";
+
+    for (auto _ : state) {
+        auto hData = new TH1D("", "Diphoton invariant mass; m_{#gamma#gamma} [GeV];Events", 30, 105, 160);
+
+        auto ntuple = RNTupleReader::Open("mini", path);
+        ProcessNTuple(ntuple.get(), hData, false /* isMC */);
+
+        delete hData;
+    }
+}
+BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_LZMA, "lzma")->Unit(benchmark::kMicrosecond)->Iterations(5);
+BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_ZSTD, "zstd")->Unit(benchmark::kMicrosecond)->Iterations(5);
+BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_None, "none")->Unit(benchmark::kMicrosecond)->Iterations(5);
+
+BENCHMARK_MAIN();

root/tree/tree/gen_atlas.cxx

@@ -0,0 +1,75 @@
+#include <ROOT/RNTupleImporter.hxx>
+#include <ROOT/RNTupleWriteOptions.hxx>
+
+#include <TSystem.h>
+
+#include <iostream>
+#include <string>
+
+#include <unistd.h>
+
+using ROOT::Experimental::RNTupleImporter;
+
+int GetCompressionSettings(std::string shorthand)
+{
+   if (shorthand == "zlib")
+      return 101;
+   if (shorthand == "lz4")
+      return 404;
+   if (shorthand == "lzma")
+      return 207;
+   if (shorthand == "zstd")
+      return 505;
+   if (shorthand == "none")
+      return 0;
+   abort();
+}
+
+void Usage(char *progname)
+{
+   std::cout << "Usage: " << progname << " -i <gg_*.root> -o <ntuple-path> -c <compression>" << std::endl;
+}
+
+int main(int argc, char **argv)
+{
+   std::string inputFile = "gg_data.root";
+   std::string outputPath = ".";
+   int compressionSettings = 0;
+   std::string compressionShorthand = "none";
+
+   int c;
+   while ((c = getopt(argc, argv, "hvi:o:c:")) != -1) {
+      switch (c) {
+      case 'h':
+      case 'v':
+         Usage(argv[0]);
+         return 0;
+      case 'i':
+         inputFile = optarg;
+         break;
+      case 'o':
+         outputPath = optarg;
+         break;
+      case 'c':
+         compressionSettings = GetCompressionSettings(optarg);
+         compressionShorthand = optarg;
+         break;
+      default:
+         fprintf(stderr, "Unknown option: -%c\n", c);
+         Usage(argv[0]);
+         return 1;
+      }
+   }
+   std::string outputFile = outputPath + "/atlas-" + compressionShorthand + ".ntuple";
+   if (!gSystem->AccessPathName(outputFile.c_str())) {
+      std::cout << "Output file " << outputFile << " already exists, nothing to do" << std::endl;
+      return 0;
+   }
+   std::cout << "Converting " << inputFile << " --> " << outputFile << std::endl;
+
+   auto importer = RNTupleImporter::Create(inputFile, "mini", outputFile);
+   ROOT::RNTupleWriteOptions options;
+   options.SetCompression(compressionSettings);
+   importer->SetWriteOptions(options);
+   importer->Import();
+}