EventList.h

Go to the documentation of this file.

#ifndef AMPGEN_EVENTLIST_H
#define AMPGEN_EVENTLIST_H
 
#include "AmpGen/ArgumentPack.h"
#include "AmpGen/EventType.h"
#include "AmpGen/MsgService.h"
#include "AmpGen/Event.h"
#include "AmpGen/Projection.h"
#include "AmpGen/Utilities.h"
#include "AmpGen/MetaUtils.h"
#include "AmpGen/Units.h"
 
#include <chrono>
#include <functional>
#include <numeric>
#include <cstddef>
#include <algorithm>
 
#include <TH1D.h>
#include <TH2D.h>
#include <TTree.h>
 
#ifdef _OPENMP
#include <omp.h>
#endif
 
namespace AmpGen {
  namespace PlotOptions {
    DECLARE_ARGUMENT(Bins, size_t);
  }
  class CompiledExpressionBase;
  class EventList {
  private:
    std::vector<Event> m_data = {};
    EventType m_eventType = {};
    std::map<std::string, unsigned int> m_extensions = {};
 
  public:
    typedef Event value_type;
    EventList() = default;
    EventList(const EventType &type);
    template <class... ARGS> EventList(const std::string &fname, const EventType &evtType, const ARGS &...args) : EventList(evtType) {
      loadFromFile(fname, ArgumentPack(args...));
    }
    template <class... ARGS> EventList(const std::string &fname, const ARGS &...args) : EventList() { loadFromFile(fname, ArgumentPack(args...)); }
    template <class... ARGS> EventList(const std::vector<std::string> &fname, const EventType &evtType, const ARGS &...args) : EventList(evtType) {
      for(auto &f : fname) loadFromFile(f, ArgumentPack(args...));
    }
    template <class... ARGS> EventList(TTree *tree, const EventType &evtType, const ARGS &...args) : EventList(evtType) {
      loadFromTree(tree, ArgumentPack(args...));
    }
    const EventList &store() const { return *this; }
    std::vector<Event>::reverse_iterator rbegin() { return m_data.rbegin(); }
    std::vector<Event>::reverse_iterator rend() { return m_data.rend(); }
    std::vector<Event>::iterator begin() { return m_data.begin(); }
    std::vector<Event>::iterator end() { return m_data.end(); }
    Event &operator[](const size_t &pos) { return m_data[pos]; }
    real_t *getEvent(const size_t &index) { return ((*this)[index]); }
    const real_t *getEvent(const size_t &index) const { return (const real_t *)((*this)[index]); }
    std::vector<Event>::const_iterator begin() const { return m_data.cbegin(); }
    std::vector<Event>::const_iterator end() const { return m_data.cend(); }
    const Event &operator[](const size_t &pos) const { return m_data[pos]; }
    EventType eventType() const { return m_eventType; }
    const Event &at(const size_t &pos) const { return m_data[pos]; }
    size_t size() const { return m_data.size(); }
    size_t aligned_size() const { return m_data.size(); }
    size_t nBlocks() const { return m_data.size(); }
    double integral() const;
    double *block(const unsigned pos) { return m_data[pos].address(); }
    real_t weight(const size_t &pos) const { return m_data[pos].weight(); }
    real_t genPDF(const size_t &pos) const { return m_data[pos].genPdf(); }
    unsigned key(const std::string &key) const {
      auto it = m_extensions.find(key);
      if(it == m_extensions.end()) return m_data[0].size() - 1;
      return it->second;
    }
    void reserve(const size_t &size);
    void resize(const size_t &size);
    void push_back(const Event &evt);
    void emplace_back(const Event &evt);
    void setEventType(const EventType &type) { m_eventType = type; }
    void add(const EventList &evts);
    void loadFromTree(TTree *tree, const ArgumentPack &args);
    void loadFromFile(const std::string &fname, const ArgumentPack &args);
    void clear();
    void extend(const std::string &key, unsigned pos) { m_extensions[key] = pos; }
 
    void setWeight(const unsigned int &pos, const double &w, const double &g = +1) {
      m_data[pos].setWeight(w);
      m_data[pos].setGenPdf(g);
    }
    void setGenPDF(const unsigned int &pos, const double &g) { m_data[pos].setGenPdf(g); }
    void erase(const std::vector<Event>::iterator &begin, const std::vector<Event>::iterator &end);
 
    TTree *tree(const std::string &name, const std::vector<std::string> &extraBranches = {}) const;
 
    TH1D *makeProjection(const Projection &projection, const ArgumentPack &args = ArgumentPack()) const;
    TH2D *makeProjection(const Projection2D &projection, const ArgumentPack &args = ArgumentPack()) const;
    std::vector<TH1D *> makeProjections(const std::vector<Projection> &projections, const ArgumentPack &args);
 
    template <class... ARGS> std::vector<TH1D *> makeDefaultProjections(const ARGS &...args) {
      auto argPack = ArgumentPack(args...);
      size_t nBins = argPack.getArg<PlotOptions::Bins>(100);
      auto proj = eventType().defaultProjections(nBins);
      return makeProjections(proj, argPack);
    }
 
    template <typename... ARGS> std::vector<TH1D *> makeProjections(const std::vector<Projection> &projections, const ARGS &...args) {
      return makeProjections(projections, ArgumentPack(args...));
    }
 
    template <typename... ARGS, typename = std::enable_if_t<!std::is_same<zeroType<ARGS...>, ArgumentPack>::value>>
    TH1D *makeProjection(const Projection &projection, const ARGS &...args) const {
      return makeProjection(projection, ArgumentPack(args...));
    }
 
    template <typename... ARGS, typename = std::enable_if_t<!std::is_same<zeroType<ARGS...>, ArgumentPack>::value>>
    TH2D *makeProjection(const Projection2D &projection, const ARGS &...args) {
      return makeProjection(projection, ArgumentPack(args...));
    }
 
    template <typename functor> EventList &transform(functor &&fcn) {
      for(auto &event : m_data) fcn(event);
      return *this;
    }
 
    template <typename functor> void filter(functor &&fcn) {
      unsigned currentSize = size();
      m_data.erase(std::remove_if(m_data.begin(), m_data.end(), fcn), m_data.end());
      INFO("Filter retains " << size() << " / " << currentSize << " events");
    }
 
    template <typename functor> unsigned count(functor &&fcn) const { return std::count_if(std::begin(*this), std::end(*this), fcn); }
  };
  DECLARE_ARGUMENT(Branches, std::vector<std::string>);      
  DECLARE_ARGUMENT(ExtraBranches, std::vector<std::string>); 
  DECLARE_ARGUMENT(IdBranches,
                   std::vector<std::string>); 
  DECLARE_ARGUMENT(EntryList, std::vector<size_t>);
  DECLARE_ARGUMENT(GetGenPdf, bool);
  DECLARE_ARGUMENT(Filter, std::string);
  DECLARE_ARGUMENT(WeightBranch, std::string);
  DECLARE_ARGUMENT(ApplySym, bool);
  DECLARE_ARGUMENT(WeightFunction, std::function<double(const Event &)>);
  DECLARE_ARGUMENT(InputUnits, AmpGen::Units);
} // namespace AmpGen
 
#endif