ROOT/RDF/PyROOTHelpers.hxx

0001 // Author: Danilo Piparo, Enrico Guiraud, Stefan Wunsch CERN  04/2018
0002
0003 /*************************************************************************
0004  * Copyright (C) 1995-2018, Rene Brun and Fons Rademakers.               *
0005  * All rights reserved.                                                  *
0006  *                                                                       *
0007  * For the licensing terms see $ROOTSYS/LICENSE.                         *
0008  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
0009  *************************************************************************/
0010
0011 #ifndef ROOT_PyROOTHelpers
0012 #define ROOT_PyROOTHelpers
0013
0014 #include "ROOT/RDataFrame.hxx"
0015
0016 #include <vector>
0017 #include <string>
0018 #include <utility>
0019
0020 namespace ROOT {
0021 namespace Internal {
0022 namespace RDF {
0023
0024 template <typename dtype>
0025 ULong64_t GetVectorAddress(std::vector<dtype> &p)
0026 {
0027    return reinterpret_cast<ULong64_t>(&p);
0028 }
0029
0030 inline ULong64_t GetAddress(std::vector<std::string> &p)
0031 {
0032    return reinterpret_cast<ULong64_t>(&p);
0033 }
0034 inline ULong64_t GetAddress(TTree &p)
0035 {
0036    return reinterpret_cast<ULong64_t>(&p);
0037 }
0038
0039 template <typename BufType, typename... ColTypes, std::size_t... Idx>
0040 void TTreeAsFlatMatrix(std::index_sequence<Idx...>, TTree &tree, std::vector<BufType> &matrix,
0041                        std::vector<std::string> &columns)
0042 {
0043    auto buffer = matrix.data();
0044
0045    auto fillMatrix = [buffer](ColTypes... cols, ULong64_t entry) {
0046       int expander[] = {(buffer[entry * sizeof...(Idx) + Idx] = cols, 0)...};
0047       (void)expander;
0048    };
0049
0050    auto columnsWithEntry = columns;
0051    columnsWithEntry.emplace_back("tdfentry_");
0052
0053    ROOT::RDataFrame dataframe(tree, columns);
0054    dataframe.Foreach(fillMatrix, columnsWithEntry);
0055 }
0056
0057 template <typename BufType, typename... ColTypes>
0058 void TTreeAsFlatMatrixHelper(TTree &tree, std::vector<BufType> &matrix, std::vector<std::string> &columns)
0059 {
0060    TTreeAsFlatMatrix<BufType, ColTypes...>(std::index_sequence_for<ColTypes...>(), tree, matrix, columns);
0061 }
0062
0063 // RDataFrame.AsNumpy helpers
0064
0065 // NOTE: This is a workaround for the missing Take action in the PyROOT interface
0066 template <typename T>
0067 ROOT::RDF::RResultPtr<std::vector<T>> RDataFrameTake(ROOT::RDF::RNode df, std::string_view column)
0068 {
0069    return df.Take<T>(column);
0070 }
0071
0072 } // namespace RDF
0073 } // namespace Internal
0074 } // namespace ROOT
0075
0076 #endif