mesytec-mnode/src/mana_plugin_root_histogram.cc

#include <TFile.h>
#include <TH1.h>
#include <filesystem>
#include <regex>
#include <spdlog/spdlog.h>
#include <vector>

#include "internal/mana_lib.hpp"
#include "internal/rxi/log.hpp"

using namespace mesytec::mnode;

struct Context
{
    std::unique_ptr<TFile> outputFile;
    std::vector<std::vector<TH1 *>> hitCounts;
    std::vector<std::vector<TH1 *>> rawHistograms;
};

static Context *g_ctx = nullptr;

MANA_DEFINE_PLUGIN_INIT(init)
{
    if (g_ctx)
    {
        log_warn("init() called multiple times. This plugin is a singleton!");
        return g_ctx;
    }
    log_set_level(LOG_INFO);
    log_debug("init");
    return g_ctx = new Context;
}

MANA_DEFINE_PLUGIN_SHUTDOWN(shutdown)
{
    log_debug("shutdown");
    if (context != g_ctx)
    {
        log_warn("shutdown() called with invalid context");
        return;
    }
    delete g_ctx;
    g_ctx = nullptr;
}

struct ObjectPath
{
    std::string path;
    std::string objectName;
    std::vector<std::string> components;

    ObjectPath(const std::string &path, const std::string &sep = "\\.")
        : path(path)
    {
        std::regex re(sep);
        std::sregex_token_iterator it(path.begin(), path.end(), re, -1);
        std::copy(it, std::sregex_token_iterator(), std::back_inserter(components));

        if (!components.empty())
        {
            objectName = components.back();
            components.pop_back();
        }
    }
};

template <typename It> TDirectory *root_mkpath_cd(It begin, It end)
{
    TDirectory *result = nullptr;
    for (auto it = begin; it != end; ++it)
    {
        result = gDirectory->mkdir(it->c_str(), "", true);
        result->cd();
    }

    return result;
}

inline std::vector<std::vector<TH1 *>> make_hitcount_histos(TDirectory *baseDir,
                                                            const nlohmann::json &jRun)
{
    std::vector<std::vector<TH1 *>> result;

    for (const auto &jEvent: jRun["events"])
    {
        std::vector<TH1 *> eventHistos;
        for (const auto &jOutput: jEvent["outputs"])
        {
            auto size = jOutput["size"].get<size_t>();
            ObjectPath path(jOutput["name"].get<std::string>());
            path.components.emplace_back("hit_counts");
            baseDir->cd();
            root_mkpath_cd(path.components.begin(), path.components.end());
            auto histoName = fmt::format("{}_hits", path.path);
            auto histoTitle = fmt::format("Hit Counts {}", path.path);
            auto th1 = new TH1F(histoName.c_str(), histoTitle.c_str(), size, 0.0, size);
            th1->GetXaxis()->SetTitle("Index");
            th1->GetYaxis()->SetTitle("Hits");
            eventHistos.emplace_back(std::move(th1));
        }
        result.emplace_back(std::move(eventHistos));
    }

    return result;
}

inline std::vector<std::vector<TH1 *>> make_raw_histos(TDirectory *baseDir,
                                                       const nlohmann::json &jRun)
{
    std::vector<std::vector<TH1 *>> result;

    for (const auto &jEvent: jRun["events"])
    {
        std::vector<TH1 *> eventHistos;
        for (const auto &jOutput: jEvent["outputs"])
        {
            auto histoCount = jOutput["size"].get<size_t>();
            auto bits = jOutput["bits"].get<unsigned>();
            auto bins = 1u << std::min(bits, 16u);
            auto maxValue = std::pow(2.0, bits);
            ObjectPath path(jOutput["name"].get<std::string>());
            baseDir->cd();
            root_mkpath_cd(path.components.begin(), path.components.end());

            for (size_t i = 0; i < histoCount; ++i)
            {
                auto histoName = fmt::format("{}_{:04d}", path.objectName, i);
                auto histoTitle = fmt::format("{}[{}]", path.path, i);
                auto th1 = new TH1F(histoName.c_str(), histoTitle.c_str(), bins, 0.0, maxValue);
                eventHistos.emplace_back(std::move(th1));
            }
        }
        result.emplace_back(std::move(eventHistos));
    }

    return result;
}

inline std::string histo_info(const std::vector<std::vector<TH1 *>> &histos)
{
    size_t histoCount = 0;
    size_t histoMem = 0;

    for (size_t ei = 0; ei < histos.size(); ++ei)
    {
        for (size_t mi = 0; mi < histos[ei].size(); ++mi)
        {
            auto histo = histos[ei][mi];
            histoCount++;
            histoMem += histo->GetNbinsX() * sizeof(float);
        }
    }

    return fmt::format("histoCount={}, histoMem={:.2f} MiB", histoCount,
                       histoMem / (1024.0 * 1024));
}

MANA_DEFINE_PLUGIN_BEGIN_RUN(begin_run)
{
    log_debug("begin_run: context=%p, descriptor_json=%s", context, descriptor_json);
    auto jRun = nlohmann::json::parse(descriptor_json);
    std::filesystem::path rp(jRun["name"].get<std::string>());
    auto filename = fmt::format("{}_histograms.root", rp.filename().replace_extension().string());
    auto ctx = reinterpret_cast<Context *>(context);
    ctx->hitCounts.clear();
    ctx->outputFile = std::make_unique<TFile>(filename.c_str(), "RECREATE");
    ctx->hitCounts = make_hitcount_histos(ctx->outputFile.get(), jRun);
    ctx->rawHistograms = make_raw_histos(ctx->outputFile.get(), jRun);
    log_info("hitCount histograms: %s", histo_info(ctx->hitCounts).c_str());
    log_info("raw histograms: %s", histo_info(ctx->rawHistograms).c_str());
    log_info("writing histograms into: %s", filename.c_str());
}

MANA_DEFINE_PLUGIN_END_RUN(end_run)
{
    log_debug("end: context=%p, descriptor_json=%s", context, descriptor_json);
    auto ctx = reinterpret_cast<Context *>(context);
    ctx->outputFile->Write();
    *ctx = {};
}

MANA_DEFINE_PLUGIN_EVENT_DATA(process_event)
{
    log_trace("event: ctx=%p, eventIndex=%d, arrayCount=%zu, totalBytes=%zu", context, eventIndex,
              arrayCount, totalBytes);
    auto ctx = reinterpret_cast<Context *>(context);

    auto &hitCountHistograms = ctx->hitCounts.at(eventIndex);
    auto &rawHistograms = ctx->rawHistograms.at(eventIndex);
    size_t rawHistoIndex = 0;

    for (size_t ai = 0; ai < arrayCount; ++ai)
    {
        auto &hitCountHisto = hitCountHistograms.at(ai);
        auto input = mana::get_span<float>(arrays[ai]);

        for (size_t i = 0; i < input.size(); ++i)
        {
            if (!std::isnan(input[i]))
            {
                hitCountHisto->Fill(i);
                rawHistograms.at(rawHistoIndex + i)->Fill(input[i]);
            }
        }

        rawHistoIndex += input.size();
    }
}

MANA_DEFINE_PLUGIN_SYSTEM_EVENT(process_system_event)
{
    log_trace("system_event: ctx=%p, size=%zu", context, size);
}

extern "C" MANA_C_SINK_PLUGIN()
{
    mana_sink_plugin_t plugin;
    plugin.init = init;
    plugin.shutdown = shutdown;
    plugin.begin_run = begin_run;
    plugin.end_run = end_run;
    plugin.process_event = process_event;
    plugin.process_system_event = process_system_event;
    return plugin;
}