BinnedStorage.h

Go to the documentation of this file.

// -*- C++ -*-
//
// This file is part of YODA -- Yet more Objects for Data Analysis
// Copyright (C) 2008-2025 The YODA collaboration (see AUTHORS for details)
//
#ifndef YODA_BinnedStorage_H
#define YODA_BinnedStorage_H
 
#include "YODA/BinnedAxis.h"
#include "YODA/Binning.h"
#include "YODA/Bin.h"
 
#include <memory>
#include <vector>
 
namespace YODA {
 
  template <class VecT>
  class BinsVecWrapper {
  public:
    class myIt;
 
    using BinType = std::decay_t<decltype(*std::declval<VecT>().begin())>;
 
    BinsVecWrapper() = delete;
    BinsVecWrapper(VecT& bins, const std::vector<size_t>& hiddenBins)
      : _bins(bins), _hiddenBins(hiddenBins) {}
 
    BinsVecWrapper(BinsVecWrapper&& other)
      : _bins(std::move(other._bins)), _hiddenBins(std::move(other._hiddenBins)) {}
 
    myIt begin() const  { return myIt(_bins, _hiddenBins); }
    myIt end()   const  { return myIt(_bins); }
 
    size_t size() const { return _bins.size() - _hiddenBins.size(); }
 
    const BinType& operator[](size_t index) const { return _bins[index]; }
 
    BinType& operator[](size_t index) { return _bins[index]; }
 
  private:
 
    VecT& _bins;
    std::vector<size_t> _hiddenBins;
  };
 
  template<typename VecT>
  class BinsVecWrapper<VecT>::myIt {
  private:
    using IterT = std::conditional_t<std::is_const<VecT>::value,
                                        typename VecT::const_iterator,
                                        typename VecT::iterator>;
  public:
 
    myIt(VecT& bins,
         const std::vector<size_t>& hiddenBins)
        : _ptr(bins.begin()),
          hiddenCurrPtr(hiddenBins.begin()),
          hiddenEndPtr(hiddenBins.end()),
          binsEndPtr(bins.end()) {
 
        if (hiddenCurrPtr != hiddenEndPtr && 0 == *hiddenCurrPtr){
 
          ++hiddenCurrPtr; 
          ++(*this);       
 
        }
      }
 
    myIt(VecT& bins)
        :  _ptr(bins.end()),
           //hiddenCurrPtr(std::vector<size_t>::const_iterator()),
           //hiddenEndPtr(std::vector<size_t>::const_iterator()),
           binsEndPtr(bins.end()) {}
 
    myIt operator++() noexcept {
      ++_ptr;
      ++currBinIdx;
 
      while (_ptr != binsEndPtr &&
            hiddenCurrPtr != hiddenEndPtr &&
            currBinIdx == *hiddenCurrPtr) {
 
        ++currBinIdx;
        ++hiddenCurrPtr;
        ++_ptr;
      }
      return *this;
    }
 
    bool operator!=(const myIt& other) const noexcept {
      return _ptr != other._ptr;
    }
 
    typename std::iterator_traits<IterT>::reference operator*() noexcept { return *_ptr; }
 
  private:
    IterT _ptr;
    std::vector<size_t>::const_iterator hiddenCurrPtr;
    std::vector<size_t>::const_iterator hiddenEndPtr;
    IterT binsEndPtr;
    size_t currBinIdx = 0;
  };
 
 
  template <typename BinContentT, typename... AxisT>
  class BinnedStorage {
  protected:
 
    using BinningT = Binning<std::decay_t<decltype(std::declval<Axis<AxisT>>())>...>;
    using BinT = Bin<sizeof...(AxisT), BinContentT, BinningT>;
    using BinsVecT = std::vector<BinT>;
    using BaseT = BinnedStorage<BinContentT, AxisT...>;
 
  public:
 
    using BinningType = BinningT;
    using BinType = BinT;
    using BinDimension = std::integral_constant<size_t, sizeof...(AxisT)>;
 
 
    // @{
 
    BinnedStorage() : _binning(std::vector<AxisT>{}...) {
      fillBins();
    }
 
    BinnedStorage(const BinningT& binning) : _binning(binning) {
      fillBins();
    }
 
    BinnedStorage(BinningT&& binning) : _binning(std::move(binning)) {
      fillBins();
    }
 
