Skip to content

Commit 83699a6

Browse files
kamahorikamahori
committed
Add RNTuple benchmark for ATLAS file
1 parent 9d751ab commit 83699a6

File tree

3 files changed

+331
-0
lines changed

3 files changed

+331
-0
lines changed

‎root/tree/tree/CMakeLists.txt‎

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,12 @@ if(ROOT_root7_FOUND)
1515
ROOTDataFrame ROOTVecOps ROOTNTuple
1616
DEPENDS libh1event
1717
)
18+
19+
RB_ADD_TOOL(gen_atlas ${CMAKE_CURRENT_SOURCE_DIR}/gen_atlas.cxx
20+
LIBRARIES Core Hist Imt RIO Tree TreePlayer
21+
ROOTDataFrame ROOTVecOps ROOTNTuple
22+
DEPENDS libh1event
23+
)
1824
endif()
1925

2026
if(ROOT_root7_FOUND AND rootbench-datafiles)
@@ -28,4 +34,15 @@ if(ROOT_root7_FOUND AND rootbench-datafiles)
2834
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_h1 -o ${RB_DATASETDIR} -c lz4 ${RB_DATASETDIR}/h1dst-lz4.root"
2935
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_h1 -o ${RB_DATASETDIR} -c zstd ${RB_DATASETDIR}/h1dst-zstd.root"
3036
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_h1 -o ${RB_DATASETDIR} -c none ${RB_DATASETDIR}/h1dst-none.root")
37+
38+
RB_ADD_GBENCHMARK(RNTupleATLASBenchmarks
39+
RNTupleATLASBenchmarks.cxx
40+
LABEL short
41+
LIBRARIES Core Hist MathCore RIO Tree ROOTNTuple
42+
DOWNLOAD_DATAFILES gg_data-zstd.root
43+
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c zlib"
44+
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c lzma"
45+
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c lz4"
46+
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c zstd"
47+
SETUP "${CMAKE_CURRENT_BINARY_DIR}/gen_atlas -i ${RB_DATASETDIR}/gg_data-zstd.root -o ${RB_DATASETDIR} -c none")
3148
endif()

