KVEventMixerN.h

#pragma once
 
#include <array>
#include <vector>
#include <deque>
#include "KVRefVec.h"
 
template<int NPart, typename ParticleClass, typename ParticleInfoStruct, int NumBins = 1>
class KVEventMixerN {
 
   struct event {
      std::vector<ParticleInfoStruct> particles;
 
      event() = default;
      event(event&&) = default;
      event(const event&) = default;
      event(const ParticleInfoStruct& p)
      {
         add(p);
      }
      void add(const ParticleInfoStruct& p)
      {
         particles.push_back(p);
      }
   };
 
   using event_buffer = std::deque<event>;
   struct bin_data {
      std::vector<event_buffer> part_events = std::vector<event_buffer>(NPart);
   };
   bin_data bins[NumBins];
 
   typename event_buffer::size_type decor_events = 10;
 
   int current_bin_number;
   std::array<int, NPart> n_part; 
   std::array<bool, NPart> mult_fixed; 
   std::array<bool, NPart> buffer_ready; 
   std::array<bool, NPart> can_decorrelate; 
 
   KVRefVec<ParticleClass> correlated_particles{NPart};
   KVRefVec<ParticleInfoStruct> uncorrelated_particles{NPart - 1}; 
 
public:
   KVEventMixerN(typename event_buffer::size_type number_of_events_to_mix = 10)
      : decor_events{number_of_events_to_mix}
   {}
 
   template<typename TreatCorFunc>
   void process_event(int, TreatCorFunc TreatCor)
   {
      // After recursive calls to process_event() have emptied the parameter pack part_iterator_list
      // (in other words, iterations over all particle types have been treated and the vector correlated_particles
      // has been filled with one particle of each type from the current event), this method is
      // called in order to dispatch the call to the user's function for correlated multiplets, TreatCor()
      TreatCor(current_bin_number, correlated_particles);
   }
 
   template<typename TreatCorFunc, typename part_iterator, typename... part_iterator_list>
   void process_event(int part_index, TreatCorFunc TreatCor, part_iterator& iter_part, part_iterator_list... part_iterators)
   {
      // This method is called recursively in order to iterate over each type of particle in the current event,
      // calculate the multiplicities (stored in n_part[] vector), and store in the buffers for the
      // current bin/event class. The correlated_particles vector is filled with NPart references to
      // a particle of each type, in order to call the user's TreatCor() function.
      ++part_index;
      bool count_and_store = true;
      // for part_index>0 we only increment the multipicity and store the particle in buffers the first
      // time that we visit, i.e. only if n_part[0]==1 && n_part[1]==1 && ... && n_part[part_index-1]==1
      if (part_index > 0) {
         for (int i = 0; i < part_index; ++i) count_and_store &= (n_part[i] == 1);
      }
      for (auto& part : iter_part) {
         if (count_and_store) {
            ++n_part[part_index]; // count multiplicity of particles
            if (n_part[part_index] == 1) bins[current_bin_number].part_events[part_index].push_back(ParticleInfoStruct(part));
            else bins[current_bin_number].part_events[part_index].back().add(ParticleInfoStruct(part));
         }
         // store reference to particle for call to TreatCor
         correlated_particles.set(part_index, &part);
         process_event(part_index, TreatCor, part_iterators...);
      }
   }
 
   void fill_all_particle_buffers(int)
   {
      // Dummy method, called at end of recursive calls to fill_all_particle_buffers() when all
      // iterators in the parameter pack part_iterators have been treated.
   }
 
   template<typename part_iterator, typename... part_iterator_list>
   void fill_all_particle_buffers(int part_index, part_iterator& iter_part, part_iterator_list... part_iterators)
   {
      // This method is called recursively to ensure that all particles from the current event are stored in
      // the decorrelating buffers, even if the event does not contain at least one particle of each of the NPart types
      // (depending on the user's event selection).
      ++part_index;
      if (part_index == 0) fill_all_particle_buffers(part_index, part_iterators...);
      else {
         if (!n_part[part_index] && (!n_part[part_index - 1] || mult_fixed[part_index - 1])) {
            mult_fixed[part_index] = true;
            for (auto& part : iter_part) {
               ++n_part[part_index]; // count multiplicity of particles
               if (n_part[part_index] == 1) bins[current_bin_number].part_events[part_index].push_back(ParticleInfoStruct(part));
               else bins[current_bin_number].part_events[part_index].back().add(ParticleInfoStruct(part));
            }
         }
         fill_all_particle_buffers(part_index, part_iterators...);
      }
   }
 
