NodeBase.h

Go to the documentation of this file.

#ifndef RDFAnalysis_NodeBase_H
#define RDFAnalysis_NodeBase_H

// Package includes
#include "RDFAnalysis/IBranchNamer.h"
#include "RDFAnalysis/Helpers.h"
#include "RDFAnalysis/SysResultPtr.h"
#include "RDFAnalysis/SysVar.h"
#include "RDFAnalysis/WeightStrategy.h"

// ROOT includes
#include "ROOT/RDataFrame.hxx"
#include <TObject.h>

// STL includes
#include <string>
#include <map>
#include <memory>
#include <array>

namespace RDFAnalysis {
  using RNode = ROOT::RDF::RNode;
  using ColumnNames_t = ROOT::RDataFrame::ColumnNames_t;

  class NodeBase
  {
    public:
      template <typename F>
          enable_ifn_string_t<F, NodeBase*> Define(
            const std::string& name,
            F f,
            const ColumnNames_t& columns = {});

      NodeBase* Define(
          const std::string& name,
          const std::string& expression);

      NodeBase* Define(
          const std::string& name,
          const std::string& expression,
          const ColumnNames_t& columns);

      template <
        std::size_t N,
        typename F,
        typename Ret_t = typename ROOT::TTraits::CallableTraits<F>::ret_type>
        std::enable_if_t<N==std::tuple_size<Ret_t>::value, NodeBase*> Define(
            const std::array<std::string, N>& names,
            F f,
            const ColumnNames_t& columns);

      const std::string& getWeight() const { return m_weight; }

      template <typename T>
        SysResultPtr<T> Fill(
            const T& model,
            const ColumnNames_t& columns,
            const std::string& weight = "",
            WeightStrategy strategy = WeightStrategy::Default);

      template <typename AccFun, typename MergeFun,
               typename ArgTypes=typename ROOT::TTraits::CallableTraits<AccFun>::arg_types,
               typename U=ROOT::TTraits::TakeFirstParameter_t<ArgTypes>>
        SysResultPtr<U> Aggregate(
            AccFun aggregator,
            MergeFun merger,
            const std::string& columnName,
            const U& aggIdentity)
        {
          return ActResult<ROOT::RDF::RResultPtr<U>, AccFun, MergeFun, std::string_view, const U&>(
              &RNode::Aggregate,
              ColumnNames_t{columnName},
              aggregator, merger, SysVarBranch(columnName), aggIdentity);
        }

      template <typename AccFun, typename MergeFun,
               typename ArgTypes=typename ROOT::TTraits::CallableTraits<AccFun>::arg_types,
               typename U=ROOT::TTraits::TakeFirstParameter_t<ArgTypes>>
        SysResultPtr<U> Aggregate(
            AccFun aggregator,
            MergeFun merger,
            const std::string& columnName)
        {
          return ActResult<ROOT::RDF::RResultPtr<U>, AccFun, MergeFun, std::string_view>(
              &RNode::Aggregate,
              ColumnNames_t{columnName},
              aggregator, merger, SysVarBranch(columnName) );
        }

      SysResultPtr<ULong64_t> Count()
      {
        return ActResult<ROOT::RDF::RResultPtr<ULong64_t>>(
            &RNode::Count,
            ColumnNames_t{});
      }

      template <typename... TrArgs, typename T, typename... Args>
        std::map<std::string, T> Act(
            std::function<T(RNode&, TrArgs...)> f,
            const ColumnNames_t& columns,
            Args&&... args);

      template <typename F, typename... Args,
               typename T=typename ROOT::TTraits::CallableTraits<F>::ret_type>
        std::enable_if_t<!is_std_function<F>::value, std::map<std::string, T>> Act(
            F&& f,
            const ColumnNames_t& columns,
            Args&&... args)
        {
          return Act(
              std::function<T(RNode&, typename sysvar_traits<Args&&>::value_type...)>(f),
              columns,
              std::forward<Args>(args)...);
        }

      template <typename T, typename... TrArgs, typename... Args>
        std::map<std::string, T> Act(
            T (RNode::*f)(TrArgs...),
            const ColumnNames_t& columns,
            Args&&... args);

      template <typename... TrArgs, typename T, typename... Args,
               typename U=typename T::Value_t>
        SysResultPtr<U> ActResult(
            std::function<T(RNode&, TrArgs...)> f,
            const ColumnNames_t& columns,
            Args&&... args)
        {
          return SysResultPtr<U>(
              namer().nominalName(),
              Act(f, columns, std::forward<Args>(args)...) );
        }

