CentralityBinner.hh

// -*- C++ -*-
#ifndef RIVET_CENTRALITYBINNER_HH
#define RIVET_CENTRALITYBINNER_HH
#include "Rivet/Config/RivetCommon.hh"
#include "Rivet/Projections/HepMCHeavyIon.hh"
#include "Rivet/Tools/RivetYODA.hh"
#include <tuple>
 
namespace Rivet {
 

  class CentralityEstimator : public Projection {
  public:

    CentralityEstimator()
        : _estimate(-1.0) {
      setName("CentralityEstimator");
      declare(HepMCHeavyIon(), "HepMC");
    }

    RIVET_DEFAULT_PROJ_CLONE(CentralityEstimator);
 
  protected:

    void project(const Event& e) {
      _estimate = -1.0;
      double imp = apply<HepMCHeavyIon>(e, "HepMC").impact_parameter();
      if (imp < 0.0) return;
      _estimate = imp > 0.0 ? 1.0 / imp : numeric_limits<double>::max();
    }

    CmpState compare(const Projection& p) const {
      return mkNamedPCmp(p, "HepMC");
    }
 
 
  public:

    double estimate() const {
      return _estimate;
    }
 
 
  protected:

    double _estimate;
  };
 

  template <typename T>
  struct CentralityBinTraits {

    static T clone(const T& t) {
      return T(t->newclone());
    }

    static void add(T& t, const T& o) {
      *t += *o;
    }

    static void scale(T& t, double f) {
      t->scaleW(f);
    }

    static void normalize(T& t, double sumw) {
      if (t->sumW() > 0.0) t->normalize(t->sumW() / sumw);
    }

    static string path(T t) {
      return t->path();
    }
  };

  struct MergeDistance {

    static double dist(double cestLo, double cestHi, double weight, double clo, double chi, double, double) {
      return (cestHi - cestLo) * weight / (cestHi * (chi - clo));
    }
  };
 

  template <typename T = Histo1DPtr, typename MDist = MergeDistance>
  class CentralityBinner : public ProjectionApplier {
  public:

    CentralityBinner(int maxbins = 200, double wlim = 0.02)
        : _currentCEst(-1.0), _maxBins(maxbins), _warnlimit(wlim), _weightsum(0.0) {
      _percentiles.insert(0.0);
      _percentiles.insert(1.0);
    }

    void setProjection(const CentralityEstimator& p, string pname) {
      declare(p, pname);
      _estimator = pname;
    }

    virtual std::string name() const {
      return "Rivet::CentralityBinner";
    }



 
    void add(T t, double cmin, double cmax, double cestmin = -1.0, double cestmax = -1.0) {
      _percentiles.insert(max(1.0 - cmax / 100.0, 0.0));
      _percentiles.insert(min(1.0 - cmin / 100.0, 1.0));
      if (_unfilled.empty() && _ready.empty()) _devnull = CentralityBinTraits<T>::clone(t);
      if (cestmin < 0.0)
        _unfilled.push_back(Bin(t, 1.0 - cmax / 100.0, 1.0 - cmin / 100.0));
      else
        _ready[t] = Bin(t, 1.0 - cmax / 100.0, 1.0 - cmin / 100.0, cestmin, cestmax);
    }

    T select(const Event& event, double weight = 1.0) {
      return select(applyProjection<CentralityEstimator>(event, _estimator).estimate(), weight);
    }

    T select(double cest, double weight = 1.0);

    void finalize();

    void normalizePerEvent() {
      for (auto& b : _ready) b.second.normalizePerEvent();
    }

    map<double, double> edges() const {
      map<double, double> ret;
      for (auto& b : _ready) {
        ret[1.0 - b.second._centLo] = b.second._cestLo;
        ret[1.0 - b.second._centHi] = b.second._cestHi;
      }
      return ret;
    }

    const T& current() const {
      return _currenT;
    }