    BinnedStorage(const std::vector<AxisT>&... edges) : _binning(edges...) {
      fillBins();
    }
 
    BinnedStorage(std::vector<AxisT>&&... edges) : _binning(std::move(edges)...) {
      fillBins();
    }
 
    BinnedStorage(std::initializer_list<AxisT>&&... edges) : _binning(std::vector<AxisT>{edges}...) {
      fillBins();
    }
 
    BinnedStorage(const Axis<AxisT>&... axes) : _binning(axes...) {
      fillBins();
    }
 
    BinnedStorage(Axis<AxisT>&&... axes) : _binning(std::move(axes)...) {
      fillBins();
    }
 
    BinnedStorage(const BinnedStorage& other) : _binning(other._binning) {
      fillBins(other._bins);
    }
 
    BinnedStorage(BinnedStorage&& other) : _binning(std::move(other._binning)) {
      fillBins(std::move(other._bins));
    }
 
    // @}
 
    // @{
 
    size_t dim() const noexcept {
      return sizeof...(AxisT) + 1;
    }
 
    BinT& bin(size_t idx) noexcept {
      return _bins.at(idx);
    }
 
    const BinT& bin(size_t idx) const noexcept {
      return _bins.at(idx);
    }
 
    BinT& bin(const std::array<size_t, sizeof...(AxisT)>& idxLocal) noexcept {
      return bin( _binning.localToGlobalIndex(idxLocal) );
    }
 
    const BinT& bin(const std::array<size_t, sizeof...(AxisT)>& idxLocal) const noexcept {
      return bin( _binning.localToGlobalIndex(idxLocal) );
    }
 
    BinT& binAt(typename BinningT::EdgeTypesTuple&& coords) noexcept {
      const size_t binIdx = _binning.globalIndexAt(coords);
      return bin(binIdx);
    }
 
    const BinT& binAt(typename BinningT::EdgeTypesTuple&& coords) const noexcept {
      const size_t binIdx = _binning.globalIndexAt(coords);
      return bin(binIdx);
    }
 
    void set(typename BinningT::EdgeTypesTuple&& coords, BinContentT&& content) noexcept {
      const size_t binIdx = _binning.globalIndexAt(coords);
      _bins[binIdx] = std::move(content);
    }
 
    void set(typename BinningT::EdgeTypesTuple&& coords, const BinContentT& content) noexcept {
      const size_t binIdx = _binning.globalIndexAt(coords);
      _bins[binIdx] = content;
    }
 
    void set(const size_t binIdx, BinContentT&& content) noexcept {
      _bins[binIdx] = std::move(content);
    }
 
    void set(const size_t binIdx, const BinContentT& content) noexcept {
      _bins[binIdx] = content;
    }
 
    std::vector<size_t> calcIndicesToSkip(const bool includeOverflows, const bool includeMaskedBins) const noexcept {
 
      // if there are no bins, exit early
      if (!_binning.numBins(!includeOverflows, !includeMaskedBins))  return {};
 
      std::vector<size_t> indicesToSkip;
 
      // define a lambda
      auto appendIndicesVec = [&indicesToSkip](std::vector<size_t>&& indicesVec) {
      indicesToSkip.insert(std::end(indicesToSkip),
                           std::make_move_iterator(std::begin(indicesVec)),
                           std::make_move_iterator(std::end(indicesVec)));
      };
 
      // only calculate the masked indices when
      // the masked bins are to be skipped over
      if(!includeOverflows) {
        appendIndicesVec(_binning.calcOverflowBinsIndices());
      }
 
      if (!includeMaskedBins) {
        appendIndicesVec(_binning.maskedBins());
      }
 
      // sort and remove duplicates
      std::sort(indicesToSkip.begin(), indicesToSkip.end());
      indicesToSkip.erase( std::unique(indicesToSkip.begin(), indicesToSkip.end()),
                           indicesToSkip.end() );
 
      return indicesToSkip;
    }
 
 
    BinsVecWrapper<BinsVecT> bins(const bool includeOverflows = false,
                                  const bool includeMaskedBins = false) noexcept {
      return BinsVecWrapper<BinsVecT>(_bins,
                calcIndicesToSkip(includeOverflows, includeMaskedBins));
    }
 
