AnalysisHandler.hh

// -*- C++ -*-
#ifndef RIVET_RivetHandler_HH
#define RIVET_RivetHandler_HH
 
#include "Rivet/Config/RivetCommon.hh"
#include "Rivet/Particle.hh"
#include "Rivet/AnalysisLoader.hh"
#include "Rivet/Tools/RivetYODA.hh"
#include "Rivet/Tools/Utils.hh"
#include "Rivet/ProjectionHandler.hh"
#include "YODA/ReaderYODA.h"
 
#include <fstream>
#include <unordered_map>
 
namespace Rivet {
 
 
  // Forward declaration and smart pointer for Analysis
  class Analysis;
  using AnaHandle = std::shared_ptr<Analysis>;
 

  class AnalysisHandler {
 
    using TypeHandlePtr = std::shared_ptr<TypeBaseHandle>;
    using TypeRegister = std::unordered_map<string, TypeHandlePtr>;
    using TypeRegisterItr = typename TypeRegister::const_iterator;
 
  public:
 
    using Annotations = std::map<std::string, std::string>;

    AnalysisHandler();

    AnalysisHandler(const AnalysisHandler&) = delete;

    AnalysisHandler& operator=(const AnalysisHandler&) = delete;

    ~AnalysisHandler();
 


    size_t numEvents() const {
      const double N = _eventCounter.get()->persistent(defaultWeightIndex())->numEntries();
      return  size_t(N + 0.5 - (N<0)); // round to nearest integer
    }

    double effNumEvents() const {
      if ((bool)_eventCounter) { return _eventCounter->effNumEntries(); }
      return _eventCounter.get()->persistent(defaultWeightIndex())->effNumEntries();
    }

    double sumW() const {
      if ((bool)_eventCounter) { return _eventCounter->sumW(); }
      return _eventCounter.get()->persistent(defaultWeightIndex())->sumW();
    }

    double sumW2() const {
      if ((bool)_eventCounter) { return _eventCounter->sumW2(); }
      return _eventCounter.get()->persistent(defaultWeightIndex())->sumW2();
    }

 


    const vector<string>& weightNames() const { return _weightNames; }

    size_t numWeights() const { return _weightNames.size(); }

    bool haveNamedWeights() const;

    void setWeightNames(const GenEvent& ge);

    void setWeightNames(const vector<string>& weightNames);

    size_t defaultWeightIndex() const { return _rivetDefaultWeightIdx; }

    vector<double> weightSumWs() const;

    void setWeightCap(const double maxWeight) { _weightCap = maxWeight; }

    void setNominalWeightName(const std::string& name) { _nominalWeightName = name; }

    void skipMultiWeights(bool skip=false) { _skipMultiWeights = skip; }

    void matchWeightNames(const std::string& patterns) { _matchWeightNames = patterns; }

    void unmatchWeightNames(const std::string& patterns) { _unmatchWeightNames = patterns; }

    void setNLOSmearing(double frac) { _NLOSmearing = frac; }

 


    Estimate0DPtr crossSection() const { return _xs; }

    void setCrossSection(const vector<pair<double,double>>& xsecs, bool isUserSupplied = false);

    void setCrossSection(const pair<double, double>& xsec, bool isUserSupplied=false);

    void setCrossSection(double xsec, double xsecerr, bool isUserSupplied=false) {
      setCrossSection({xsec, xsecerr}, isUserSupplied);
    }

    void updateCrossSection();

    void notifyEndOfFile() { _isEndOfFile = true; }

    double nominalCrossSection() const;

    double nominalCrossSectionError() const;

 


    AnalysisHandler& setRunBeams(const ParticlePair& beams);

    const ParticlePair& runBeams() const { return _beams; }

    PdgIdPair runBeamIDs() const;

    pair<double,double> runBeamEnergies() const;

    double runSqrtS() const;

    void setCheckBeams(bool check=true) { _checkBeams = check; }

    // void setCheckConsistency(bool check=true) { _checkConsistency = check; }
    // Check event consistency with the run, usually determined from the first event
    // bool consistentWithRun(Event& event) {

 

 
    // Get all the annotation names
    std::vector<std::string> annotations() const {
      return _beaminfo->annotations();
    }