‎root/tree/tree/RNTupleATLASBenchmarks.cxx‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
#include <benchmark/benchmark.h>
2+
3+
#include <fcntl.h>
4+
#include <sys/stat.h>
5+
#include <sys/types.h>
6+
#include <unistd.h>
7+
8+
#include <algorithm>
9+
#include <cassert>
10+
#include <cstdio>
11+
#include <iostream>
12+
#include <future>
13+
#include <memory>
14+
#include <string>
15+
#include <vector>
16+
17+
#include <ROOT/RDataFrame.hxx>
18+
#include <ROOT/RNTuple.hxx>
19+
#include <ROOT/RNTupleDS.hxx>
20+
#include <ROOT/RNTupleOptions.hxx>
21+
#include <Compression.h>
22+
#include <TApplication.h>
23+
#include <TBranch.h>
24+
#include <TCanvas.h>
25+
#include <TClassTable.h>
26+
#include <TF1.h>
27+
#include <TFile.h>
28+
#include <TH1D.h>
29+
#include <TLatex.h>
30+
#include <TLegend.h>
31+
#include <TLorentzVector.h>
32+
#include <TMath.h>
33+
#include <TStyle.h>
34+
#include <TSystem.h>
35+
#include <TTree.h>
36+
#include <TTreeReader.h>
37+
#include <TTreePerfStats.h>
38+
39+
#include <Math/Vector4D.h>
40+
41+
#include <rootbench/RBConfig.h>
42+
43+
44+
static float ComputeInvariantMass(float pt0, float pt1, float eta0, float eta1, float phi0, float phi1, float e0, float e1)
45+
{
46+
ROOT::Math::PtEtaPhiEVector p1(pt0, eta0, phi0, e0);
47+
ROOT::Math::PtEtaPhiEVector p2(pt1, eta1, phi1, e1);
48+
return (p1 + p2).mass() / 1000.0;
49+
}
50+
51+
52+
static void ProcessNTuple(ROOT::Experimental::RNTupleReader *ntuple, TH1D *hMass, bool isMC)
53+
{
54+
55+
auto viewTrigP = ntuple->GetView<bool>("trigP");
56+
auto viewPhotonN = ntuple->GetView<std::uint32_t>("photon_n");
57+
auto viewPhotonIsTightId = ntuple->GetView<std::vector<bool>>("photon_isTightID");
58+
auto viewPhotonPt = ntuple->GetView<std::vector<float>>("photon_pt");
59+
auto viewPhotonEta = ntuple->GetView<std::vector<float>>("photon_eta");
60+
auto viewPhotonPhi = ntuple->GetView<std::vector<float>>("photon_phi");
61+
auto viewPhotonE = ntuple->GetView<std::vector<float>>("photon_E");
62+
auto viewPhotonPtCone30 = ntuple->GetView<std::vector<float>>("photon_ptcone30");
63+
auto viewPhotonEtCone20 = ntuple->GetView<std::vector<float>>("photon_etcone20");
64+
65+
auto viewScaleFactorPhoton = ntuple->GetView<float>("scaleFactor_PHOTON");
66+
auto viewScaleFactorPhotonTrigger = ntuple->GetView<float>("scaleFactor_PhotonTRIGGER");
67+
auto viewScaleFactorPileUp = ntuple->GetView<float>("scaleFactor_PILEUP");
68+
auto viewMcWeight = ntuple->GetView<float>("mcWeight");
69+
70+
unsigned nevents = 0;
71+
for (auto e : ntuple->GetEntryRange()) {
72+
nevents++;
73+
74+
if (!viewTrigP(e)) continue;
75+
76+
std::vector<size_t> idxGood;
77+
auto isTightId = viewPhotonIsTightId(e);
78+
auto pt = viewPhotonPt(e);
79+
auto eta = viewPhotonEta(e);
80+
81+
for (size_t i = 0; i < viewPhotonN(e); ++i) {
82+
if (!isTightId[i]) continue;
83+
if (pt[i] <= 25000.) continue;
84+
if (abs(eta[i]) >= 2.37) continue;
85+
if (abs(eta[i]) >= 1.37 && abs(eta[i]) <= 1.52) continue;
86+
idxGood.push_back(i);
87+
}
88+
if (idxGood.size() != 2) continue;
89+
90+
auto ptCone30 = viewPhotonPtCone30(e);
91+
auto etCone20 = viewPhotonEtCone20(e);
92+
93+
bool isIsolatedPhotons = true;
94+
for (int i = 0; i < 2; ++i) {
95+
if ((ptCone30[idxGood[i]] / pt[idxGood[i]] >= 0.065) ||
96+
(etCone20[idxGood[i]] / pt[idxGood[i]] >= 0.065))
97+
{
98+
isIsolatedPhotons = false;
99+
break;
100+
}
101+
}
102+
if (!isIsolatedPhotons) continue;
103+
104+
auto phi = viewPhotonPhi(e);
105+
auto E = viewPhotonE(e);
106+
107+
float myy = ComputeInvariantMass(
108+
pt[idxGood[0]], pt[idxGood[1]],
109+
eta[idxGood[0]], eta[idxGood[1]],
110+
phi[idxGood[0]], phi[idxGood[1]],
111+
E[idxGood[0]], E[idxGood[1]]);
112+
113+
if (pt[idxGood[0]] / 1000. / myy <= 0.35) continue;
114+
if (pt[idxGood[1]] / 1000. / myy <= 0.25) continue;
115+
if (myy <= 105) continue;
116+
if (myy >= 160) continue;
117+
118+
if (isMC) {
119+
auto weight = viewScaleFactorPhoton(e) * viewScaleFactorPhotonTrigger(e) *
120+
viewScaleFactorPileUp(e) * viewMcWeight(e);
121+
hMass->Fill(myy, weight);
122+
} else {
123+
hMass->Fill(myy);
124+
}
125+
}
126+
}
127+
128+
129+
static void BM_RNTuple_ATLAS(benchmark::State &state, const std::string &comprAlgorithm)
130+
{
131+
using RNTupleReader = ROOT::Experimental::RNTupleReader;
132+
std::string path = RB::GetDataDir() + "/atlas-" + comprAlgorithm + ".ntuple";
133+
134+
for (auto _ : state) {
135+
auto hData = new TH1D("", "Diphoton invariant mass; m_{#gamma#gamma} [GeV];Events", 30, 105, 160);
136+
137+
auto ntuple = RNTupleReader::Open("mini", path);
138+
ProcessNTuple(ntuple.get(), hData, false /* isMC */);
139+
140+
delete hData;
141+
}
142+
}
143+
BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_LZ4, "lz4")->Unit(benchmark::kMicrosecond)->Iterations(5);
144+
BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_ZLIB, "zlib")->Unit(benchmark::kMicrosecond)->Iterations(5);
145+
BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_LZMA, "lzma")->Unit(benchmark::kMicrosecond)->Iterations(5);
146+
BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_ZSTD, "zstd")->Unit(benchmark::kMicrosecond)->Iterations(5);
147+
BENCHMARK_CAPTURE(BM_RNTuple_ATLAS, BM_RNTuple_ATLAS_None, "none")->Unit(benchmark::kMicrosecond)->Iterations(5);
148+
149+
BENCHMARK_MAIN();

