Cutflow.hh

#ifndef RIVET_Cutflow_HH
#define RIVET_Cutflow_HH
 
#include "Rivet/Tools/RivetYODA.hh"
#include "Rivet/Tools/Utils.hh"
 
#include <iomanip>
#include <iostream>
#include <vector>
#include <sstream>
#include <string>
 
namespace Rivet {
 
  namespace {
    template<typename T>
    std::vector<T> operator+(std::vector<T>&& res, const std::vector<T>& vadd) {
      res.insert(std::end(res), std::begin(vadd), std::end(vadd));
      return std::move(res);
    }
  }
 
  class Cutflow : public YODA::BinnedHisto<std::string> {
  public:
 
    using BaseT = YODA::BinnedHisto<std::string>;
    using BinningT = BaseT::BinningT;
    using FillType = BaseT::FillType;
    using BinType = BaseT::BinT;
    using Ptr = std::shared_ptr<Cutflow>;
    using AnalysisObject::operator =;
 
    Cutflow() : BaseT(), icurr(0) {}
 
    Cutflow(const BinningT& binning) : BaseT(binning), icurr(0) { }
 
    Cutflow(BinningT&& binning) : BaseT(std::move(binning)), icurr(0) { }
 
    Cutflow(const std::vector<std::string>& edges, const std::string& path = "")
      : BaseT(std::vector<std::string>{""} + edges, path), icurr(0) { }
 
    virtual int fill(FillType&& coords, const double weight = 1.0, const double fraction = 1.0) {
      int pos = BaseT::fill(std::move(coords), weight, fraction);
      icurr = (size_t)pos; // cannot be negative for discrete AOs
      return pos;
    }
 
    //
    virtual int fill(std::string val, const double weight = 1.0, const double fraction = 1.0) {
      int pos = BaseT::fill({val}, weight, fraction);
      icurr = (size_t)pos; // cannot be negative for discrete AOs
      return pos;
    }
 
    virtual int fillinit(const double weight=1.0, const double fraction=1.0) {
      return fill(""s, weight, fraction);
    }
 
    virtual int fillnext(const bool cutresult, const double weight = 1.0, const double fraction = 1.0) {
      if (!cutresult)  return 0;
      const std::string edge = BaseT::bin(++icurr).xEdge();
      return fill(edge, weight, fraction);
    }
 
    virtual int fillnext(const double weight = 1.0, const double fraction = 1.0) {
      return fillnext(true, weight, fraction);
    }
 
    virtual int fillnext(const std::vector<bool>& cutresults, const double weight = 1.0, const double fraction = 1.0) {
      if (icurr+cutresults.size() > BaseT::numBins()+1) {
        throw RangeError("Number of filled cut results needs to match the Cutflow construction (in cutflow '"+BaseT::path()+"')");
      }
      for (size_t i = 0; i < cutresults.size(); ++i) {
        if (!cutresults[i])  return 0;
        fillnext(weight, fraction); // also incremenrs icurr
      }
      return (int)icurr;
    }
 
    size_t currentIndex() const {  return icurr; }
 
    void scale(const double factor) {
      BaseT::scaleW(factor);
    }
 
    void normalizeStep(const std::string& edge, const double norm) {
      const auto& b = BaseT::binAt(edge);
      if (b.sumW() == 0) {
        MSG_WARNING("Failed to scale Cutflow " << BaseT::path() << " as bin " << edge << " is empty");
        return;
      }
      scale(norm/b.sumW());
    }
 
    void normalizeFirst(const double norm) {
      normalizeStep(""s, norm);
    }
 
    string str() const {
      using namespace std;
      stringstream ss;
      ss << fixed << std::setprecision(1) << BaseT::bin(1).sumW();
      const size_t weight0len = ss.str().length();
      ss << fixed << std::setprecision(1) << BaseT::bin(1).effNumEntries();
      const size_t count0len = ss.str().length();
      ss.str("");
      ss << BaseT::path() << " cut-flow:\n";
      size_t maxnamelen = 0;
      for (const string& edge : BaseT::xEdges()) {
        maxnamelen = std::max(edge.length(), maxnamelen);
      }
      ss << setw(maxnamelen+5) << "" << "   "
         << setw(weight0len) << right << "Weight" << "   "
         << setw(count0len) << right << "Count" << "    "
         << setw(6) << right << "A_cumu" << "    "
         << setw(6) << right << "A_incr";
      //for (size_t i = 0; i <= ncuts; ++i) {
      const double wtot = BaseT::bin(1).sumW();
      double wlast = wtot;
      for (const auto& bin : BaseT::bins()) {
        const size_t i = bin.index();
        const int pcttot = (wtot == 0) ? -1 : round(100*bin.sumW()/wtot);
        const int pctinc = (i == 1 || wlast == 0) ? -1 : round(100*bin.sumW()/wlast);
        wlast = bin.sumW();
        stringstream ss2;
        ss2 << fixed << setprecision(1) << bin.sumW();
        const string weightstr = ss2.str(); ss2.str("");
        ss2 << fixed << setprecision(1) << bin.effNumEntries();
        const string countstr = ss2.str(); ss2.str("");
        ss2 << fixed << setprecision(3) << pcttot << "%";
        const string pcttotstr = ss2.str(); ss2.str("");
        ss2 << fixed << setprecision(3) << pctinc << "%";
        const string pctincstr = ss2.str();
        ss << "\n"
           << setw(maxnamelen+5) << left << (i == 1 ? "" : "Pass "+BaseT::bin(i).xEdge()) << "   "
           << setw(weight0len) << right << weightstr << "    "
           << setw(count0len) << right << countstr << "    "
           << setw(6) << right << (pcttot < 0 ? "- " : pcttotstr) << "    "
           << setw(6) << right << (pctinc < 0 ? "- " : pctincstr);
      }
      return ss.str();
    }
 