    double estimator() const {
      return _currentCEst;
    }
 
    vector<T> allObjects() {
      vector<T> ret;
      for (auto& fb : _flexiBins) ret.push_back(fb._t);
      if (!ret.empty()) return ret;
      for (auto b : _ready) ret.push_back(b.second._t);
      return ret;
    }
 
  private:

    struct FlexiBin {

      FlexiBin(T& t, double cest = 0.0, double weight = 0.0)
          : _t(t), _cestLo(cest), _cestHi(cest), _weightsum(weight), _n(1), _m(0) { }

      FlexiBin(double cest)
          : _cestLo(cest), _cestHi(cest), _weightsum(0.0), _n(0), _m(0) { }

      void merge(const FlexiBin& fb) {
        _cestLo = min(_cestLo, fb._cestLo);
        _cestHi = max(_cestHi, fb._cestHi);
        _weightsum += fb._weightsum;
        CentralityBinTraits<T>::add(_t, fb._t);
        _n += fb._n;
        _m += fb._m + 1;
      }

      bool operator<(const FlexiBin& fb) const {
        return _cestLo < fb._cestLo;
      }

      bool inRange(double cest) const {
        return cest == _cestLo || (_cestLo < cest && cest < _cestHi);
      }

      T _t;

      double _cestLo, _cestHi;

      mutable double _weightsum;

      mutable int _n;

      mutable int _m;
    };
 
    struct Bin {

      Bin()
          : _centLo(-1.0),
            _centHi(-1.0),
            _cestLo(-1.0),
            _cestHi(-1.0),
            _weightsum(0.0),
            _underflow(0.0),
            _overflow(0.0),
            _ambiguous(0),
            _ambweight(0.0) { }

      Bin(T t, double centLo, double centHi, double cestLo = -1.0, double cestHi = -1.0)
          : _t(t),
            _centLo(centLo),
            _centHi(centHi),
            _cestLo(cestLo),
            _cestHi(cestHi),
            _weightsum(0.0),
            _underflow(0.0),
            _overflow(0.0),
            _ambiguous(0.0),
            _ambweight(0.0) { }

      bool inRange(double cest) const {
        return _cestLo >= 0 && _cestLo <= cest && (_cestHi < 0.0 || cest <= _cestHi);
      }

      void normalizePerEvent() {
        CentralityBinTraits<T>::normalize(_t, _weightsum);
      }

      T _t;

      double _centLo, _centHi;

      double _cestLo, _cestHi;

      double _weightsum;

      double _underflow;

      double _overflow;

      double _ambiguous;

      double _ambweight;
    };
 
  protected:

    typedef set<FlexiBin> FlexiBinSet;

    typename FlexiBinSet::iterator _findBin(double cest) {
      if (_flexiBins.empty()) return _flexiBins.end();
      auto it = _flexiBins.lower_bound(FlexiBin(cest));
      if (it->_cestLo == cest) return it;
      if (it != _flexiBins.begin()) --it;
      if (it->_cestLo < cest && cest < it->_cestHi) return it;
      return _flexiBins.end();
    }

    string _estimator;

    T _currenT;

    double _currentCEst;

    int _maxBins;

    double _warnlimit;

    vector<Bin> _unfilled;

    FlexiBinSet _flexiBins;

    double _weightsum;

    set<double> _percentiles;

    map<T, Bin> _ready;

    T _devnull;
 
  public:

    void debug();
    void fulldebug();
  };
 

  template <>
  struct CentralityBinTraits<Profile1DPtr> {
 
    typedef Profile1DPtr T;

    static T clone(const T& t) {
      return Profile1DPtr(t->newclone());
    }

    static void add(T& t, const T& o) {
      *t += *o;
    }

    static void scale(T& t, double f) {
      t->scaleW(f);
    }
 
    static void normalize(T& t, double sumw) { }

    static string path(T t) {
      return t->path();
    }
  };
 