‎root/tree/tree/gen_atlas.cxx‎

Lines changed: 165 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,165 @@
1+
#include <ROOT/RField.hxx>
2+
#include <ROOT/RNTuple.hxx>
3+
#include <ROOT/RNTupleModel.hxx>
4+
#include <ROOT/RNTupleOptions.hxx>
5+
6+
#include <TBranch.h>
7+
#include <TBranchElement.h>
8+
#include <TBranchSTL.h>
9+
#include <TCanvas.h>
10+
#include <TFile.h>
11+
#include <TH1F.h>
12+
#include <TLeaf.h>
13+
#include <TTree.h>
14+
15+
#include <cassert>
16+
#include <iostream>
17+
#include <memory>
18+
#include <string>
19+
#include <vector>
20+
21+
#include <unistd.h>
22+
23+
using RNTupleModel = ROOT::Experimental::RNTupleModel;
24+
using RFieldBase = ROOT::Experimental::Detail::RFieldBase;
25+
using RNTupleWriter = ROOT::Experimental::RNTupleWriter;
26+
using RNTupleWriteOptions = ROOT::Experimental::RNTupleWriteOptions;
27+
28+
int GetCompressionSettings(std::string shorthand)
29+
{
30+
if (shorthand == "zlib")
31+
return 101;
32+
if (shorthand == "lz4")
33+
return 404;
34+
if (shorthand == "lzma")
35+
return 207;
36+
if (shorthand == "zstd")
37+
return 505;
38+
if (shorthand == "none")
39+
return 0;
40+
abort();
41+
}
42+
43+
44+
static void SplitPath(const std::string &path, std::string *basename, std::string *suffix)
45+
{
46+
size_t idx_dot = path.find_last_of(".");
47+
if (idx_dot == std::string::npos) {
48+
*basename = path;
49+
suffix->clear();
50+
} else {
51+
*basename = path.substr(0, idx_dot);
52+
*suffix = path.substr(idx_dot + 1);
53+
}
54+
}
55+
56+
57+
std::string StripSuffix(const std::string &path) {
58+
std::string basename;
59+
std::string suffix;
60+
SplitPath(path, &basename, &suffix);
61+
return basename;
62+
}
63+
64+
65+
std::string GetFileName(const std::string &path) {
66+
const std::string::size_type idx = path.find_last_of('/');
67+
if (idx != std::string::npos)
68+
return path.substr(idx+1);
69+
else
70+
return path;
71+
}
72+
73+
74+
void Usage(char *progname) {
75+
std::cout << "Usage: " << progname << " -i <gg_*.root> -o <ntuple-path> -c <compression>" << std::endl;
76+
}
77+
78+
79+
int main(int argc, char **argv) {
80+
std::string inputFile = "gg_data.root";
81+
std::string outputPath = ".";
82+
int compressionSettings = 0;
83+
std::string compressionShorthand = "none";
84+
85+
int c;
86+
while ((c = getopt(argc, argv, "hvi:o:c:")) != -1) {
87+
switch (c) {
88+
case 'h':
89+
case 'v':
90+
Usage(argv[0]);
91+
return 0;
92+
case 'i':
93+
inputFile = optarg;
94+
break;
95+
case 'o':
96+
outputPath = optarg;
97+
break;
98+
case 'c':
99+
compressionSettings = GetCompressionSettings(optarg);
100+
compressionShorthand = optarg;
101+
break;
102+
default:
103+
fprintf(stderr, "Unknown option: -%c\n", c);
104+
Usage(argv[0]);
105+
return 1;
106+
}
107+
}
108+
std::string outputFile = outputPath + "/" + "atlas" + "-" + compressionShorthand + ".ntuple";
109+
110+
std::unique_ptr<TFile> f(TFile::Open(inputFile.c_str()));
111+
assert(f && ! f->IsZombie());
112+
113+
auto model = RNTupleModel::Create();
114+
115+
auto tree = f->Get<TTree>("mini");
116+
for (auto b : TRangeDynCast<TBranch>(*tree->GetListOfBranches())) {
117+
assert(b);
118+
119+
TLeaf *l = static_cast<TLeaf*>(b->GetListOfLeaves()->First());
120+
121+
auto field = RFieldBase::Create(l->GetName(), l->GetTypeName());
122+
123+
if (typeid(*b) == typeid(TBranchSTL) || typeid(*b) == typeid(TBranchElement)) {
124+
if (field->GetType() == "std::vector<bool>") {
125+
std::vector<bool> **v = new std::vector<bool> *();
126+
tree->SetBranchAddress(b->GetName(), v);
127+
model->GetDefaultEntry()->CaptureValue(field->CaptureValue(*v));
128+
model->GetRootField()->Attach(std::unique_ptr<RFieldBase>(field));
129+
} else if (field->GetType() == "std::vector<float>") {
130+
std::vector<float> **v = new std::vector<float> *();
131+
tree->SetBranchAddress(b->GetName(), v);
132+
model->GetDefaultEntry()->CaptureValue(field->CaptureValue(*v));
133+
model->GetRootField()->Attach(std::unique_ptr<RFieldBase>(field));
134+
} else if (field->GetType() == "std::vector<std::int32_t>") {
135+
std::vector<std::int32_t> **v = new std::vector<std::int32_t> *();
136+
tree->SetBranchAddress(b->GetName(), v);
137+
model->GetDefaultEntry()->CaptureValue(field->CaptureValue(*v));
138+
model->GetRootField()->Attach(std::unique_ptr<RFieldBase>(field));
139+
} else if (field->GetType() == "std::vector<std::uint32_t>") {
140+
std::vector<std::uint32_t> **v = new std::vector<std::uint32_t> *();
141+
tree->SetBranchAddress(b->GetName(), v);
142+
model->GetDefaultEntry()->CaptureValue(field->CaptureValue(*v));
143+
model->GetRootField()->Attach(std::unique_ptr<RFieldBase>(field));
144+
} else {
145+
assert(false);
146+
}
147+
} else {
148+
model->AddField(std::unique_ptr<RFieldBase>(field));
149+
void *fieldDataPtr = model->GetDefaultEntry()->GetValue(l->GetName()).GetRawPtr();
150+
tree->SetBranchAddress(b->GetName(), fieldDataPtr);
151+
}
152+
}
153+
154+
RNTupleWriteOptions options;
155+
options.SetCompression(compressionSettings);
156+
auto ntuple = RNTupleWriter::Recreate(std::move(model), "mini", outputFile, options);
157+
158+
auto nEntries = tree->GetEntries();
159+
for (decltype(nEntries) i = 0; i < nEntries; ++i) {
160+
tree->GetEntry(i);
161+
ntuple->Fill();
162+
}
163+
164+
tree->ResetBranchAddresses();
165+
}

0 commit comments

Comments
 (0)