    const BinsVecWrapper<const BinsVecT> bins(const bool includeOverflows = false,
                                              const bool includeMaskedBins = false) const noexcept {
      return BinsVecWrapper<const BinsVecT>(_bins,
                calcIndicesToSkip(includeOverflows, includeMaskedBins));
    }
 
    // @}
 
    // @{
 
    const BinningT& binning() const noexcept {
      return _binning;
    }
 
    size_t binDim() const noexcept {
      return _binning.dim();
    }
 
    size_t numBins(const bool includeOverflows = false, const bool includeMaskedBins = false) const noexcept {
      return _binning.numBins(includeOverflows, includeMaskedBins);
    }
 
    size_t numBinsAt(const size_t axisN, const bool includeOverflows = false) const noexcept {
      size_t nOverflows = includeOverflows? 0 : _binning.countOverflowBins(axisN);
      return _binning.numBinsAt(axisN) - nOverflows;
    }
 
    void reset() noexcept { clearBins(); }
 
    void clearBins() noexcept {
      _bins.clear();
      fillBins();
    }
 
    void maskBins(const std::vector<size_t>& indicesToMask, const bool status = true) noexcept {
      _binning.maskBins(indicesToMask, status);
    }
 
    void maskBin(const size_t indexToMask, const bool status = true) noexcept {
      _binning.maskBin(indexToMask, status);
    }
 
    void maskSlice(const size_t dim, const size_t idx, const bool status = true) {
      _binning.maskSlice(dim, idx, status);
    }
 
    void maskBinAt(typename BinningT::EdgeTypesTuple&& coords, const bool status = true) noexcept {
      _binning.maskBinAt(coords, status);
    }
 
    bool isMasked(const size_t binIndex) const noexcept {
      return _binning.isMasked(binIndex);
    }
 
    std::vector<size_t> maskedBins() const noexcept {
      return _binning.maskedBins();
    }
 
    bool isVisible(const size_t binIndex) const noexcept {
      return _binning.isVisible(binIndex);
    }
 
    template <size_t... AxisNs, class RetT = void>
    auto mergeBins(
      std::decay_t<decltype(AxisNs, std::declval<std::pair<size_t, size_t>>())>... mergeRanges)
      noexcept -> std::enable_if_t<MetaUtils::is_detected_v<MetaUtils::operatorTraits::addition_assignment_t, BinContentT>, RetT>
    {
      auto mergeStorageBins =
        [&binning = BaseT::_binning, &binStorage = BaseT::_bins](auto I, const auto& mergeRangePair){
          assert(mergeRangePair.first < mergeRangePair.second);
 
          auto append = [&binStorage](const auto& pivotBinsIndices, const auto& binsIndicesToMerge){
            assert(pivotBinsIndices.size() == binsIndicesToMerge.size());
            /*for (const auto& k : binsIndicesToMerge) { std::cout << " " << k; }
            std::cout << std::endl;
            for (const auto& k : pivotBinsIndices) { std::cout << " " << k; }
            std::cout << std::endl;*/
 
            // first merge the bins based on old set of indices
            // unless the bins are masked
            for (size_t i = 0; i < pivotBinsIndices.size(); ++i) {
              auto& pivotBin = binStorage[pivotBinsIndices[i]];
              pivotBin += std::move(binStorage[binsIndicesToMerge[i]]);
            }
            // then erase the bins (which will change the set of indices)
            binStorage.erase(
              std::remove_if(binStorage.begin(), binStorage.end(), [&](const auto& b) {
                return std::find(binsIndicesToMerge.begin(), binsIndicesToMerge.end(), b.index()) != binsIndicesToMerge.end();
              }), binStorage.end());
          };
 
          ssize_t nBinRowsToBeMerged = mergeRangePair.second - mergeRangePair.first;
 
          size_t currBinRowIdx = mergeRangePair.first;
          size_t nextBinRowIdx = mergeRangePair.first + 1;
          //std::cout << nBinRowsToBeMerged << " " << currBinRowIdx << " " << nextBinRowIdx << std::endl;
 
          while (nBinRowsToBeMerged--) {
            append(binning.sliceIndices(I, mergeRangePair.first), binning.sliceIndices(I, nextBinRowIdx));
            binning.template mergeBins<I>({currBinRowIdx, nextBinRowIdx});
          }
        };
 
      ((void)mergeStorageBins(std::integral_constant<size_t, AxisNs>(), mergeRanges), ...);
 
    }
 
 
    template<size_t axisN, template<typename...> typename BinnedT, typename Func,
             typename = std::enable_if_t< (axisN < sizeof...(AxisT) && sizeof...(AxisT)>=2) >>
    auto mkBinnedSlices(Func&& how2add, const bool includeOverflows=false) const {
 
      size_t vecN = BaseT::numBinsAt(axisN, includeOverflows);
      auto binnedSlice = _mkBinnedT<BinnedT>(_binning.template _getAxesExcept<axisN>());
      std::vector<decltype(binnedSlice)> rtn(vecN, binnedSlice);
      for (size_t i = 0; i < vecN; ++i) {
 
        auto mkSlice = [&oldBins = _bins, &how2add, &binnedSlice = rtn[i]](const auto& indicesToCopy) {
          assert(binnedSlice.numBins(true) == indicesToCopy.size());
 
          // for any given pivot, add the content
          // from the old slice to the new slice
          for (size_t i = 0; i < binnedSlice.numBins(true); ++i) {
            auto& pivotBin = binnedSlice.bin(i);
            auto& binToCopy = oldBins[indicesToCopy[i]];
            how2add(pivotBin, binToCopy);
          }
        };
 
        // get bin slice for any given bin i along the axis that is to be
        // sliced, then make the estimates for the new binning
        mkSlice(_binning.sliceIndices(axisN, i + !includeOverflows));
 
      }
      return rtn;
    }
 
 
    BinnedStorage& operator = (const BinnedStorage& other) noexcept {
      if (this != &other) {
        _binning = other._binning;
        fillBins(other._bins);
      }
      return *this;
    }
 
