EventMixingFinalState.hh

// -*- C++ -*-
#ifndef RIVET_EventMixingFinalState_HH
#define RIVET_EventMixingFinalState_HH
 
#include "Rivet/Projection.hh"
#include "Rivet/Projections/ParticleFinder.hh"
#include "Rivet/Tools/Random.hh"
#include <deque>
#include <algorithm>
 
namespace Rivet {
 


  template <class RandomAccessIterator,
            class WeightIterator, class RandomNumberGenerator>
  void weighted_shuffle(RandomAccessIterator first, RandomAccessIterator last,
                        WeightIterator fw, WeightIterator lw, RandomNumberGenerator& g) {
    while(first != last && fw != lw) {
      std::discrete_distribution<int> weightDist(fw, lw);
      int i = weightDist(g);
      if(i){
        std::iter_swap(first, next(first, i));
        std::iter_swap(fw, next(fw, i));
      }
      ++first;
      ++fw;
    }
  }

  typedef pair<Particles, double> MixEvent;
  typedef map<double, std::deque<MixEvent> > MixMap;

  class EventMixingBase : public Projection {
  protected:
    EventMixingBase(const Projection & mixObsProj, const ParticleFinder& mix,
                    size_t nMixIn, double oMin, double oMax, double deltao,
                    const size_t defaultIdx) : nMix(nMixIn), unitWeights(true) {
      // The base class contructor should be called explicitly in derived classes
      // to add projections below.
      setName("EventMixingBase");
      declare(mixObsProj,"OBS");
      declare(mix,"MIX");
      MSG_WARNING("EventMixing is not fully validated. Use with caution.");
 
      _defaultWeightIdx = defaultIdx;
      // Set up the map for mixing events.
      for(double o = oMin; o < oMax; o+=deltao )
        mixEvents[o] = std::deque<MixEvent>();
    }

    using Projection::operator =;
 
 
  public:

    bool hasMixingEvents() const {
      MixMap::const_iterator mixItr = mixEvents.lower_bound(mObs);
      if(mixItr == mixEvents.end() || mixItr->second.size() < nMix + 1)
        return false;
      return true;
    }

    vector<MixEvent> getMixingEvents() const {
      if (!hasMixingEvents())
        return vector<MixEvent>();
      MixMap::const_iterator mixItr = mixEvents.lower_bound(mObs);
      return vector<MixEvent>(mixItr->second.begin(), mixItr->second.end() - 1);
    }

    virtual const Particles particles() const {
      // Test if we have enough mixing events.
      if (!hasMixingEvents())
        return Particles();
      // Get mixing events for the current, projected mixing observable.
      MixMap::const_iterator mixItr = mixEvents.lower_bound(mObs);
      vector<MixEvent> mixEvents(mixItr->second.begin(), mixItr->second.end() - 1);
      // Make the vector of mixed particles.
      Particles mixParticles;
      vector<double> weights;
      size_t pSize = 0;
      for (size_t i = 0; i < mixEvents.size(); ++i)
        pSize+=mixEvents[i].first.size();
      mixParticles.reserve(pSize);
      weights.reserve(pSize);
      // Put the particles in the vector.
      for (size_t i = 0; i < mixEvents.size(); ++i) {
        mixParticles.insert(mixParticles.end(), mixEvents[i].first.begin(), mixEvents[i].first.end());
        vector<double> tmp(mixEvents[i].first.size(), mixEvents[i].second);
        weights.insert(weights.end(), tmp.begin(), tmp.end());
      }
 
      // Shuffle the particles.
      if (unitWeights) {
        // Use the thread safe random number generator.
        //auto rnd = [&] (int i) {return rng()()%i;};
        std::shuffle(mixParticles.begin(), mixParticles.end(), rng());
        return mixParticles;
      } else {
        weighted_shuffle(mixParticles.begin(), mixParticles.end(), weights.begin(), weights.end(), rng());
        Particles tmp = vector<Particle>(mixParticles.begin(), mixParticles.begin() + size_t(ceil(mixParticles.size() / 2)));
        return tmp;
      }
    }
 
 
  protected:

    virtual void calculateMixingObs(const Projection* mProj) = 0;
 

    void project(const Event& e) {
      const Projection* mixObsProjPtr = &apply<Projection>(e, "OBS");
      calculateMixingObs(mixObsProjPtr);
      MixMap::iterator mixItr = mixEvents.lower_bound(mObs);
      if (mixItr == mixEvents.end()){
        // We are out of bounds.
        MSG_DEBUG("Mixing observable out of bounds.");
        return;
      }
      const Particles mix = apply<ParticleFinder>(e, "MIX").particles();
      mixItr->second.push_back(make_pair(mix,e.weights()[_defaultWeightIdx]));
      // Assume unit weights until we see otherwise.
      if (unitWeights && e.weights()[_defaultWeightIdx] != 1.0 ) {
        unitWeights = false;
        nMix *= 2;
      }
      if (mixItr->second.size() > nMix + 1)
        mixItr->second.pop_front();
    }
 

    CmpState compare(const Projection& p) const {
      return mkNamedPCmp(p,"OBS");
    }
 

    double mObs;
 
 
  protected:

    size_t nMix;

    MixMap mixEvents;

    bool unitWeights;
 
    size_t _defaultWeightIdx;
 
  };
 

  class EventMixingFinalState : public EventMixingBase {
  public:
    EventMixingFinalState(const ParticleFinder & mixObsProj,
                          const ParticleFinder& mix, size_t nMixIn, double oMin, double oMax,
                          double deltao, const size_t defaultIdx) :
      EventMixingBase(mixObsProj, mix, nMixIn, oMin, oMax, deltao, defaultIdx) {
      setName("EventMixingFinalState");
    }
 
    RIVET_DEFAULT_PROJ_CLONE(EventMixingFinalState);

    using Projection::operator =;
 
 
  protected:

    virtual void calculateMixingObs(const Projection* mProj) {
      mObs = ((ParticleFinder*) mProj)->particles().size();
    }
 
  };
 

  class EventMixingCentrality : public EventMixingBase {
  public:
 
    EventMixingCentrality(const CentralityProjection& mixObsProj,
                          const ParticleFinder& mix, size_t nMixIn, double oMin, double oMax,
                          double deltao, const size_t defaultIdx) :
      EventMixingBase(mixObsProj, mix, nMixIn, oMin, oMax, deltao, defaultIdx) {
      setName("EventMixingCentrality");
    }
 
    RIVET_DEFAULT_PROJ_CLONE(EventMixingCentrality);

    using Projection::operator =;
 
  protected:

    virtual void calculateMixingObs(const Projection* mProj) {
      mObs = ((CentralityProjection*) mProj)->operator()();
    }
 
  };
 
 
}
 
#endif