    bool hasAnnotation(const std::string& name) const {
      return _beaminfo->hasAnnotation(name);
    }

    const std::string& annotation(const std::string& name) const {
      return _beaminfo->annotation(name);
    }

    const std::string& annotation(const std::string& name, const std::string& defaultreturn) const {
      return _beaminfo->annotation(name, defaultreturn);
    }

    template <typename T>
    const T annotation(const std::string& name) const {
      return _beaminfo->annotation<T>(name);
    }

    template <typename T>
    const T annotation(const std::string& name, T&& defaultreturn) const {
      return _beaminfo->annotation<T>(name, std::forward<T>(defaultreturn));
    }

    template <typename T>
    void setAnnotation(const std::string& name, T&& value) {
      _beaminfo->setAnnotation<T>(name, std::forward<T>(value));
    }
 

    void setAnnotations(const Annotations& anns) {
      _beaminfo->setAnnotations(anns);
    }

    void rmAnnotation(const std::string& name) {
      _beaminfo->rmAnnotation(name);
    }
 

    void clearAnnotations() {
      _beaminfo->clearAnnotations();
    }

 


    template<typename T>
    void registerType() {
      const std::string name = T().type();
      const TypeRegisterItr& res = _register.find(name);
      if (res == _register.end()) {
        _register[name] = make_shared<TypeHandle<T>>();
      }
      _reader.registerType<T>(); // also let YODA know
    }

    bool copyAO(YODA::AnalysisObjectPtr src, YODA::AnalysisObjectPtr dst, const double scale=1.0);

    bool addAO(YODA::AnalysisObjectPtr src, YODA::AnalysisObjectPtr& dst, const double scale);

 


    std::vector<std::string> analysisNames() const;

    std::vector<std::string> stdAnalysisNames() const;

    const std::map<std::string, AnaHandle>& analysesMap() const {
      return _analyses;
    }

    std::vector<AnaHandle> analyses() const {
      std::vector<AnaHandle> rtn;
      rtn.reserve(_analyses.size());
      for (const auto& apair : _analyses) rtn.push_back(apair.second);
      return rtn;
    }

    AnaHandle analysis(const std::string& analysisname) {
      if ( _analyses.find(analysisname) == _analyses.end() )
        throw LookupError("No analysis named '" + analysisname + "' registered in AnalysisHandler");
      try {
        return _analyses[analysisname];
      } catch (...) {
        throw LookupError("No analysis named '" + analysisname + "' registered in AnalysisHandler");
      }
    }

    AnalysisHandler& addAnalysis(Analysis* analysis);

    AnalysisHandler& addAnalysis(const std::string& analysisname);

    AnalysisHandler& addAnalysis(const std::string& analysisname, std::map<string, string> pars);

    AnalysisHandler& addAnalyses(const std::vector<std::string>& analysisnames);
 

    AnalysisHandler& removeAnalysis(const std::string& analysisname);

    AnalysisHandler& removeAnalyses(const std::vector<std::string>& analysisnames);

 


    void init(const GenEvent& event);

    void analyze(GenEvent& event);

    void analyze(GenEvent* event);

    void finalize();

 


    void collapseEventGroup();

    void readData(std::istream& istr, const string& fmt, bool preload = true);

    void readData(const std::string& filename, bool preload = true);

    vector<YODA::AnalysisObjectPtr> getYodaAOs(const bool includeraw=false, const bool mkinert=true) const;

    vector<YODA::AnalysisObjectPtr> getRawAOs() const;

    vector<std::string> getRawAOPaths() const;

    const YODA::AnalysisObjectPtr getPreload(const string& path) const {
      auto it = _preloads.find(path);
      if ( it == _preloads.end() ) return nullptr;
      return it->second;
    }

    void writeData(std::ostream& ostr, const string& fmt) const;

    void writeData(const string& filename) const;

    void setFinalizePeriod(const string& dumpfile, int period) {
      _dumpPeriod = period;
      _dumpFile = dumpfile;
    }