  template <>
  struct CentralityBinTraits<Profile2DPtr> {
 
    typedef Profile2DPtr T;

    static T clone(const T& t) {
      return Profile2DPtr(t->newclone());
    }

    static void add(T& t, const T& o) {
      *t += *o;
    }

    static void scale(T& t, double f) {
      t->scaleW(f);
    }
 
    static void normalize(T& t, double sumw) { }

    static string path(T t) {
      return t->path();
    }
  };
 
  template <typename T>
  struct CentralityBinTraits<vector<T>> {

    static vector<T> clone(const vector<T>& tv) {
      vector<T> rtv;
      for (auto t : tv) rtv.push_back(CentralityBinTraits<T>::clone(t));
      return rtv;
    }

    static void add(vector<T>& tv, const vector<T>& ov) {
      for (int i = 0, N = tv.size(); i < N; ++i) CentralityBinTraits::add(tv[i], ov[i]);
    }

    static void scale(vector<T>& tv, double f) {
      for (auto t : tv) CentralityBinTraits<T>::scale(t, f);
    }
 
    static void normalize(vector<T>& tv, double sumw) {
      for (auto t : tv) CentralityBinTraits<T>::normalize(t, sumw);
    }

    static string path(const vector<T>& tv) {
      string ret = "(vector:";
      for (auto t : tv) {
        ret += " ";
        ret += CentralityBinTraits<T>::path(t);
      }
      ret += ")";
      return ret;
    }
  };
 
  template <size_t I, typename... Types>
  struct TupleCentralityBinTraitsHelper {
 
    typedef tuple<Types...> Tuple;
    typedef typename tuple_element<I - 1, Tuple>::type T;
 
    static void clone(Tuple& ret, const Tuple& tup) {
      get<I - 1>(ret) = CentralityBinTraits<T>::clone(get<I - 1>(tup));
      TupleCentralityBinTraitsHelper<I - 1, Types...>::clone(ret, tup);
    }
 
    static void add(Tuple& tup, const Tuple& otup) {
      CentralityBinTraits<T>::add(get<I - 1>(tup), get<I - 1>(otup));
      TupleCentralityBinTraitsHelper<I - 1, Types...>::add(tup, otup);
    }
 
    static void scale(Tuple& tup, double f) {
      CentralityBinTraits<T>::scale(get<I - 1>(tup), f);
      TupleCentralityBinTraitsHelper<I - 1, Types...>::scale(tup, f);
    }
 
    static void normalize(Tuple& tup, double sumw) {
      CentralityBinTraits<T>::normalize(get<I - 1>(tup), sumw);
      TupleCentralityBinTraitsHelper<I - 1, Types...>::normalize(tup, sumw);
    }
 
    static string path(const Tuple& tup) {
      return " " + CentralityBinTraits<T>::path(get<I - 1>(tup))
          + TupleCentralityBinTraitsHelper<I - 1, Types...>::path(tup);
    }
  };
 
  template <typename... Types>
  struct TupleCentralityBinTraitsHelper<0, Types...> {
 
    typedef tuple<Types...> Tuple;
 
    static void clone(Tuple&, const Tuple&) { }
    static void add(Tuple& tup, const Tuple& otup) { }
    static void scale(Tuple& tup, double f) { }
    static void normalize(Tuple& tup, double sumw) { }
    static string path(const Tuple& tup) {
      return "";
    }
  };
 
  template <typename... Types>
  struct CentralityBinTraits<tuple<Types...>> {
 
    typedef tuple<Types...> Tuple;
    static const size_t N = tuple_size<Tuple>::value;

    static Tuple clone(const Tuple& tup) {
      Tuple ret;
      TupleCentralityBinTraitsHelper<N, Types...>::clone(ret, tup);
      return ret;
    }

    static void add(Tuple& tup, const Tuple& otup) {
      TupleCentralityBinTraitsHelper<N, Types...>::add(tup, otup);
    }

