root/TMVA/RTensorUtils.hxx

0001 #ifndef TMVA_RTENSOR_UTILS
0002 #define TMVA_RTENSOR_UTILS
0003
0004 #include <vector>
0005 #include <string>
0006
0007 #include "TMVA/RTensor.hxx"
0008 #include "ROOT/RDataFrame.hxx"
0009 #include "ROOT/RDF/RInterface.hxx"
0010
0011 namespace TMVA {
0012 namespace Experimental {
0013
0014 /// \brief Convert the content of an RDataFrame to an RTensor
0015 /// \param[in] dataframe RDataFrame node
0016 /// \param[in] columns Vector of column names
0017 /// \param[in] layout Memory layout
0018 /// \return RTensor with content from selected columns
0019 template <typename T, typename U>
0020 RTensor<T>
0021 AsTensor(U &dataframe, std::vector<std::string> columns = {}, MemoryLayout layout = MemoryLayout::RowMajor)
0022 {
0023    // If no columns are specified, get all columns from dataframe
0024    if (columns.size() == 0) {
0025       columns = dataframe.GetColumnNames();
0026    }
0027
0028    // Book actions to read-out columns of dataframe in vectors
0029    using ResultPtr = ROOT::RDF::RResultPtr<std::vector<T>>;
0030    std::vector<ResultPtr> resultPtrs;
0031    for (auto &col : columns) {
0032       resultPtrs.emplace_back(dataframe.template Take<T>(col));
0033    }
0034
0035    // Copy data to tensor based on requested memory layout
0036    const auto numCols = resultPtrs.size();
0037    const auto numEntries = resultPtrs[0]->size();
0038    RTensor<T> x({numEntries, numCols}, layout);
0039    const auto data = x.GetData();
0040    if (layout == MemoryLayout::RowMajor) {
0041       for (std::size_t i = 0; i < numEntries; i++) {
0042          const auto entry = data + numCols * i;
0043          for (std::size_t j = 0; j < numCols; j++) {
0044             entry[j] = resultPtrs[j]->at(i);
0045          }
0046       }
0047    } else if (layout == MemoryLayout::ColumnMajor) {
0048       for (std::size_t i = 0; i < numCols; i++) {
0049          // TODO: Replace by RVec<T>::insert as soon as available.
0050          std::memcpy(data + numEntries * i, &resultPtrs[i]->at(0), numEntries * sizeof(T));
0051       }
0052    } else {
0053       throw std::runtime_error("Memory layout is not known.");
0054    }
0055
0056    // Remove dimensions of 1
0057    x.Squeeze();
0058
0059    return x;
0060 }
0061
0062 } // namespace TMVA::Experimental
0063 } // namespace TMVA
0064
0065 #endif