Skip to content

Commit 7c1166b

Browse files
jackal1-66jackal1-66
committed
First evtpool merger
1 parent 16b6488 commit 7c1166b

2 files changed

Lines changed: 211 additions & 0 deletions

File tree

‎Generators/CMakeLists.txt‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -165,6 +165,14 @@ o2_add_test_root_macro(share/egconfig/pythia8_userhooks_charm.C
165165
LABELS generators)
166166
endif()
167167

168+
o2_add_executable(merge-evtpool
169+
COMPONENT_NAME generators
170+
SOURCES src/MergeEventPool.cxx
171+
PUBLIC_LINK_LIBRARIES O2::SimulationDataFormat
172+
ROOT::Core
173+
ROOT::Tree
174+
Boost::program_options)
175+
168176
o2_data_file(COPY share/external DESTINATION Generators)
169177
o2_data_file(COPY share/egconfig DESTINATION Generators)
170178
o2_data_file(COPY share/TPCLoopers DESTINATION Generators)

‎Generators/src/MergeEventPool.cxx‎

Lines changed: 203 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,203 @@
1+
// Copyright 2019-2026 CERN and copyright holders of ALICE O2.
2+
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
3+
// All rights not expressly granted are reserved.
4+
//
5+
// This software is distributed under the terms of the GNU General Public
6+
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
7+
//
8+
// In applying this license CERN does not waive the privileges and immunities
9+
// granted to it by virtue of its status as an Intergovernmental Organization
10+
// or submit itself to any jurisdiction.
11+
12+
/// \brief Merges multiple event-pool files (e.g. evtpool.root / genevents_Kine.root,
13+
/// produced by o2-sim --noGeant) into a single "o2sim" tree.
14+
///
15+
/// Unlike a naive TFileMerger-style concatenation (as done by O2DPG's root_merger.py),
16+
/// this tool renumbers MCEventHeader's event ID so that it stays unique and monotonically
17+
/// increasing across the merged output, instead of resetting per input file. No other
18+
/// remapping is needed: each tree entry already holds one full, self-contained event, so
19+
/// MCTrack mother/daughter indices (which are local to that entry) remain valid as-is.
20+
21+
#include "SimulationDataFormat/MCTrack.h"
22+
#include "SimulationDataFormat/MCEventHeader.h"
23+
#include "SimulationDataFormat/TrackReference.h"
24+
#include <fairlogger/Logger.h>
25+
#include <TFile.h>
26+
#include <TTree.h>
27+
#include <TBranch.h>
28+
#include <boost/program_options.hpp>
29+
#include <filesystem>
30+
#include <memory>
31+
#include <string>
32+
#include <vector>
33+
34+
namespace bpo = boost::program_options;
35+
namespace fs = std::filesystem;
36+
37+
namespace
38+
{
39+
const char* kTrackBranch = "MCTrack";
40+
const char* kHeaderBranch = "MCEventHeader.";
41+
const char* kTrackRefBranch = "TrackRefs";
42+
43+
// Checks that every input file exists, is readable, and has a tree with the branches
44+
// this tool needs to merge (MCTrack + MCEventHeader are mandatory, TrackRefs optional
45+
// but must be consistently present or absent across all files). Nothing is written
46+
// until all inputs pass this check.
47+
bool checkFiles(std::vector<std::string> const& files, std::string const& treename, bool& hasTrackRefs)
48+
{
49+
bool ok = true;
50+
bool first = true;
51+
for (auto const& f : files) {
52+
if (!fs::exists(f)) {
53+
LOG(error) << "Input file " << f << " does not exist";
54+
ok = false;
55+
continue;
56+
}
57+
std::unique_ptr<TFile> file(TFile::Open(f.c_str(), "READ"));
58+
if (!file || file->IsZombie()) {
59+
LOG(error) << "Cannot open " << f;
60+
ok = false;
61+
continue;
62+
}
63+
auto tree = (TTree*)file->Get(treename.c_str());
64+
if (!tree) {
65+
LOG(error) << "No tree named '" << treename << "' found in " << f;
66+
ok = false;
67+
continue;
68+
}
69+
const bool hasTracks = tree->GetBranch(kTrackBranch) != nullptr;
70+
const bool hasHeader = tree->GetBranch(kHeaderBranch) != nullptr;
71+
const bool hasRefs = tree->GetBranch(kTrackRefBranch) != nullptr;
72+
if (!hasTracks || !hasHeader) {
73+
LOG(error) << "File " << f << " is missing the required '" << kTrackBranch << "' and/or '" << kHeaderBranch << "' branch";
74+
ok = false;
75+
continue;
76+
}
77+
if (first) {
78+
hasTrackRefs = hasRefs;
79+
first = false;
80+
} else if (hasTrackRefs != hasRefs) {
81+
LOG(error) << "Inconsistent schema: file " << f << (hasRefs ? " has" : " lacks") << " a '" << kTrackRefBranch << "' branch, unlike previous input files";
82+
ok = false;
83+
}
84+
LOG(info) << " OK " << f << " (" << tree->GetEntries() << " events)";
85+
}
86+
return ok;
87+
}
88+
89+
// Merges the already-validated input files into outfile, giving every event a fresh,
90+
// globally unique, monotonically increasing event ID (starting at startId).
91+
Long64_t mergeFiles(std::vector<std::string> const& files, std::string const& treename,
92+
std::string const& outfile, bool hasTrackRefs, UInt_t startId)
93+
{
94+
TFile fout(outfile.c_str(), "RECREATE");
95+
if (fout.IsZombie()) {
96+
LOG(fatal) << "Cannot create output file " << outfile;
97+
return -1;
98+
}
99+
100+
auto tracks = std::make_unique<std::vector<o2::MCTrack>>();
101+
auto header = std::make_unique<o2::dataformats::MCEventHeader>();
102+
auto trackrefs = std::make_unique<std::vector<o2::TrackReference>>();
103+
auto* tracksPtr = tracks.get();
104+
auto* headerPtr = header.get();
105+
auto* trackrefsPtr = trackrefs.get();
106+
107+
auto outTree = new TTree(treename.c_str(), treename.c_str());
108+
outTree->Branch(kTrackBranch, &tracksPtr);
109+
outTree->Branch(kHeaderBranch, &headerPtr);
110+
if (hasTrackRefs) {
111+
outTree->Branch(kTrackRefBranch, &trackrefsPtr);
112+
}
113+
114+
UInt_t nextEventId = startId;
115+
Long64_t totalEvents = 0;
116+
117+
for (auto const& f : files) {
118+
std::unique_ptr<TFile> fin(TFile::Open(f.c_str(), "READ"));
119+
auto tin = (TTree*)fin->Get(treename.c_str());
120+
121+
tin->SetBranchAddress(kTrackBranch, &tracksPtr);
122+
tin->SetBranchAddress(kHeaderBranch, &headerPtr);
123+
if (hasTrackRefs) {
124+
tin->SetBranchAddress(kTrackRefBranch, &trackrefsPtr);
125+
}
126+
127+
const Long64_t nEntries = tin->GetEntries();
128+
LOG(info) << "Merging " << nEntries << " events from " << f << " (event ID " << nextEventId << ".." << (nextEventId + nEntries - 1) << ")";
129+
for (Long64_t i = 0; i < nEntries; ++i) {
130+
tin->GetEntry(i);
131+
header->SetEventID(nextEventId++);
132+
outTree->Fill();
133+
}
134+
totalEvents += nEntries;
135+
}
136+
137+
fout.cd();
138+
outTree->Write("", TObject::kWriteDelete);
139+
fout.Close();
140+
141+
return totalEvents;
142+
}
143+
} // namespace
144+
145+
int main(int argc, char* argv[])
146+
{
147+
bpo::options_description options("o2-generators-merge-evtpool options");
148+
options.add_options()
149+
("input,i", bpo::value<std::string>()->required(), "comma-separated list of input event-pool ROOT files")
150+
("output,o", bpo::value<std::string>()->required(), "output ROOT file with the merged event pool")
151+
("treename,t", bpo::value<std::string>()->default_value("o2sim"), "name of the tree to merge")
152+
("start-id", bpo::value<UInt_t>()->default_value(1), "event ID assigned to the first merged event")
153+
("help,h", "produce help message");
154+
155+
bpo::variables_map vm;
156+
try {
157+
bpo::store(bpo::parse_command_line(argc, argv, options), vm);
158+
if (vm.count("help")) {
159+
LOG(info) << options;
160+
return 0;
161+
}
162+
bpo::notify(vm);
163+
} catch (const bpo::error& e) {
164+
LOG(fatal) << "Error parsing command-line arguments: " << e.what() << "\n\n"
165+
<< options;
166+
return 1;
167+
}
168+
169+
std::vector<std::string> infiles;
170+
{
171+
std::stringstream ss(vm["input"].as<std::string>());
172+
std::string tok;
173+
while (std::getline(ss, tok, ',')) {
174+
if (!tok.empty()) {
175+
infiles.push_back(tok);
176+
}
177+
}
178+
}
179+
if (infiles.empty()) {
180+
LOG(fatal) << "No input files given";
181+
return 1;
182+
}
183+
184+
const std::string outfile = vm["output"].as<std::string>();
185+
const std::string treename = vm["treename"].as<std::string>();
186+
const UInt_t startId = vm["start-id"].as<UInt_t>();
187+
188+
LOG(info) << "Validating " << infiles.size() << " input file(s) ...";
189+
bool hasTrackRefs = false;
190+
if (!checkFiles(infiles, treename, hasTrackRefs)) {
191+
LOG(fatal) << "Validation failed; not writing any output";
192+
return 1;
193+
}
194+
195+
LOG(info) << "Merging into " << outfile << " ...";
196+
const Long64_t total = mergeFiles(infiles, treename, outfile, hasTrackRefs, startId);
197+
if (total < 0) {
198+
return 1;
199+
}
200+
201+
LOG(info) << "Done: wrote " << total << " events (event ID " << startId << ".." << (startId + total - 1) << ") to " << outfile;
202+
return 0;
203+
}

0 commit comments

Comments
 (0)