    static void scale(Tuple& tup, double f) {
      TupleCentralityBinTraitsHelper<N, Types...>::scale(tup, f);
    }
 
    static void normalize(Tuple& tup, double sumw) {
      TupleCentralityBinTraitsHelper<N, Types...>::normalize(tup, sumw);
    }

    static string path(const Tuple& tup) {
      string ret = "(tuple:";
      ret += TupleCentralityBinTraitsHelper<N, Types...>::path(tup);
      ret += ")";
      return ret;
    }
  };
 
  template <typename T, typename MDist>
  T CentralityBinner<T, MDist>::select(double cest, double weight) {
    _currenT = _devnull;
    _currentCEst = cest;
    _weightsum += weight;
 
    // If estimator is negative, something has gone wrong.
    if (_currentCEst < 0.0) return _currenT;
 
    // If we already have finalized the limits on the centrality
    // estimator, we just add the weights to their bins and return the
    // corresponding AnalysisObject.
    if (_unfilled.empty()) {
      for (auto& b : _ready)
        if (b.second.inRange(_currentCEst)) {
          b.second._weightsum += weight;
          return b.second._t;
        }
      return _currenT;
    }
 
    auto it = _findBin(cest);
    if (it == _flexiBins.end()) {
      _currenT = CentralityBinTraits<T>::clone(_unfilled.begin()->_t);
      it = _flexiBins.insert(FlexiBin(_currenT, _currentCEst, weight)).first;
    }
    else {
      it->_weightsum += weight;
      ++(it->_n);
      _currenT = it->_t;
    }
 
    if ((int)_flexiBins.size() <= _maxBins) return _currenT;
 
 
    set<double>::iterator citn = _percentiles.begin();
    set<double>::iterator cit0 = citn++;
    auto selectit = _flexiBins.end();
    double mindist = -1.0;
    double acc = 0.0;
    auto next = _flexiBins.begin();
    auto prev = next++;
    for (; next != _flexiBins.end(); prev = next++) {
      acc += prev->_weightsum / _weightsum;
      if (acc > *citn) {
        cit0 = citn++;
        continue;
      }
      if (acc + next->_weightsum / _weightsum > *citn) continue;
      double dist = MDist::dist(prev->_cestLo, next->_cestHi, next->_weightsum + prev->_weightsum, *cit0,
                                *citn, next->_n + prev->_n, next->_m + prev->_m);
      if (mindist < 0.0 || dist < mindist) {
        selectit = prev;
        mindist = dist;
      }
    }
 
    if (selectit == _flexiBins.end()) return _currenT;
    auto mergeit = selectit++;
    FlexiBin merged = *mergeit;
    merged.merge(*selectit);
    if (merged.inRange(cest) || selectit->inRange(cest)) _currenT = merged._t;
    _flexiBins.erase(mergeit);
    _flexiBins.erase(selectit);
    _flexiBins.insert(merged);
 
    return _currenT;
  }
 
 
  template <typename T, typename MDist>
  void CentralityBinner<T, MDist>::finalize() {
 
    // Take the contents of the dynamical binning and fill the original
    // AnalysisObjects.
 
    double clo = 0.0;
    for (const FlexiBin& fb : _flexiBins) {
      double chi = min(clo + fb._weightsum / _weightsum, 1.0);
      for (Bin& bin : _unfilled) {
        double olo = bin._centLo;
        double ohi = bin._centHi;
        if (clo > ohi || chi <= olo) continue;
        // If we only have partial overlap we need to scale
        double lo = max(olo, clo);
        double hi = min(ohi, chi);
        T t = CentralityBinTraits<T>::clone(fb._t);
        double frac = (hi - lo) / (chi - clo);
        CentralityBinTraits<T>::scale(t, frac);
        CentralityBinTraits<T>::add(bin._t, t);
        bin._weightsum += fb._weightsum * frac;
        if (clo <= olo) bin._cestLo = fb._cestLo + (fb._cestHi - fb._cestLo) * (olo - clo) / (chi - clo);
        if (clo < olo) {
          bin._underflow = clo;
          bin._ambiguous += fb._n * frac;
          bin._ambweight += fb._weightsum * frac * (1.0 - frac);
        }
        if (chi > ohi) {
          bin._cestHi = fb._cestLo + (fb._cestHi - fb._cestLo) * (ohi - clo) / (chi - clo);
          bin._overflow = chi;
          bin._ambiguous += fb._n * frac;
          bin._ambweight += fb._weightsum * frac * (1.0 - frac);
        }
      }
      clo = chi;
    }
    _flexiBins.clear();
    for (Bin& bin : _unfilled) {
      if (bin._overflow == 0.0) bin._overflow = 1.0;
      _ready[bin._t] = bin;
      if (bin._ambweight / bin._weightsum > _warnlimit)
        MSG_WARNING("Analysis object \""
                    << CentralityBinTraits<T>::path(bin._t)
                    << "\", contains events with centralities between " << bin._underflow * 100.0 << " and "
                    << bin._overflow * 100.0 << "% (" << int(bin._ambiguous + 0.5)
                    << " ambiguous events with effectively " << 100.0 * bin._ambweight / bin._weightsum
                    << "% of the weights)."
                    << "Consider increasing the number of bins.");
    }
    _unfilled.clear();
  }
 
