florian/mesytec-mnode
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

implement some rntuple thing. no idea what the output file actually contains

Browse source
This commit is contained in:
Florian Lüke 2024-12-30 04:32:40 +01:00
parent 0ef89eac14
commit fe73397582
1 changed files with 43 additions and 30 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

73
src/mana_plugin_root_rntuple_writer.cc
View file

@ -3,6 +3,7 @@
#include <ROOT/RNTupleWriter.hxx>
#include <ROOT/RVec.hxx>
#include <TFile.h>
#include <algorithm>
#include <filesystem>
#include <regex>
#include <spdlog/spdlog.h>
@ -25,9 +26,8 @@ using RNTupleWriter = ROOT::Experimental::RNTupleWriter;
struct Context
{
std::unique_ptr<TFile> outputFile;
std::unique_ptr<RNTupleModel> model;
std::vector<std::vector<std::shared_ptr<ROOT::RVecF>>> eventRVecs;
std::vector<std::shared_ptr<std::vector<float>>> rvecs;
std::unique_ptr<RNTupleWriter> writer;
};
static Context *g_ctx = nullptr;
@ -56,48 +56,48 @@ MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown)
g_ctx = nullptr;
}
inline void make_things(const nlohmann::json &jRun)
{
std::vector<std::vector<std::shared_ptr<ROOT::RVecF>>> result;
for (const auto &jEvent: jRun["events"])
{
auto dir = gDirectory->mkdir(jEvent["name"].get<std::string>().c_str(), "", true);
auto model = RNTupleModel::Create();
std::vector<std::shared_ptr<ROOT::RVecF>> rvecs;
for (const auto &jOutput: jEvent["outputs"])
{
auto rvec = model->MakeField<ROOT::RVecF>(jOutput["name"].get<std::string>());
rvec->resize(jOutput["size"].get<size_t>());
rvecs.emplace_back(std::move(rvec));
}
result.emplace_back(std::move(rvecs));
}
}
MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run)
{
assert(context);
assert(context == g_ctx);
log_debug("begin_run: context=%p, descriptor_json=%s", context, descriptor_json);
auto jRun = nlohmann::json::parse(descriptor_json);
if (jRun["events"].size() > 1)
{
log_error("root-rntuple-writer: currently only supports a single event/trigger");
return;
}
auto ctx = reinterpret_cast<Context *>(context);
ctx->rvecs.clear();
auto jEvent = jRun["events"][0];
auto model = RNTupleModel::Create();
for (const auto &jOutput: jEvent["outputs"])
{
auto name = jOutput["name"].get<std::string>();
std::replace(std::begin(name), std::end(name), '.', '_');
auto rvec = model->MakeField<std::vector<float>>(name);
rvec->resize(jOutput["size"].get<size_t>());
ctx->rvecs.emplace_back(std::move(rvec));
}
std::filesystem::path rp(jRun["name"].get<std::string>());
auto filename =
fmt::format("{}_rntuple_events.root", rp.filename().replace_extension().string());
auto ctx = reinterpret_cast<Context *>(context);
ctx->outputFile = std::make_unique<TFile>(filename.c_str(), "RECREATE");
ctx->outputFile->cd();
make_things(jRun);
ctx->model = ROOT::Experimental::RNTupleModel::Create();
auto field = ctx->model->MakeField<float>("value");
log_info("writing rntuple into: %s", filename.c_str());
ctx->writer =
RNTupleWriter::Recreate(std::move(model), jEvent["name"].get<std::string>(), filename);
}
MANA_DEFINE_PLUGIN_END_RUN(end_run)
{
assert(context);
assert(context == g_ctx);
log_debug("end: context=%p, descriptor_json=%s", context, descriptor_json);
auto ctx = reinterpret_cast<Context *>(context);
ctx->outputFile->Write();
*ctx = {};
}
@ -105,7 +105,20 @@ MANA_DEFINE_PLUGIN_EVENT_DATA(process_event)
{
log_trace("event: ctx=%p, eventIndex=%d, arrayCount=%zu, totalBytes=%zu", context, eventIndex,
arrayCount, totalBytes);
if (eventIndex != 0)
return;
auto ctx = reinterpret_cast<Context *>(context);
for (size_t ai = 0; ai < arrayCount; ++ai)
{
auto &rvec = ctx->rvecs.at(ai);
auto input = mana::get_span<float>(arrays[ai]);
rvec->resize(input.size());
std::copy(std::begin(input), std::end(input), std::begin(*rvec));
}
ctx->writer->Fill();
}
MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)