Generator.h

Go to the documentation of this file.

#ifndef AMPGEN_GENERATOR_H
#define AMPGEN_GENERATOR_H
 
#include "AmpGen/EventList.h"
#if ENABLE_AVX
#include "AmpGen/EventListSIMD.h"
#endif
#include "AmpGen/simd/utils.h"
#include "AmpGen/EventType.h"
#include "AmpGen/PhaseSpace.h"
#include "AmpGen/TreePhaseSpace.h"
#include "AmpGen/Utilities.h"
#include "AmpGen/ProfileClock.h"
#include "AmpGen/ProgressBar.h"
#include "AmpGen/MetaUtils.h"
 
namespace AmpGen {
  template <typename phaseSpace_t = PhaseSpace> class Generator {
#if ENABLE_AVX
    using eventlist_t = EventListSIMD;
#else
    using eventlist_t = EventList;
#endif
 
  private:
    EventType m_eventType;
    phaseSpace_t m_gps;
    size_t m_generatorBlock = {5000000};
    TRandom *m_rnd = {gRandom};
    bool m_normalise = {true};
 
  public:
    template <typename... ARGS> explicit Generator(const ARGS &...args) : m_gps(args...) {
      m_eventType = m_gps.eventType();
      if(m_rnd != gRandom) setRandom(m_rnd);
    }
 
    phaseSpace_t phsp() { return m_gps; }
 
    void setRandom(TRandom *rand) {
      m_rnd = rand;
      m_gps.setRandom(m_rnd);
    }
    void setBlockSize(const size_t &blockSize) { m_generatorBlock = blockSize; }
    void setNormFlag(const bool &normSetting) { m_normalise = normSetting; }
 
    void fillEventListPhaseSpace(eventlist_t &events, const size_t &N) {
      if constexpr(std::is_same<phaseSpace_t, PhaseSpace>::value) {
        constexpr auto w = utils::size<real_v>::value;
        for(unsigned i = 0; i != N; ++i) {
          double *addr = reinterpret_cast<double *>(events.block(i / w)) + i % w;
          m_gps.fill(addr, w);
        }
      } else {
        auto it = events.begin();
        while(it != events.end()) {
          *it = m_gps.makeEvent();
          ++it;
        }
      }
    }
    template <typename pdf_t> double getMax(const eventlist_t &events, pdf_t &pdf) const {
      double max = 0.;
      for(const auto &evt : events) {
        auto value = evt.genPdf();
        if(std::isnan(value)) {
          ERROR("PDF for event is nan: " << value);
          evt.print();
          pdf.debug(evt);
        } else if(value > max)
          max = value;
      }
      DEBUG("Returning normalisation constant = " << max);
      return max;
    }
 
    template <typename eventList_t, typename pdf_t> void fillEventList(pdf_t &pdf, eventList_t &list, const size_t &N) {
      if(m_rnd == nullptr) {
        ERROR("Random generator not set!");
        return;
      }
      double maxProb = m_normalise ? 0 : 1;
      auto size0 = list.size();
      double totalGenerated = 0;
      pdf.reset(true);
      ProgressBar pb(60, detail::trimmedString(__PRETTY_FUNCTION__));
      ProfileClock t_phsp, t_eval, t_acceptReject, t_total;
      std::vector<bool> efficiencyReport(m_generatorBlock, false);
 
      while(list.size() - size0 < N) {
        eventlist_t mc(m_eventType);
        mc.resize(m_generatorBlock);
        t_phsp.start();
        fillEventListPhaseSpace(mc, m_generatorBlock);
        t_phsp.stop();
        t_eval.start();
        pdf.setEvents(mc);
        pdf.prepare();
        auto previousSize = list.size();
#ifdef _OPENMP
#pragma omp parallel for
#endif
        for(size_t block = 0; block < mc.nBlocks(); ++block) { mc.setGenPDF(block, pdf(mc.block(block), block) / mc.genPDF(block)); }
        maxProb = maxProb == 0 ? 1.5 * getMax(mc, pdf) : maxProb;
        DEBUG("Norm: " << maxProb);
        // if constexpr ( std::is_same<phaseSpace_t, TreePhaseSpace>::value ) m_gps.recalculate_weights(mc);
 
        t_eval.stop();
        t_acceptReject.start();
        totalGenerated += mc.size();
        for(const auto &event : mc) {
          if(event.genPdf() > maxProb) {
            std::cout << std::endl;
            WARNING("PDF value exceeds norm value: " << event.genPdf() << " > " << maxProb);
            event.print();
          }
          if(event.genPdf() > maxProb * m_rnd->Rndm()) {
            list.push_back(event);
            list.rbegin()->setGenPdf(pdf(event));
            efficiencyReport[event.index()] = true;
          } else
            efficiencyReport[event.index()] = false;
          if(list.size() - size0 == N) break;
        }
        t_acceptReject.stop();
        double efficiency = 100. * (list.size() - previousSize) / (double)m_generatorBlock;
        pb.print(double(list.size()) / double(N),
                 " ε[gen] = " + mysprintf("%.4f", efficiency) + "% , " + std::to_string(int(t_total.count() / 1000.)) + " seconds");
        if(list.size() == previousSize) {
          ERROR("No events generated, PDF: " << type_string<pdf_t>() << " is likely to be malformed");
          break;
        }
      }
      pb.finish();
      t_total.stop();
      INFO("Generated " << N << " events in " << t_total << " ms");
      INFO("Generating phase space : " << t_phsp << " ms");
      INFO("Evaluating PDF         : " << t_eval << " ms");
      INFO("Accept/reject          : " << t_acceptReject << " ms");
      INFO("Efficiency             = " << double(N) * 100. / totalGenerated << " %");
    }
    template <typename pdf_t, typename = typename std::enable_if<!std::is_integral<pdf_t>::value>::type> EventList generate(pdf_t &pdf, const size_t &nEvents) {
      eventlist_t evts(m_eventType);
      fillEventList(pdf, evts, nEvents);
      EventList output(m_eventType);
      for(const auto &event : evts) output.emplace_back(event);
      return output;
    }
    EventList generate(const size_t &nEvents) {
      eventlist_t evts(m_eventType);
      evts.resize(nEvents);
      fillEventListPhaseSpace(evts, nEvents);
      EventList output(m_eventType);
      for(const auto &event : evts) output.emplace_back(event);
      return output;
    }
  };
 
  template <class FCN> class PDFWrapper {
  public:
    void prepare() {};
    void setEvents(AmpGen::EventList & /*evts*/) {};
    double prob_unnormalised(const AmpGen::Event &evt) const { return m_fcn(evt); }
    explicit PDFWrapper(const FCN &fcn) : m_fcn(fcn) {}
    size_t size() const { return 0; }
    void reset(const bool & /*flag*/ = false) {};
 
  private:
    FCN m_fcn;
  };
 
  template <class FCN> PDFWrapper<FCN> make_pdf(const FCN &fcn) { return PDFWrapper<FCN>(fcn); }

  extern "C" void python__generate(const char *eventType, double *out, const unsigned int size);
} // namespace AmpGen
#endif