  template <typename T, typename MDist>
  void CentralityBinner<T, MDist>::fulldebug() {
    cerr << endl;
    double acc = 0.0;
    set<double>::iterator citn = _percentiles.begin();
    set<double>::iterator cit0 = citn++;
    int i = 0;
    for (auto it = _flexiBins.begin(); it != _flexiBins.end();) {
      ++i;
      auto curr = it++;
      double w = curr->_weightsum / _weightsum;
      acc += w;
      if (curr == _flexiBins.begin() || it == _flexiBins.end() || acc > *citn)
        cerr << "*";
      else
        cerr << " ";
      if (acc > *citn) cit0 = citn++;
      cerr << setw(6) << i << setw(12) << acc - w << setw(12) << acc << setw(8) << curr->_n << setw(8)
           << curr->_m << setw(12) << curr->_cestLo << setw(12) << curr->_cestHi << endl;
    }
    cerr << "Number of sampler bins: " << _flexiBins.size() << endl;
  }
 
  template <typename T, typename MDist>
  void CentralityBinner<T, MDist>::debug() {
    cerr << endl;
    double acc = 0.0;
    int i = 0;
    set<double>::iterator citn = _percentiles.begin();
    set<double>::iterator cit0 = citn++;
    for (auto it = _flexiBins.begin(); it != _flexiBins.end();) {
      auto curr = it++;
      ++i;
      double w = curr->_weightsum / _weightsum;
      acc += w;
      if (curr == _flexiBins.begin() || it == _flexiBins.end() || acc > *citn) {
        if (acc > *citn) cit0 = citn++;
        cerr << setw(6) << i << setw(12) << acc - w << setw(12) << acc << setw(8) << curr->_n << setw(8)
             << curr->_m << setw(12) << curr->_cestLo << setw(12) << curr->_cestHi << endl;
      }
    }
    cerr << "Number of sampler bins: " << _flexiBins.size() << endl;
  }

  class GeneratedCentrality : public CentralityEstimator {
 
  public:

    GeneratedCentrality() {
      declare(HepMCHeavyIon(), "HI");
    }

    RIVET_DEFAULT_PROJ_CLONE(GeneratedCentrality);
 
  protected:

    void project(const Event& e) {
      _estimate = apply<HepMCHeavyIon>(e, "HI").centrality();
    }

    CmpState compare(const Projection& p) const {
      return mkNamedPCmp(p, "GeneratedCentrality");
    }
  };
 
 
}
 
#endif