Flow.h

Go to the documentation of this file.

// SPDX-License-Identifier: GPLv3-or-later
//
// SPDX-FileCopyrightText: Patrick L.S. Connor <patrick.connor@desy.de>
// SPDX-FileCopyrightText: Louis Moureaux <louis.moureaux@cern.ch>
 
#pragma once
 
#include "exceptions.h"
#include "FriendUtils.h"
#include "Looper.h"
 
#include <TChain.h>
#include <TFile.h>
#include <TH1.h>
#include <TTree.h>
 
#include <any>
#include <array>
#include <filesystem>
#include <map>
#include <memory>
#include <ranges>
#include <source_location>
#include <string>
#include <vector>
 
namespace Darwin::Tools {
 
enum BranchMode {
    mandatory  , 
    facultative  
}; 
 
using enum BranchMode;
 
using Slice = std::pair<int, int>; 
 
inline std::ostream& operator<< (std::ostream& Stream, const Darwin::Tools::Slice& slice)
{
    return Stream << slice.second << '/' << slice.first;
}
 
class Flow {
 
    int steering; 
 
    std::vector<std::filesystem::path> inputs; 
 
    std::unique_ptr<ChainSlice> tIn; 
    std::shared_ptr<TFile> fOut; 
    std::unique_ptr<TTree> tOut; 
 
    std::map<std::string, std::any> branches; 
 
    template<typename T> 
    std::shared_ptr<T*> GetBranch (const std::string& name) const 
    try {
        using namespace std;
        if (steering & verbose)
            cout << "Flow: retrieving an existing branch" << endl;
        auto& branch = branches.at(name);
        return any_cast<shared_ptr<T*>>(branch);
    }
    catch (const std::bad_any_cast& e) {
        BOOST_THROW_EXCEPTION(e);
    }
 
    template<typename T> T * NoBranch (const std::string& name, BranchMode mode)
    {
        if (mode == facultative)
            return nullptr;
 
        namespace DE = Darwin::Exceptions;
        std::string what = name + " branch could not be found";
        BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tIn) );
    }
 
public:
    inline std::string DumpActiveBranches () const
    {
        auto keys = branches | std::views::keys;
        return std::accumulate(keys.begin(), keys.end(), std::string(),
                [](const std::string& str, const std::string& branch_name) {
                    return str + ' ' + branch_name;
                });
    }
 
    Flow (int = none, 
          const std::vector<std::filesystem::path>& = {} 
         );
 
    ~Flow ();
 
    ChainSlice * GetInputTree
            (const Slice = {1, 0}, 
             const std::string& = "events"); 
 
    template<typename THX = TH1, size_t N>
    std::array<std::unique_ptr<THX>, N> GetInputHists
            (const std::array<std::string, N>& names = {} 
            )
    {
        using namespace std;
        namespace fs = filesystem;
 
        if (inputs.size() == 0)
            BOOST_THROW_EXCEPTION( runtime_error("Empty input list") );
 
        array<unique_ptr<THX>, N> sums;
        for (const fs::path& input: inputs) {
            auto fIn = make_unique<TFile>(input.c_str(), "READ");
            for (size_t i = 0; i < N; ++i) {
                const string& name = names[i];
                unique_ptr<THX> h(fIn->Get<THX>(name.c_str()));
                if (!h) {
                    namespace DE = Darwin::Exceptions;
                    BOOST_THROW_EXCEPTION(
                            DE::BadInput(Form("`%s` cannot be found in (one of) the "
                                              " file(s).", name.c_str()), fIn));
                }
                if (sums[i])
                    sums[i]->Add(h.get());
                else {
                    sums[i] = std::move(h);
                    sums[i]->SetDirectory(nullptr);
                }
            }
        }
        return sums;
    }
 
    template<typename THX = TH1, typename... Args>
    auto GetInputHists (const Args... args)
    {
        constexpr const size_t N = sizeof...(args);
        std::array<std::string, N> names {{ args... }};
        return GetInputHists<THX, N>(names);
    }
 
    template<typename THX = TH1>
    std::unique_ptr<THX> GetInputHist
            (const std::string& name) 
    {
        auto hists = GetInputHists<THX,1>({name});
        THX * hist = hists.front().release();
        return std::unique_ptr<THX>(hist);
    }
 
    TTree * GetOutputTree (std::shared_ptr<TFile> = {}, 
            const std::source_location = std::source_location::current()); 
 
    TTree * GetOutputTree (const std::filesystem::path&, 
            const std::source_location = std::source_location::current()); 
 
    inline void SetOutputFile (std::shared_ptr<TFile> fOut) { this->fOut = fOut; }
 
    inline TFile * GetOutputFile () { return fOut.get(); }
 
    std::pair<TFile *, TTree *> GetOutput (const std::filesystem::path&, 
            const std::source_location = std::source_location::current()); 
 
    template<typename T> 
    T * GetBranchReadOnly (
                const std::string& name, 
                BranchMode mode = mandatory) 
    {
        using namespace std;
 
        if (!tIn)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetInputTree()` should "
                            "be called before declaring a read-only branch") );
 
        namespace DE = Darwin::Exceptions;
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else {
            t = make_shared<T*>();
            if (steering & verbose)
                cout << "Flow: loading branch `" << name << "`" << endl;
 
            if (tIn->GetBranch(name.c_str()) == nullptr)
                return NoBranch<T>(name, mode);
 
            int err = tIn->SetBranchAddress(name.c_str(), t.get());
            if (steering & verbose)
                cout << "Flow: `TTree::SetBranchAddress()` returned " << to_string(err)
                     << " (check `TTree::ESetBranchAddressStatus` for the meaning)."
                     << endl;
            if (err < 0) {
                string what = "`"s + name + "` branch could not be set. "s;
                if (mode == facultative) {
                    if (steering & verbose)
                        cout << orange << "Flow: " << what << def << endl;
                    return nullptr;
                }
                BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tIn) );
            }
            branches.insert({name, t});
        }
        return *t;
    }
 
    template<typename T> 
    T * GetBranchWriteOnly (const std::string& name) 
    {
        using namespace std;
 
        if (!tOut)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetOutputTree()` should "
                                                    "be called before") );
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else
            t = make_shared<T*>();
 
        if (steering & verbose)
            cout << "Flow: setting up new branch for `" << name << "`" << endl;
        if (tOut->Branch(name.c_str(), t.get()) == nullptr) {
            namespace DE = Darwin::Exceptions;
            string what = name + " branch could not be set up";
            BOOST_THROW_EXCEPTION( DE::BadInput(what.c_str(), *tOut) );
        }
 
        if (!branches.contains(name))
            branches.insert({name, t});
        return *t;
    }
 
    template<typename T> 
    T * GetBranchReadWrite (const std::string& name, 
                            BranchMode mode = mandatory) 
    {
        using namespace std;
 
        if (!tIn)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetInputTree()` should "
                            "be called before declaring a read-write branch") );
 
        if (!tOut)
            BOOST_THROW_EXCEPTION( invalid_argument("`GetOutputTree()` should "
                            "be called before declaring a read-write branch") );
 
        shared_ptr<T*> t;
 
        if (branches.contains(name))
            t = GetBranch<T>(name);
        else {
            if (GetBranchReadOnly<T>(name, mode) == nullptr)
                return NoBranch<T>(name, mode);
            if (steering & Friend)
                GetBranchWriteOnly<T>(name);
            t = any_cast<shared_ptr<T*>>(branches[name]);
        }
        return *t;
    }
};
 
} // namespace Darwin::Tools
 
using Darwin::Tools::operator<<;