      template <typename F, typename... Args,
               typename T=typename ROOT::TTraits::CallableTraits<F>::ret_type,
               typename U=typename T::Value_t>
        std::enable_if_t<!is_std_function<F>::value, SysResultPtr<U>> ActResult(
            F&& f,
            const ColumnNames_t& columns,
            Args&&... args)
        {
          return ActResult(
              std::function<T(RNode&, typename sysvar_traits<Args&&>::value_type...)>(f),
              columns,
              std::forward<Args>(args)...);
        }

      template <typename T, typename... TrArgs, typename... Args,
                typename U=typename T::Value_t>
        SysResultPtr<U> ActResult(
            T (RNode::*f)(TrArgs...),
            const ColumnNames_t& columns,
            Args&&... args)
        {
          return SysResultPtr<U>(
              namer().nominalName(),
              Act(f, columns, std::forward<Args>(args)...) );
        }

      const std::string& name() const { return m_name; }

      const std::string& cutflowName() const { return m_cutflowName; }

      bool isAnonymous() const { return m_name.empty(); }

      bool isMC() const { return m_isMC; }

      const std::map<std::string, RNode>& rnodes() const { return m_rnodes; }
      std::map<std::string, RNode>& rnodes() { return m_rnodes; }

      const IBranchNamer& namer() const { return *m_namer; }

      auto objects() { return as_range(m_objects); }
      auto objects() const { return as_range(m_objects); }

      virtual bool isRoot() const = 0;

    protected:
      struct NamerInitialiser {
        NamerInitialiser() {} //> no-op
        NamerInitialiser(
            IBranchNamer& namer,
            const std::map<std::string, ROOT::RDF::RNode>& rnodes) {
          namer.readBranchList(rnodes);
        }
      };

      template <std::size_t I, std::size_t N, typename... Elements>
        std::enable_if_t<I!=0, void> unwindDefine(
            const std::array<std::string, N>& names,
            const std::string& fullName,
            const std::tuple<Elements...>*);

      template <std::size_t I, std::size_t N, typename... Elements>
        std::enable_if_t<I == 0, void> unwindDefine(
            const std::array<std::string, N>& names,
            const std::string& fullName,
            const std::tuple<Elements...>*);


      template <typename F>
        enable_ifn_string_t<F, std::map<std::string, RNode>> makeChildRNodes(
            F f,
            const ColumnNames_t& columns = {},
            const std::string& cutflowName = "");

      std::map<std::string, RNode> makeChildRNodes(
          const std::string& expression,
          const std::string& cutflowName = "");

      std::map<std::string, RNode> makeChildRNodes(
          const std::string& expression,
          const ColumnNames_t& columns,
          const std::string& cutflowName = "");

      NodeBase(
          const RNode& rnode,
          std::unique_ptr<IBranchNamer>&& namer,
          bool isMC,
          const std::string& name = "ROOT",
          const std::string& cutflowName = "Number of events",
          const std::string& weight = "",
          WeightStrategy strategy = WeightStrategy::Default);

      template <typename W>
        NodeBase(
            const RNode& rnode,
            std::unique_ptr<IBranchNamer>&& namer,
            bool isMC,
            const std::string& name,
            const std::string& cutflowName,
            W w,
            const ColumnNames_t& columns,
            WeightStrategy strategy = WeightStrategy::Default);

      NodeBase(
          NodeBase& parent,
          std::map<std::string, RNode>&& rnodes,
          const std::string& name,
          const std::string& cutflowName,
          const std::string& weight,
          WeightStrategy strategy = WeightStrategy::Default);
      template <typename W>
        NodeBase(
            NodeBase& parent,
            std::map<std::string, RNode>&& rnodes,
            const std::string& name,
            const std::string& cutflowName,
            W w,
            const ColumnNames_t& columns,
            WeightStrategy strategy = WeightStrategy::Default);

      template <typename F>
        enable_ifn_string_t<F, std::string> setWeight(
            F f,
            const ColumnNames_t& columns,
            NodeBase* parent,
            WeightStrategy strategy);

      std::string setWeight(
          const std::string& expression,
          NodeBase* parent,
          WeightStrategy strategy);

      std::string nameWeight();

      std::map<std::string, RNode> m_rnodes;      

      std::unique_ptr<IBranchNamer> m_namer;

      NamerInitialiser m_namerInit;

      bool m_isMC;

      std::string m_name;

      std::string m_cutflowName;

      RNode* m_rootRNode = nullptr;

      std::string m_weight;

      std::vector<SysResultPtr<TObject>> m_objects;
  }; //> end class NodeBase
} //> end namespace RDFAnalysis
#include "RDFAnalysis/NodeBase.icc"
#endif //> !RDFAnalysis_NodeBase_H