    void print(std::ostream& os) const {
      os << str() << std::flush;
    }
 
  protected:
 
    Log& getLog() const {
      return Rivet::Log::getLog("Rivet.Cutflow");
    }
 
    size_t icurr;
 
  };
 
 
  template <>
  class FillCollector<Cutflow> : public Cutflow {
  public:
 
    using YAO = Cutflow;
    using Ptr = shared_ptr<FillCollector<YAO>>;
    using YAO::operator =;
 
    FillCollector() : YAO() { }
 
    FillCollector(typename YAO::Ptr yao) : YAO(yao->binning()) {
      YAO::setPath(yao->path());
    }
 
    int fill(typename YAO::FillType&& fillCoords,
             const double weight=1.0, const double fraction=1.0) {
      (void)fraction; // suppress unused variable warning
      YAO::icurr = YAO::_binning.globalIndexAt(fillCoords);
      _fills.insert(_fills.end(), { std::move(fillCoords), weight } );
      return (int)YAO::icurr;
    }
 
    int fill(const std::string val, const double weight = 1.0, const double fraction = 1.0) {
      return fill(typename YAO::FillType{val}, weight, fraction);
    }
 
    int fillnext(const bool cutresult, const double weight = 1.0, const double fraction = 1.0) {
      if (!cutresult)  return 0;
      return fill(YAO::bin(++YAO::icurr).xEdge(), weight, fraction);
    }
 
    int fillnext(const double weight = 1.0, const double fraction = 1.0) {
      return fillnext(true, weight, fraction);
    }
 
    int fillnext(const std::vector<bool>& cutresults, const double weight = 1.0, const double fraction = 1.0) {
      if (YAO::icurr+cutresults.size() > YAO::numBins()+1) {
        throw RangeError("Number of filled cut results needs to match the Cutflow construction (in cutflow '"+YAO::path()+"')");
      }
      for (size_t i = 0; i < cutresults.size(); ++i) {
        if (!cutresults[i])  return 0;
        fillnext(weight, fraction);
      }
      return (int)YAO::icurr;
    }
 
    int fillinit(const double weight=1.0, const double fraction=1.0) {
      return fill(""s, weight, fraction);
    }
 
    void reset() noexcept { _fills.clear(); }
 
    const Fills<YAO>& fills() const { return _fills; }
 
  private:
 
    Fills<YAO> _fills;
 
  };
 
 
  using CutflowPtr = MultiplexPtr<Multiplexer<Cutflow>>;
 
  inline std::ostream& operator << (std::ostream& os, const Cutflow& cf) {
    return os << cf.str();
  }
 
  inline std::ostream& operator << (std::ostream& os, const CutflowPtr& cf) {
    return os << cf->str();
  }
 
 
 