    void setNoFinalizePeriod() {
      setFinalizePeriod("DUMMY", -1);
    }

    void setBootstrapFilename(const string& filename) {
      _bootstrapfilename = filename;
    }

    vector<pair<string,size_t>> fillLayout() const;

    vector<bool> fillOutcomes() const;

    vector<double> fillFractions() const;

 
    void mergeYODAs(const vector<string>& aofiles,
                    const vector<string>& delopts=vector<string>(),
                    const vector<string>& addopts=vector<string>(),
                    const vector<string>& matches=vector<string>(),
                    const vector<string>& unmatches=vector<string>(),
                    const bool equiv=false,
                    const bool reentrantOnly = true,
                    const bool checkbeams = true);

    void merge(AnalysisHandler &other);

    void loadAOs(const vector<string>& aoPaths, const vector<double>& aoData);


 
    vector<double> serializeContent(bool fixed_length = false) {
      if (!_initialised)
        throw Error("AnalysisHandler has not been initialised!");
 
      collapseEventGroup();
 
      // Loop over raw AOs and work out the size of the content data
      const vector<YODA::AnalysisObjectPtr> raos = getRawAOs();
      size_t total = 0;
      for (size_t i = 0; i < raos.size(); ++i) {
        total += raos[i]->lengthContent(fixed_length)+1;
      }
      total += _beaminfo->numBins()+1;
 
      // Loop over raw AOs and retrieve the content data
      std::vector<double> data; // serialized data vector
      data.reserve(total); // pre-allocate enough memory
      // Add beam IDs
      data.push_back(_beaminfo->numBins());
      for (const string& beamID : _beaminfo->xEdges()) {
        data.push_back(_beamInfoLabelToID(beamID));
      }
      // Add raw YODA AO content
      for (size_t i = 0; i < raos.size(); ++i) {
        vector<double> tmp = raos[i]->serializeContent(fixed_length);
        data.push_back(tmp.size()); // length of the AO
        data.insert(std::end(data),
                    std::make_move_iterator(std::begin(tmp)),
                    std::make_move_iterator(std::end(tmp)));
      }
      return data;
    }
 
    void deserializeContent(const vector<double>& data, size_t nprocs = 0) {
      if (!_initialised)
        throw Error("AnalysisHandler has not been initialised!");
 
      collapseEventGroup();
 
      // get Rivet AOs for access to raw AO pointers
      vector<MultiplexAOPtr> raos = getRivetAOs();
 
 
      // beam info first
      size_t iAO = 0, iW = 0, nBeams = data[0], offset = 1;
      if (nprocs)  nBeams /= nprocs;
      const auto itr = data.cbegin();
      // set beam IDs
      vector<int> edges{itr+offset, itr+offset+nBeams};
      vector<string> labels; labels.reserve(edges.size());
      size_t id = 0;
      for (int edge : edges) {
        if (nprocs >= 2)  edge /= nprocs;
        labels.push_back(_mkBeamInfoLabel(++id, edge));
      }
 
      _beaminfo = make_shared<YODA::BinnedEstimate<string>>(labels, "/TMP/_BEAMPZ");
      offset += nBeams;
      // set beam momenta
      size_t beamLen = *(itr + offset); ++offset;
      if (nprocs)  beamLen /= nprocs;
      std::vector<double> energies{itr+offset, itr+offset+beamLen};
      if (nprocs >= 2) {
        for (double& e : energies) { e /= nprocs; }
      }
      _beaminfo->deserializeContent(energies);
      for (auto& b : _beaminfo->bins(true))  b.rmErrs();
      offset += beamLen;
 
      // then the multiweighted AOs
      while (offset < data.size()) {
        if (iW < numWeights())  raos[iAO].get()->setActiveWeightIdx(iW);
        else {
          raos[iAO].get()->unsetActiveWeight();
          iW = 0; ++iAO; // move on to next AO
          raos[iAO].get()->setActiveWeightIdx(iW);
        }
 
        // obtain content length and set content iterators
        size_t aoLen = *(itr + offset); ++offset;
        if (nprocs)  aoLen /= nprocs;
        auto first = itr + offset;
        auto last = first + aoLen;
        // load data into AO
        raos[iAO].get()->activeAO()->deserializeContent(std::vector<double>{first, last});
 
        ++iW; offset += aoLen; // increment offset
      }
      raos[iAO].get()->unsetActiveWeight();
      // Reset cross-section bookkeeping
      _ntrials = 0.0;
      _fileCounter = CounterPtr(weightNames(), Counter("_FILECOUNT"));
      _xserr = CounterPtr(weightNames(), Counter("XSECERR"));
      if (nprocs >= 2) {
        for (size_t iW = 0; iW < numWeights(); ++iW) {
          *_fileCounter.get()->persistent(iW) = *_eventCounter.get()->persistent(iW);
          _xs.get()->persistent(iW)->scale(1.0/nprocs);
        }
      }
    }

 