    BinnedStorage& operator = (BinnedStorage&& other) noexcept {
      if (this != &other) {
        _binning = std::move(other._binning);
        fillBins(std::move(other._bins));
      }
      return *this;
    }
 
    bool operator == (const BinnedStorage& other) const noexcept {
      return _binning.isCompatible(other._binning);
    }
 
    bool operator != (const BinnedStorage& other) const noexcept {
        return ! operator == (other);
    }
 
    std::vector<size_t> _global2local(size_t idx) const noexcept {
      const auto& indices = _binning.globalToLocalIndices(idx);
      return std::vector<size_t>(indices.begin(), indices.end());
    }
 
    size_t _local2global(const std::vector<size_t>& indices) const {
      assert(indices.size() == sizeof...(AxisT));
      std::array<size_t, sizeof...(AxisT)> arr;
      std::copy_n(std::make_move_iterator(indices.begin()), sizeof...(AxisT), arr.begin());
      return _binning.localToGlobalIndex(arr);
    }
 
 
 
    protected:
 
      // @{
 
      void fillBins() noexcept {
        _bins.reserve(_binning.numBins());
 
        for (size_t i = 0; i < _binning.numBins(); ++i) {
          _bins.emplace_back(i, _binning);
        }
      }
 
      void fillBins(const BinsVecT& bins) noexcept {
        _bins.clear();
        _bins.reserve(_binning.numBins());
        for (const auto& b : bins) {
          _bins.emplace_back(b, _binning);
        }
      }
 
      void fillBins(BinsVecT&& bins) noexcept {
        _bins.clear();
        _bins.reserve(_binning.numBins());
        for (auto&& b : bins) {
          _bins.emplace_back(std::move(b), _binning);
        }
      }
 
      template<template<typename...> typename BinnedT, class... newAxes, size_t... As>
      auto _mkBinnedT_aux(const std::tuple<newAxes...>& t, std::index_sequence<As...>) const {
        return BinnedT<typename newAxes::EdgeT...>(std::get<As>(t)...);
      }
 
      template <template<typename...> typename BinnedT, class... newAxes>
      auto _mkBinnedT(const std::tuple<newAxes...>& t) const {
        return _mkBinnedT_aux<BinnedT>(t, std::make_index_sequence<sizeof...(newAxes)>{});
      }
 
      // @}
 
      BinsVecT _bins;
 
      BinningT _binning;
 
  };
 
} // namespace YODA
 
#endif