SumPDF.h

Go to the documentation of this file.

#ifndef AMPGEN_SUMPDF_H
#define AMPGEN_SUMPDF_H
 
#include "AmpGen/MetaUtils.h"
#include "AmpGen/MsgService.h"
#include "AmpGen/ProfileClock.h"
#include "AmpGen/KeyedFunctors.h"
#include "AmpGen/KahanSum.h"
#include <tuple>
#if ENABLE_AVX
#include "AmpGen/simd/utils.h"
#endif
 
namespace AmpGen {
  class EventList;
  class EventListSIMD;
 
  template <class eventListType, class... pdfTypes> class SumPDF {
  private:
    typedef typename eventListType::value_type eventValueType; 
    std::tuple<pdfTypes...> m_pdfs;                            
    const eventListType *m_events = {nullptr};                 
 
  public:
    SumPDF() = default;

    SumPDF(const pdfTypes &...pdfs) : m_pdfs(std::tuple<pdfTypes...>(pdfs...)) {}

 
    double getVal() {
      std::vector<real_v> tmp(m_events->nBlocks());
      fill_likelihood(tmp.data());
      KahanSum<real_v> sum;
      for(unsigned block = 0; block != tmp.size(); ++block) sum += tmp[block];
      auto rt = -2 * utils::sum_elements(sum());
      return rt;
    }
    void fill_likelihood(real_v *output) {
      for_each(m_pdfs, [](auto &f) { f.prepare(); });
      if constexpr(std::is_same<eventListType, EventList>::value) {
#pragma omp parallel for
        for(unsigned int i = 0; i < m_events->size(); ++i) {
          auto prob = ((*this))((*m_events)[i]);
          auto w = (*m_events)[i].weight();
          output[i] = w * std::log(prob);
        }
      }
#if ENABLE_AVX
      if constexpr(std::is_same<eventListType, EventListSIMD>::value) {
#pragma omp parallel for
        for(unsigned block = 0; block < m_events->nBlocks(); ++block) { output[block] = m_events->weight(block) * AVX::log(this->operator()(nullptr, block)); }
      }
#endif
    }

    real_v operator()(const real_v *evt, const unsigned block) {
      real_v prob = 0.;
      for_each(this->m_pdfs, [&prob, &evt, block](const auto &f) mutable { prob += f(evt, block); });
      return prob;
    }

    double operator()(const eventValueType &evt) {
      double prob = 0;
      for_each(this->m_pdfs, [&prob, &evt](const auto &f) mutable { prob += f(evt); });
      return prob;
    }

    void setEvents(eventListType &events) {
      m_events = &events;
      for_each(m_pdfs, [&events](auto &f) { f.setEvents(events); });
    }

    std::size_t nPDFs() const { return sizeof...(pdfTypes); }

    std::tuple<pdfTypes...> pdfs() const { return m_pdfs; }
 
    std::function<double(const eventValueType &)> evaluator(const eventListType *events) const {
      std::vector<double> values(events->size());
      for_each(this->m_pdfs, [events, &values](const auto &pdf) mutable {
        auto eval = pdf.evaluator(events);
        for(unsigned i = 0; i != events->size(); ++i) values[i] += eval(events->at(i));
      });
      return arrayToFunctor<double, typename eventListType::value_type>(values);
    }
    KeyedFunctors<double(eventValueType)> componentEvaluator(const eventListType *events) const {
      KeyedFunctors<double(eventValueType)> view;
      for_each(this->m_pdfs, [&view, &events](const auto &pdf) mutable {
        auto eval = pdf.evaluator(events);
        view.add([eval](const auto &event) { return eval(event); }, type_string(pdf), "");
      });
      return view;
    }
  };

  template <class eventListType = EventList, class... pdfTypes> auto make_pdf(pdfTypes &&...pdfs) {
    // return SumPDF<eventListType, pdfTypes...>( std::forward<pdfTypes>( pdfs )... );
    return SumPDF<eventListType, pdfTypes...>(pdfs...);
  }
 
  template <class eventListType = EventList, class... pdfTypes> auto make_likelihood(eventListType &events, pdfTypes &&...pdfs) {
    auto rt = SumPDF<eventListType, pdfTypes...>(std::forward<pdfTypes>(pdfs)...);
    rt.setEvents(events);
    return rt;
  }
} // namespace AmpGen
 
#endif