Flattener.h

Go to the documentation of this file.

// SPDX-License-Identifier: GPLv3-or-later
//
// SPDX-FileCopyrightText: Patrick L.S. Connor <patrick.connor@cern.ch>
 
#pragma once
 
#include <Flow.h>
#include <GenericObject.h>
#include <GenericPlugin.h>
#include <MetaInfo.h>
#include <SelfAwareBranch.h>
#include <Variation.h>
 
#include <list>
#include <memory>
#include <ranges>
#include <set>
#include <string>
#include <type_traits>
 
namespace pt = boost::property_tree;
namespace DT = Darwin::Tools;
 
namespace Darwin::Physics {
 
class BranchFlattener {
protected:
    Variation v;
    BranchFlattener (const Variation& v) : v(v) { }
public:
    virtual ~BranchFlattener () = default;
    virtual void Fill () = 0;
    virtual void Clear () = 0;
};
 
template<typename T> class ScaledObjectBranchFlattener : public BranchFlattener {
 
    std::vector<T> * branchIn; 
    std::list<std::unique_ptr<DT::SelfAwareBranch<const T&>>>
                            branchesOut; 
    std::unique_ptr<DT::SelfAwareBranch<const std::vector<T>&>> multiplicity;
 
public:
 
    template<typename Ret> void Add (Tools::Flow& flow,
                                     const std::string& name,
                                     const std::function<Ret(const T&)>& f)
    {
        branchesOut.emplace_back(std::make_unique<DT::SelfAwareVectorBranch<Ret, const T&>>(flow, name, f));
    }
 
    ScaledObjectBranchFlattener (Tools::Flow& flow, const Variation& v,
                             const std::string& nameIn, 
                             const std::string& nameOut 
                             )
        : BranchFlattener(v), branchIn(nullptr)
    {
        using namespace std;
 
        branchIn = flow.GetBranchReadOnly<vector<T>>(nameIn, DT::facultative);
        if (!branchIn) return;
 
        Add<float>(flow, nameOut + "_pt"  , [this](const T& el) { return el.CorrP4(this->v).pt  (); });
        Add<float>(flow, nameOut + "_eta" , [this](const T& el) { return el.CorrP4(this->v).eta (); });
        Add<float>(flow, nameOut + "_phi" , [this](const T& el) { return el.CorrP4(this->v).phi (); });
        Add<float>(flow, nameOut + "_mass", [this](const T& el) { return el.CorrP4(this->v).mass(); });
 
        auto getMult = [](const vector<T>& branchIn) { return branchIn.size(); };
        multiplicity = make_unique<DT::SelfAwarePlainBranch<int, const vector<T>&>>(flow, "n" + nameOut, getMult);
    }
 
    void Fill () override
    {
        if (!branchIn) return;
 
        for (const auto& br: branchesOut)
            for (const auto& el: *branchIn)
                br->Fill(el);
 
        multiplicity->Fill(*branchIn);
    }
 
    void Clear () override
    {
        if (!branchIn) return;
 
        for (const auto& br: branchesOut)
            br->Clear();
 
        multiplicity->Clear();
    }
 
    virtual ~ScaledObjectBranchFlattener () = default;
};
 
template<typename T> class EventObjectBranchFlattener : public BranchFlattener {
 
    T * branchIn; 
    std::list<std::unique_ptr<DT::SelfAwareBranch<const T&>>>
                            branchesOut; 
 
public:
 
    template<typename Ret> void Add (Tools::Flow& flow,
                                     const std::string& name,
                                     const std::function<Ret(const T&)>& f)
    {
        branchesOut.emplace_back(std::make_unique<DT::SelfAwarePlainBranch<Ret, const T&>>(flow, name, f));
    }
 
    EventObjectBranchFlattener (Tools::Flow& flow, const Variation& v,
                                const std::string& nameIn, 
                                const std::string& nameOut 
                                )
        : BranchFlattener(v), branchIn(nullptr)
    {
        using namespace std;
 
        branchIn = flow.GetBranchReadOnly<T>(nameIn, DT::facultative);
        if (!branchIn) return;
 
        Add<double>(flow, nameOut, [this](const T& obj) { return obj.Weight(this->v); });
    }
 
    void Fill () override
    {
        for (const auto& br: branchesOut)
            br->Fill(*branchIn);
    }
 
    void Clear () override
    {
        for (const auto& br: branchesOut)
            br->Clear();
    }
 
    virtual ~EventObjectBranchFlattener () = default;
};
 
class TreeFlattener : public DT::GenericPlugin {
 
protected:
 
    std::list<std::unique_ptr<BranchFlattener>> branch_flatteners;
 
    virtual int GetParamsVeto () noexcept override { return DT::syst; }
 
public:
 
    TreeFlattener () = default;
    virtual ~TreeFlattener () = default;
 
    virtual void ExecuteAtBegin (const std::set<Variation>& = {nominal}) override
    {
        for (const auto& br: branch_flatteners)
            br->Clear();
    }
 
    virtual void ExecuteAtEnd (const std::set<Variation>& = {nominal}) override
    {
        for (const auto& br: branch_flatteners)
            br->Fill();
    }
};
 
} // end of Darwin::Physics namespace