    enum class Stage { OTHER, INIT, FINALIZE };

    Stage stage() const { return _stage; }

 
  private:


    Log& getLog() const;

    vector<MultiplexAOPtr> getRivetAOs() const;

    void stripOptions(YODA::AnalysisObjectPtr ao, const vector<string>& delopts) const;

    void mergeAOS(map<string, YODA::AnalysisObjectPtr> &allaos,
                  const map<string, YODA::AnalysisObjectPtr> &newaos,
                  map<string, std::array<double,4>> &allxsecs,
                  const vector<string>& delopts=vector<string>(),
                  const vector<string>& optAnas=vector<string>(),
                  const vector<string>& optKeys=vector<string>(),
                  const vector<string>& optVals=vector<string>(),
                  const bool equiv=false,
                  const bool overwrite_xsec = false,
                  const double user_xsec = 1.0);
 

    void loadAOs(const map<string, YODA::AnalysisObjectPtr>& allAOs,
                 const bool unscale = false, const bool reentrantOnly = true);

    void _setRunBeamInfo(const ParticlePair& beams);

    string _mkBeamInfoLabel(size_t n, PdgId id) {
      return "BEAM"+std::to_string(n)+"("+std::to_string(id)+")";
    }

    PdgId _beamInfoLabelToID(const string& label) {
      size_t pos = label.find("(");
      string beamID = label.substr(pos+1, label.size()-pos-2);
      return std::stoi(beamID);
    }
 

 
 
  private:

    Stage _stage = Stage::OTHER;

    std::map<std::string, AnaHandle> _analyses;

    map<string,YODA::AnalysisObjectPtr> _preloads;

    vector<YODA::AnalysisObjectPtr> _finalizedAOs;

    template<typename... Args>
    void registerDefaultTypes();

    TypeRegister _register;

    YODA::Reader& _reader = YODA::ReaderYODA::create();
 


    std::vector<std::string> _weightNames;
    std::vector<std::valarray<double> > _subEventWeights;
    //size_t _numWeightTypes; // always == WeightVector.size()

    std::vector<size_t> _weightIndices;

    CounterPtr _eventCounter;

    Estimate0DPtr _xs;

    CounterPtr _xserr;

    YODA::BinnedEstimatePtr<string> _beaminfo;

    double _ntrials;

    CounterPtr _fileCounter;

    bool _isEndOfFile;

    bool _subeventWarning;

    std::pair<double,double> _userxs;

    ParticlePair _beams;

    bool _initialised;

    bool _checkBeams{true};

    bool _ignoreEventGroups{false};

    bool _skipMultiWeights{false};

    std::string _matchWeightNames;

    std::string _unmatchWeightNames;

    std::string _nominalWeightName;

    double _weightCap;

    double _NLOSmearing;

    int _eventNumber;

    size_t _defaultWeightIdx;

    size_t _rivetDefaultWeightIdx;

    int _customDefaultWeightIdx;

    int _dumpPeriod;

    string _dumpFile;

    bool _dumping;

    ofstream _fbootstrap;

    std::string _bootstrapfilename;

    ProjectionHandler _projHandler;

 
  };
 
 
}
 
#endif