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 <algorithm>
#include <deque>
 
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