  class Cutflows : public YODA::FillableStorage<2, CutflowPtr, std::string>,
                     public YODA::Fillable {
  public:
 
    using BaseT = YODA::FillableStorage<2, CutflowPtr, std::string>;
    using BinContentT = CutflowPtr;
    using BinT = YODA::Bin<1, BinContentT, BaseT::BinningT>;
    using FillDim = std::integral_constant<size_t, 2>;
 
 
    Cutflows() : BaseT(groupAdapter<FillDim{}, BinContentT, std::string>) { }
 
    Cutflows(const std::vector<std::string>& edges)
         : BaseT(YODA::Axis<std::string>(edges), groupAdapter<FillDim{}, BinContentT, std::string>) { }
 
    Cutflows(std::initializer_list<std::string>&& edges)
         : BaseT(YODA::Axis<std::string>(std::move(edges)), groupAdapter<FillDim{}, BinContentT, std::string>) { }
 
 
 
    int fill(const std::string& grpCoord, const std::string& edge, const double weight = 1.0, const double fraction = 1.0) {
      auto& cfl = BaseT::binAt({grpCoord});
      if (!cfl.raw())  return -1; // nullptr if bin not booked
      return cfl->fill({edge}, weight, fraction);
    }
 
    int groupfillinit(const double weight=1.0, const double fraction=1.0) {
      bool pass = true;
      for (auto& cfl : BaseT::bins()) {
        if (!cfl.get())  continue;
        pass &= cfl->fillinit(weight, fraction);
      }
      return (int)pass;
    }
 
    int groupfillnext(const bool cutresult, const double weight = 1.0, const double fraction = 1.0) {
      if (!cutresult)  return 0;
      bool pass = true;
      for (auto& cfl : BaseT::bins()) {
        if (!cfl.get())  continue;
        pass &= cfl->fillnext(cutresult, weight, fraction);
      }
      return (int)pass;
    }
 
    int groupfillnext(const double weight = 1.0, const double fraction = 1.0) {
      return groupfillnext(true, weight, fraction);
    }
 
    int groupfillnext(const std::vector<bool>& cutresults, const double weight = 1.0, const double fraction = 1.0) {
      bool pass = true;
      for (auto& cfl : BaseT::bins()) {
        if (!cfl.get())  continue;
        pass &= cfl->fillnext(cutresults, weight, fraction);
      }
      return (int)pass;
    }
 
    int fillnext(const std::string& grpCoord, const double weight=1.0, const double fraction=1.0) {
      return BaseT::binAt(grpCoord)->fillnext(weight, fraction);
    }
 
    int fillnext(const std::string& grpCoord, const bool cutresult, const double weight=1.0, const double fraction=1.0) {
      return BaseT::binAt(grpCoord)->fillnext(cutresult, weight, fraction);
    }
 
    int fillnext(const std::string& grpCoord, const std::vector<bool>& cutresults,
                 const double weight = 1.0, const double fraction = 1.0) {
      return BaseT::binAt(grpCoord)->fillnext(cutresults, weight, fraction);
    }
 
 
 
    void reset() noexcept { BaseT::reset(); }
 
 
 
    size_t fillDim() const noexcept { return FillDim{}; }
 
    size_t dim() const noexcept { return fillDim()+1; }
 
    double numEntries(const bool includeOverflows=true) const noexcept {
      double n = 0;
      for (const auto& cfl : BaseT::bins(includeOverflows)) {
        if (!cfl.get())  continue;
        n += cfl->numEntries(includeOverflows);
      }
      return n;
    }
 
    double effNumEntries(const bool includeOverflows=true) const noexcept {
      double n = 0;
      for (const auto& cfl : BaseT::bins(includeOverflows)) {
        if (!cfl.get())  continue;
        n += cfl->effNumEntries(includeOverflows);
      }
      return n;
    }
 
    double sumW(const bool includeOverflows=true) const noexcept {
      double sumw = 0;
      for (const auto& cfl : BaseT::bins(includeOverflows)) {
        if (!cfl.get())  continue;
        sumw += cfl->sumW(includeOverflows);
      }
      return sumw;
    }
 
    double sumW2(const bool includeOverflows=true) const noexcept {
      double sumw2 = 0;
      for (const auto& cfl : BaseT::bins(includeOverflows)) {
        if (!cfl.get())  continue;
        sumw2 += cfl->sumW2(includeOverflows);
      }
      return sumw2;
    }
 
    double integral(const bool includeOverflows=true) const noexcept {
      return sumW(includeOverflows);
    }
 
    double integralError(const bool includeOverflows=true) const noexcept {
      return sqrt(sumW2(includeOverflows));
    }
 
 
 
    void scaleW(const double scalefactor) noexcept {
      for (auto& cfl : BaseT::bins(true, true)) {
        if (!cfl.get())  continue;
        cfl->scaleW(scalefactor);
      }
    }
 
    void scale(const double scalefactor) noexcept {
      for (auto& cfl : BaseT::bins(true, true)) {
        if (!cfl.get())  continue;
        cfl->scale(scalefactor);
      }
    }
 
    void normalizeStep(const std::string& edge, const double norm) {
      for (auto& cfl : BaseT::bins(true, true)) {
        if (!cfl.get())  continue;
        if (cfl->binAt(edge).sumW() != 0.)  cfl->normalizeStep(edge, norm);
      }
    }
 
    void normalizeFirst(const double norm) {
      normalizeStep(""s, norm);
    }
 
 
 
    string str() const {
      std::stringstream ss;
      for (auto& cf : BaseT::bins(true, true)) {
        if (!cf.get())  continue;
        ss << cf << "\n\n";
      }
      return ss.str();
    }
 
    void print(std::ostream& os) const {
      os << str() << std::flush;
    }
 
 
  };
 
 
  using CutflowsPtr = std::shared_ptr<Cutflows>;
 
  inline std::ostream& operator << (std::ostream& os, const Cutflows& cfs) {
    return os << cfs.str();
  }
 
  inline std::ostream& operator << (std::ostream& os, const CutflowsPtr& cfs) {
    return os << cfs->str();
  }
 
 
}
 
#endif