   template<typename TreatNCorFunc>
   void recursive_decorrelation(ParticleClass& part, int part_index, int decor_index, int decor_vec_index, TreatNCorFunc TreatNCor)
   {
      if (decor_index == NPart) {
         TreatNCor(current_bin_number, part, uncorrelated_particles);
      }
      else if (decor_index != part_index) {
         // loop over events & particles in buffer
         auto& event_buffer = bins[current_bin_number].part_events[decor_index];
         for (size_t i = 0; i < decor_events; ++i) {
            for (auto& _part : event_buffer[i].particles) {
               // add to uncorrelated particle vector
               uncorrelated_particles.set(decor_vec_index, &_part);
               recursive_decorrelation(part, part_index, (decor_index + 1), (decor_vec_index + 1), TreatNCor);
            }
         }
      }
      else // can't decorrelate particle with particles of same type
         recursive_decorrelation(part, part_index, (decor_index + 1), decor_vec_index, TreatNCor);
   }
 
   template<typename TreatNCorFunc>
   void decorrelate(int, TreatNCorFunc)
   {}
 
   template<typename TreatNCorFunc, typename part_iterator, typename... part_iterator_list>
   void decorrelate(int part_index, TreatNCorFunc TreatNCor, part_iterator& iter_part, part_iterator_list... part_iterators)
   {
      // DECORRELATION
      // Each particle of type part_index in the current event is coupled with all the particles in the buffers
      // of the other types of particles as long as they are all sufficiently full i.e. with decor_events in them
      ++part_index;
      if (n_part[part_index] && can_decorrelate[part_index]) {
         for (auto& part : iter_part) {
            recursive_decorrelation(part, part_index, 0, 0, TreatNCor);
         }
      }
      decorrelate(part_index, TreatNCor, part_iterators...);
   }
 
   template<typename TreatCorFunc, typename TreatNCorFunc, typename part_iterator, typename... part_iterator_list>
   void ProcessEvent(int bin_number, TreatCorFunc TreatCor, TreatNCorFunc TreatNCor, const part_iterator& iter_part, part_iterator_list... part_iterators)
   {
      current_bin_number = bin_number;
      // set all particle multiplicities to zero
      for (int i = 0; i < NPart; ++i) {
         n_part[i] = 0;
         mult_fixed[i] = false;
         can_decorrelate[i] = true;
      }
 
      // recursively iterate over all N-tuples of particles, updating multiplicities, filling buffers,
      // and calling the TreatCor function
      process_event(-1, TreatCor, iter_part, part_iterators...);
 
      // make sure all particles from event are added to mixing buffers
      // in process_event(), if there are no A particles, B,C,D,... are not even looked at
      //                     if there are A particles but not B, C,D,... are not even looked at
      //                     if there are A & B particles but not C, D,... are not even looked at etc.
      fill_all_particle_buffers(-1, iter_part, part_iterators...);
 
      // now check each buffer is ready for decorrelation -
      // either:
      //         the buffer contains decor_events (the current event was not added)
      // or
      //         the buffer contains decor_events+1 (the current event was added)
      for (int i = 0; i < NPart; ++i) {
         auto buf_siz = bins[current_bin_number].part_events[i].size();
         buffer_ready[i] = ((buf_siz == decor_events) && !n_part[i]) || ((buf_siz == decor_events + 1) && n_part[i]);
         // this buffer can be used to decorrelate all particle types except itself
         for (int j = 0; j < NPart; ++j) {
            if (j == i) continue;
            can_decorrelate[j] &= buffer_ready[i];
         }
      }
 
      // DECORRELATION
      decorrelate(-1, TreatNCor, iter_part, part_iterators...);
 
      // now we remove (if necessary) the oldest event in each buffer, to keep only decor_events
      for (int i = 0; i < NPart; ++i) {
         if (bins[current_bin_number].part_events[i].size() > decor_events) {
            // remove oldest (first) event
            bins[current_bin_number].part_events[i].pop_front();
         }
      }
   }
};