EIC code displayed by LXR master/​include/​root/​TMVA/​BatchGenerator/​RChunkLoader.hxx	Source navigation Diff markup Identifier search General search
	Version: master test Revert   Change

File indexing completed on 2025-12-15 10:28:51

 // Author: Dante Niewenhuis, VU Amsterdam 07/2023
 // Author: Kristupas Pranckietis, Vilnius University 05/2024
 // Author: Nopphakorn Subsa-Ard, King Mongkut's University of Technology Thonburi (KMUTT) (TH) 08/2024
 // Author: Vincenzo Eduardo Padulano, CERN 10/2024
 // Author: Martin Føll, University of Oslo (UiO) & CERN 05/2025
 
 /*************************************************************************
  * Copyright (C) 1995-2025, Rene Brun and Fons Rademakers.               *
  * All rights reserved.                                                  *
  *                                                                       *
  * For the licensing terms see $ROOTSYS/LICENSE.                         *
  * For the list of contributors see $ROOTSYS/README/CREDITS.             *
  *************************************************************************/
 
 #ifndef TMVA_RCHUNKLOADER
 #define TMVA_RCHUNKLOADER
 
 #include <vector>
 #include <random>
 
 #include "TMVA/RTensor.hxx"
 #include "TMVA/BatchGenerator/RChunkConstructor.hxx"
 #include "ROOT/RDataFrame.hxx"
 #include "ROOT/RDF/Utils.hxx"
 #include "ROOT/RVec.hxx"
 
 #include "ROOT/RLogger.hxx"
 
 namespace TMVA {
 namespace Experimental {
 namespace Internal {
 
 // clang-format off
 /**
 \class ROOT::TMVA::Experimental::Internal::RChunkLoaderFunctor
 \ingroup tmva
 \brief Loading chunks made in RChunkLoader into tensors from data from RDataFrame.
 */
 
 template <typename... ColTypes>
 class RChunkLoaderFunctor {
    // clang-format on   
    std::size_t fOffset{};
    std::size_t fVecSizeIdx{};
    float fVecPadding{};
    std::vector<std::size_t> fMaxVecSizes{};
    TMVA::Experimental::RTensor<float> &fChunkTensor;
 
    std::size_t fNumChunkCols;
 
    int fI;
    int fNumColumns;
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Copy the content of a column into RTensor when the column consits of vectors 
    template <typename T, std::enable_if_t<ROOT::Internal::RDF::IsDataContainer<T>::value, int> = 0>
    void AssignToTensor(const T &vec, int i, int numColumns)
    {
       std::size_t max_vec_size = fMaxVecSizes[fVecSizeIdx++];
       std::size_t vec_size = vec.size();
       if (vec_size < max_vec_size) // Padding vector column to max_vec_size with fVecPadding
       {
          std::copy(vec.begin(), vec.end(), &fChunkTensor.GetData()[fOffset + numColumns * i]);
          std::fill(&fChunkTensor.GetData()[fOffset + numColumns * i + vec_size],
                    &fChunkTensor.GetData()[fOffset + numColumns * i + max_vec_size], fVecPadding);
       } else // Copy only max_vec_size length from vector column
       {
          std::copy(vec.begin(), vec.begin() + max_vec_size, &fChunkTensor.GetData()[fOffset + numColumns * i]);
       }
       fOffset += max_vec_size;
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Copy the content of a column into RTensor when the column consits of single values 
    template <typename T, std::enable_if_t<!ROOT::Internal::RDF::IsDataContainer<T>::value, int> = 0>
    void AssignToTensor(const T &val, int i, int numColumns)
    {
       fChunkTensor.GetData()[fOffset + numColumns * i] = val;
       fOffset++;
       // fChunkTensor.GetData()[numColumns * i] = val;
    }
 
 public:
    RChunkLoaderFunctor(TMVA::Experimental::RTensor<float> &chunkTensor, std::size_t numColumns,
                        const std::vector<std::size_t> &maxVecSizes, float vecPadding, int i)
       : fChunkTensor(chunkTensor), fMaxVecSizes(maxVecSizes), fVecPadding(vecPadding), fI(i), fNumColumns(numColumns)
    {
    }
 
    void operator()(const ColTypes &...cols)
    {
       fVecSizeIdx = 0;
       (AssignToTensor(cols, fI, fNumColumns), ...);
    }
 };
 
 // clang-format off
 /**
 \class ROOT::TMVA::Experimental::Internal::RChunkLoader
 \ingroup tmva
 \brief Building and loading the chunks from the blocks and chunks constructed in RChunkConstructor
 
 In this class the blocks are stiches together to form chunks that are loaded into memory. The blocks used to create each chunk comes from different parts of the dataset. This is achieved by shuffling the blocks before distributing them into chunks. The purpose of this process is to reduce bias during machine learning training by ensuring that the data is well mixed. The dataset is also spit into training and validation sets with the user-defined validation split fraction.
 */
 
 template <typename... Args>
 class RChunkLoader {
 private:
    // clang-format on   
    std::size_t fNumEntries;
    std::size_t fChunkSize;
    std::size_t fBlockSize;
    float fValidationSplit;
 
    std::vector<std::size_t> fVecSizes;
    std::size_t fSumVecSizes;
    std::size_t fVecPadding;
    std::size_t fNumChunkCols;
 
    std::size_t fNumTrainEntries;
    std::size_t fNumValidationEntries;
 
    ROOT::RDF::RNode &f_rdf;
    std::vector<std::string> fCols;
    std::size_t fNumCols;
    std::size_t fSetSeed;
 
    bool fNotFiltered;
    bool fShuffle;
 
    ROOT::RDF::RResultPtr<std::vector<ULong64_t>> fEntries;
 
    std::unique_ptr<RChunkConstructor> fTraining;
    std::unique_ptr<RChunkConstructor> fValidation;
 
 public:
    RChunkLoader(ROOT::RDF::RNode &rdf, std::size_t numEntries,
                 ROOT::RDF::RResultPtr<std::vector<ULong64_t>> rdf_entries, const std::size_t chunkSize,
                 const std::size_t blockSize, const float validationSplit, const std::vector<std::string> &cols,
                 const std::vector<std::size_t> &vecSizes = {}, const float vecPadding = 0.0, bool shuffle = true,
                 const std::size_t setSeed = 0)
       : f_rdf(rdf),
         fNumEntries(numEntries),
         fEntries(rdf_entries),
         fCols(cols),
         fVecSizes(vecSizes),
         fVecPadding(vecPadding),
         fChunkSize(chunkSize),
         fBlockSize(blockSize),
         fValidationSplit(validationSplit),
         fNotFiltered(f_rdf.GetFilterNames().empty()),
         fShuffle(shuffle),
         fSetSeed(setSeed)
    {
       fNumCols = fCols.size();
       fSumVecSizes = std::accumulate(fVecSizes.begin(), fVecSizes.end(), 0);
 
       fNumChunkCols = fNumCols + fSumVecSizes - fVecSizes.size();
 
       // number of training and validation entries after the split
       fNumValidationEntries = static_cast<std::size_t>(fValidationSplit * fNumEntries);
       fNumTrainEntries = fNumEntries - fNumValidationEntries;
 
       fTraining = std::make_unique<RChunkConstructor>(fNumTrainEntries, fChunkSize, fBlockSize);
       fValidation = std::make_unique<RChunkConstructor>(fNumValidationEntries, fChunkSize, fBlockSize);
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Distribute the blocks into training and validation datasets
    void SplitDataset()
    {
       std::random_device rd;
       std::mt19937 g;
 
       if (fSetSeed == 0) {
          g.seed(rd());
       } else {
          g.seed(fSetSeed);
       }
 
       std::vector<Long_t> BlockSizes = {};
 
       // fill the training and validation block sizes
       for (size_t i = 0; i < fTraining->NumberOfDifferentBlocks.size(); i++) {
          BlockSizes.insert(BlockSizes.end(), fTraining->NumberOfDifferentBlocks[i], fTraining->SizeOfBlocks[i]);
       }
 
       for (size_t i = 0; i < fValidation->NumberOfDifferentBlocks.size(); i++) {
          BlockSizes.insert(BlockSizes.end(), fValidation->NumberOfDifferentBlocks[i], fValidation->SizeOfBlocks[i]);
       }
 
       // make an identity permutation map
       std::vector<Long_t> indices(BlockSizes.size());
 
       for (int i = 0; i < indices.size(); ++i) {
          indices[i] = i;
       }
 
       // shuffle the identity permutation to create a new permutation
       if (fShuffle) {
          std::shuffle(indices.begin(), indices.end(), g);
       }
 
       // use the permuation to shuffle the vector of block sizes  
       std::vector<Long_t> PermutedBlockSizes(BlockSizes.size());
       for (int i = 0; i < BlockSizes.size(); ++i) {
          PermutedBlockSizes[i] = BlockSizes[indices[i]];
       }
 
       // create a vector for storing the boundaries of the blocks
       std::vector<Long_t> BlockBoundaries(BlockSizes.size());
 
       // get the boundaries of the blocks with the partial sum of the block sizes
       // insert 0 at the beginning for the lower boundary of the first block
       std::partial_sum(PermutedBlockSizes.begin(), PermutedBlockSizes.end(), BlockBoundaries.begin());
       BlockBoundaries.insert(BlockBoundaries.begin(), 0);
 
       // distribute the neighbouring block boudaries into pairs to get the intevals for the blocks
       std::vector<std::pair<Long_t, Long_t>> BlockIntervals;
       for (size_t i = 0; i < BlockBoundaries.size() - 1; ++i) {
          BlockIntervals.emplace_back(BlockBoundaries[i], BlockBoundaries[i + 1]);
       }
 
       // use the inverse of the permutation above to order the block intervals in the same order as
       // the original vector of block sizes
       std::vector<std::pair<Long_t, Long_t>> UnpermutedBlockIntervals(BlockIntervals.size());
       for (int i = 0; i < BlockIntervals.size(); ++i) {
          UnpermutedBlockIntervals[indices[i]] = BlockIntervals[i];
       }
 
       // distribute the block intervals between training and validation
       fTraining->BlockIntervals.insert(fTraining->BlockIntervals.begin(), UnpermutedBlockIntervals.begin(),
                                        UnpermutedBlockIntervals.begin() + fTraining->NumberOfBlocks);
       fValidation->BlockIntervals.insert(fValidation->BlockIntervals.begin(),
                                          UnpermutedBlockIntervals.begin() + fTraining->NumberOfBlocks,
                                          UnpermutedBlockIntervals.end());
 
       // distribute the different block intervals types for training and validation
       fTraining->DistributeBlockIntervals();
       fValidation->DistributeBlockIntervals();
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Create training chunks consisiting of block intervals of different types 
    void CreateTrainingChunksIntervals()
    {
 
       std::random_device rd;
       std::mt19937 g;
 
       if (fSetSeed == 0) {
          g.seed(rd());
       } else {
          g.seed(fSetSeed);
       }
 
       // shuffle the block intervals within each type of block
       if (fShuffle) {
          std::shuffle(fTraining->FullBlockIntervalsInFullChunks.begin(),
                       fTraining->FullBlockIntervalsInFullChunks.end(), g);
          std::shuffle(fTraining->LeftoverBlockIntervalsInFullChunks.begin(),
                       fTraining->LeftoverBlockIntervalsInFullChunks.end(), g);
          std::shuffle(fTraining->FullBlockIntervalsInLeftoverChunks.begin(),
                       fTraining->FullBlockIntervalsInLeftoverChunks.end(), g);
          std::shuffle(fTraining->LeftoverBlockIntervalsInLeftoverChunks.begin(),
                       fTraining->LeftoverBlockIntervalsInLeftoverChunks.end(), g);
       }
 
       // reset the chunk intervals and sizes before each epoch
       fTraining->ChunksIntervals = {};
       fTraining->ChunksSizes = {};
 
       // create the chunks each consisiting of block intervals
       fTraining->CreateChunksIntervals();
 
       if (fShuffle) {
          std::shuffle(fTraining->ChunksIntervals.begin(), fTraining->ChunksIntervals.end(), g);
       }
 
       fTraining->SizeOfChunks();
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Create training chunks consisiting of block intervals of different types 
    void CreateValidationChunksIntervals()
    {
       std::random_device rd;
       std::mt19937 g;
 
       if (fSetSeed == 0) {
          g.seed(rd());
       } else {
          g.seed(fSetSeed);
       }
 
       if (fShuffle) {
          std::shuffle(fValidation->FullBlockIntervalsInFullChunks.begin(),
                       fValidation->FullBlockIntervalsInFullChunks.end(), g);
          std::shuffle(fValidation->LeftoverBlockIntervalsInFullChunks.begin(),
                       fValidation->LeftoverBlockIntervalsInFullChunks.end(), g);
          std::shuffle(fValidation->FullBlockIntervalsInLeftoverChunks.begin(),
                       fValidation->FullBlockIntervalsInLeftoverChunks.end(), g);
          std::shuffle(fValidation->LeftoverBlockIntervalsInLeftoverChunks.begin(),
                       fValidation->LeftoverBlockIntervalsInLeftoverChunks.end(), g);
       }
 
       fValidation->ChunksIntervals = {};
       fValidation->ChunksSizes = {};
 
       fValidation->CreateChunksIntervals();
 
       if (fShuffle) {
          std::shuffle(fValidation->ChunksIntervals.begin(), fValidation->ChunksIntervals.end(), g);
       }
 
       fValidation->SizeOfChunks();
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Load the nth chunk from the training dataset into a tensor
    /// \param[in] TrainChunkTensor RTensor for the training chunk
    /// \param[in] chunk Index of the chunk in the dataset
    void LoadTrainingChunk(TMVA::Experimental::RTensor<float> &TrainChunkTensor, std::size_t chunk)
    {
 
       std::random_device rd;
       std::mt19937 g;
 
       if (fSetSeed == 0) {
          g.seed(rd());
       } else {
          g.seed(fSetSeed);
       }
 
       std::size_t chunkSize = fTraining->ChunksSizes[chunk];
 
       if (chunk < fTraining->Chunks) {
          TMVA::Experimental::RTensor<float> Tensor({chunkSize, fNumChunkCols});
          TrainChunkTensor = TrainChunkTensor.Resize({{chunkSize, fNumChunkCols}});
 
          // make an identity permutation map        
          std::vector<int> indices(chunkSize);
          std::iota(indices.begin(), indices.end(), 0);
 
          // shuffle the identity permutation to create a new permutation         
          if (fShuffle) {
             std::shuffle(indices.begin(), indices.end(), g);
          }
 
          // fill a chunk by looping over the blocks in a chunk (see RChunkConstructor)
          std::size_t chunkEntry = 0;
          std::vector<std::pair<Long_t, Long_t>> BlocksInChunk = fTraining->ChunksIntervals[chunk];
 
          std::sort(BlocksInChunk.begin(), BlocksInChunk.end(),
                    [](const std::pair<Long_t, Long_t>& a, const std::pair<Long_t, Long_t>& b) {
                       return a.first < b.first;
                    });
          
          for (std::size_t i = 0; i < BlocksInChunk.size(); i++) {
 
             // Use the block start and end entry to load into the chunk if the dataframe is not filtered
             if (fNotFiltered) {
                RChunkLoaderFunctor<Args...> func(Tensor, fNumChunkCols, fVecSizes, fVecPadding, chunkEntry);
                ROOT::Internal::RDF::ChangeBeginAndEndEntries(f_rdf, BlocksInChunk[i].first, BlocksInChunk[i].second);
 
                f_rdf.Foreach(func, fCols);
                chunkEntry += BlocksInChunk[i].second - BlocksInChunk[i].first;
             }
 
             // use the entry column of the dataframe as a map to load the entries that passed the filters
             else {
                std::size_t blockSize = BlocksInChunk[i].second - BlocksInChunk[i].first;
                for (std::size_t j = 0; j < blockSize; j++) {
                   RChunkLoaderFunctor<Args...> func(Tensor, fNumChunkCols, fVecSizes, fVecPadding, chunkEntry);
                   ROOT::Internal::RDF::ChangeBeginAndEndEntries(f_rdf, (*fEntries)[BlocksInChunk[i].first + j],
                                                                 (*fEntries)[BlocksInChunk[i].first + j + 1]);
                   f_rdf.Foreach(func, fCols);
                   chunkEntry++;
                }
             }
          }
 
          // shuffle data in RTensor with the permutation map defined above
          for (std::size_t i = 0; i < chunkSize; i++) {
             std::copy(Tensor.GetData() + indices[i] * fNumChunkCols,
                       Tensor.GetData() + (indices[i] + 1) * fNumChunkCols,
                       TrainChunkTensor.GetData() + i * fNumChunkCols);
          }
       }
    }
 
    //////////////////////////////////////////////////////////////////////////
    /// \brief Load the nth chunk from the validation dataset into a tensor
    /// \param[in] ValidationChunkTensor RTensor for the validation chunk
    /// \param[in] chunk Index of the chunk in the dataset
    void LoadValidationChunk(TMVA::Experimental::RTensor<float> &ValidationChunkTensor, std::size_t chunk)
    {
 
       std::random_device rd;
       std::mt19937 g;
 
       if (fSetSeed == 0) {
          g.seed(rd());
       } else {
          g.seed(fSetSeed);
       }
 
       std::size_t chunkSize = fValidation->ChunksSizes[chunk];
 
       if (chunk < fValidation->Chunks) {
          TMVA::Experimental::RTensor<float> Tensor({chunkSize, fNumChunkCols});
          ValidationChunkTensor = ValidationChunkTensor.Resize({{chunkSize, fNumChunkCols}});
 
          // make an identity permutation map        
          std::vector<int> indices(chunkSize);
          std::iota(indices.begin(), indices.end(), 0);
 
          // shuffle the identity permutation to create a new permutation
          if (fShuffle) {
             std::shuffle(indices.begin(), indices.end(), g);
          }
 
          std::size_t chunkEntry = 0;
          std::vector<std::pair<Long_t, Long_t>> BlocksInChunk = fValidation->ChunksIntervals[chunk];
 
          std::sort(BlocksInChunk.begin(), BlocksInChunk.end(),
                    [](const std::pair<Long_t, Long_t>& a, const std::pair<Long_t, Long_t>& b) {
                       return a.first < b.first;
                    });
          
          for (std::size_t i = 0; i < BlocksInChunk.size(); i++) {
 
             // use the block start and end entry to load into the chunk if the dataframe is not filtered
             if (fNotFiltered) {
                RChunkLoaderFunctor<Args...> func(Tensor, fNumChunkCols, fVecSizes, fVecPadding, chunkEntry);
                ROOT::Internal::RDF::ChangeBeginAndEndEntries(f_rdf, BlocksInChunk[i].first, BlocksInChunk[i].second);
                f_rdf.Foreach(func, fCols);
                chunkEntry += BlocksInChunk[i].second - BlocksInChunk[i].first;
             }
 
             // use the entry column of the dataframe as a map to load the entries that passed the filters
             else {
                std::size_t blockSize = BlocksInChunk[i].second - BlocksInChunk[i].first;
                for (std::size_t j = 0; j < blockSize; j++) {
                   RChunkLoaderFunctor<Args...> func(Tensor, fNumChunkCols, fVecSizes, fVecPadding, chunkEntry);
                   ROOT::Internal::RDF::ChangeBeginAndEndEntries(f_rdf, (*fEntries)[BlocksInChunk[i].first + j],
                                                                 (*fEntries)[BlocksInChunk[i].first + j + 1]);
 
                   f_rdf.Foreach(func, fCols);
                   chunkEntry++;
                }
             }
          }
 
          // shuffle data in RTensor with the permutation map defined above         
          for (std::size_t i = 0; i < chunkSize; i++) {
             std::copy(Tensor.GetData() + indices[i] * fNumChunkCols,
                       Tensor.GetData() + (indices[i] + 1) * fNumChunkCols,
                       ValidationChunkTensor.GetData() + i * fNumChunkCols);
          }
       }
    }
 
    std::vector<std::size_t> GetTrainingChunkSizes() { return fTraining->ChunksSizes; }
    std::vector<std::size_t> GetValidationChunkSizes() { return fValidation->ChunksSizes; }
 
    std::size_t GetNumTrainingEntries() { return fNumTrainEntries; }
    std::size_t GetNumValidationEntries() { return fNumValidationEntries; }
 
    void CheckIfUnique(TMVA::Experimental::RTensor<float> &Tensor)
    {
       auto tensorSize = Tensor.GetSize();
       TMVA::Experimental::RTensor<float> SqueezeTensor = Tensor.Reshape({1, tensorSize}).Squeeze();
 
       std::list<int> allEntries;
       for (int i = 0; i < tensorSize; i++) {
          allEntries.push_back(SqueezeTensor(0, i));
       }
       allEntries.sort();
       allEntries.unique();
       if (allEntries.size() == tensorSize) {
          std::cout << "Tensor consists of only unique elements" << std::endl;
       }
    };
 
    void CheckIfOverlap(TMVA::Experimental::RTensor<float> &Tensor1, TMVA::Experimental::RTensor<float> &Tensor2)
    {
       auto tensorSize1 = Tensor1.GetSize();
       TMVA::Experimental::RTensor<float> SqueezeTensor1 = Tensor1.Reshape({1, tensorSize1}).Squeeze();
 
       std::list<int> allEntries1;
       for (int i = 0; i < tensorSize1; i++) {
          allEntries1.push_back(SqueezeTensor1(0, i));
       }
 
       auto tensorSize2 = Tensor2.GetSize();
       TMVA::Experimental::RTensor<float> SqueezeTensor2 = Tensor2.Reshape({1, tensorSize2}).Squeeze();
 
       std::list<int> allEntries2;
       for (int i = 0; i < tensorSize2; i++) {
          allEntries2.push_back(SqueezeTensor2(0, i));
       }
 
       std::set<int> result;
 
       // Call the set_intersection(), which computes the
       // intersection of set1 and set2 and
       // inserts the result into the 'result' set
       std::set<int> set1(allEntries1.begin(), allEntries1.end());
       std::set<int> set2(allEntries2.begin(), allEntries2.end());
       std::set_intersection(set1.begin(), set1.end(), set2.begin(), set2.end(), inserter(result, result.begin()));
       // std::list<int> result = intersection(allEntries1, allEntries2);
 
       if (result.size() == 0) {
          std::cout << "No overlap between the tensors" << std::endl;
       } else {
          std::cout << "Intersection between tensors: ";
          for (int num : result) {
             std::cout << num << " ";
          }
          std::cout << std::endl;
       }
    };
 
    std::size_t GetNumTrainingChunks() { return fTraining->Chunks; }
 
    std::size_t GetNumValidationChunks() { return fValidation->Chunks; }
 };
 
 } // namespace Internal
 } // namespace Experimental
 } // namespace TMVA
 #endif // TMVA_RCHUNKLOADER

This page was automatically generated by the 2.3.7 LXR engine. The LXR team