JetUtils.hh

#ifndef RIVET_JETUTILS_HH
#define RIVET_JETUTILS_HH
 
#include "Rivet/Jet.hh"
#include "Rivet/Tools/ParticleBaseUtils.hh"
 
namespace Rivet {
 


 
  inline PseudoJets mkPseudoJets(const Particles& ps) {
    PseudoJets rtn; rtn.reserve(ps.size());
    for (const Particle& p : ps) rtn.push_back(p.pseudojet());
    return rtn;
  }
 
  inline PseudoJets mkPseudoJets(const Jets& js) {
    PseudoJets rtn; rtn.reserve(js.size());
    for (const Jet& j : js) rtn.push_back(j.pseudojet());
    return rtn;
  }
 
  inline Jets mkJets(const PseudoJets& pjs) {
    Jets rtn; rtn.reserve(pjs.size());
    for (const PseudoJet& pj : pjs) rtn.push_back(pj);
    return rtn;
  }

 


  using JetSelector = function<bool(const Jet&)>;
  using JetSorter = function<bool(const Jet&, const Jet&)>;
 

  struct BoolJetFunctor {
    virtual bool operator()(const Jet& p) const = 0;
    virtual ~BoolJetFunctor() {}
  };
 

  struct BoolJetAND : public BoolJetFunctor {
    BoolJetAND(const std::vector<JetSelector>& sels) : selectors(sels) {}
    BoolJetAND(const JetSelector& a, const JetSelector& b) : selectors({a,b}) {}
    BoolJetAND(const JetSelector& a, const JetSelector& b, const JetSelector& c) : selectors({a,b,c}) {}
    bool operator()(const Jet& j) const {
      for (const JetSelector& sel : selectors) if (!sel(j)) return false;
      return true;
    }
    std::vector<JetSelector> selectors;
  };

  inline BoolJetAND operator && (const JetSelector& a, const JetSelector& b) {
    return BoolJetAND(a, b);
  }
 

  struct BoolJetOR : public BoolJetFunctor {
    BoolJetOR(const std::vector<JetSelector>& sels) : selectors(sels) {}
    BoolJetOR(const JetSelector& a, const JetSelector& b) : selectors({a,b}) {}
    BoolJetOR(const JetSelector& a, const JetSelector& b, const JetSelector& c) : selectors({a,b,c}) {}
    bool operator()(const Jet& j) const {
      for (const JetSelector& sel : selectors) if (sel(j)) return true;
      return false;
    }
    std::vector<JetSelector> selectors;
  };

  inline BoolJetOR operator || (const JetSelector& a, const JetSelector& b) {
    return BoolJetOR(a, b);
  }
 

  struct BoolJetNOT : public BoolJetFunctor {
    BoolJetNOT(const JetSelector& sel) : selector(sel) {}
    bool operator()(const Jet& j) const { return !selector(j); }
    JetSelector selector;
  };

  inline BoolJetNOT operator ! (const JetSelector& a) {
    return BoolJetNOT(a);
  }
 
 

  struct HasBTag : BoolJetFunctor {
    HasBTag(const Cut& c=Cuts::open(), double dR=-1) : cut(c), deltaR(dR) {}
    // HasBTag(const std::function<bool(const Jet& j)>& f) : selector(f) {}
    bool operator() (const Jet& j) const { return j.bTagged(cut, deltaR); }
    // const std::function<bool(const Jet& j)> selector;
    const Cut cut;
    const double deltaR;
  };
  using hasBTag = HasBTag;

  struct HasCTag : BoolJetFunctor {
    HasCTag(const Cut& c=Cuts::open(), double dR=-1) : cut(c), deltaR(dR) {}
    // HasCTag(const std::function<bool(const Jet& j)>& f) : selector(f) {}
    bool operator() (const Jet& j) const { return j.cTagged(cut, deltaR); }
    // const std::function<bool(const Jet& j)> selector;
    const Cut cut;
    const double deltaR;
  };
  using hasCTag = HasCTag;

  struct HasTauTag : BoolJetFunctor {
    HasTauTag(const Cut& c=Cuts::open(), double dR=-1) : cut(c), deltaR(dR) {}
    // HasTauTag(const std::function<bool(const Jet& j)>& f) : selector(f) {}
    bool operator() (const Jet& j) const { return j.tauTagged(cut, deltaR); }
    // const std::function<bool(const Jet& j)> selector;
    const Cut cut;
    const double deltaR;
  };
  using hasTauTag = HasTauTag;

  struct HasNoTag : BoolJetFunctor {
    HasNoTag(const Cut& c=Cuts::open(), double dR=-1, bool quarktagsonly=false) : cut(c), deltaR(dR), qtagsonly(quarktagsonly) {}
    HasNoTag(const Cut& c=Cuts::open(), bool quarktagsonly=false) : HasNoTag(c, -1, quarktagsonly) {}
    // HasNoTag(const std::function<bool(const Jet& j)>& f) : selector(f) {}
    bool operator() (const Jet& j) const { return !j.bTagged(cut, deltaR) && !j.cTagged(cut, deltaR) && (qtagsonly || !j.tauTagged(cut, deltaR)); }
    // const std::function<bool(const Jet& j)> selector;
    const Cut cut;
    const double deltaR;
    bool qtagsonly;
  };
  using hasNoTag = HasNoTag;

 


  Jets& iselect(Jets& jets, const Cut& c);
 

  inline Jets select(const Jets& jets, const Cut& c) {
    Jets rtn = jets;
    return iselect(rtn, c);
  }
 

  inline Jets select(const Jets& jets, const Cut& c, Jets& out) {
    out = select(jets, c);
    return out;
  }
 
 

  Jets& idiscard(Jets& jets, const Cut& c);
 

  inline Jets discard(const Jets& jets, const Cut& c) {
    Jets rtn = jets;
    return idiscard(rtn, c);
  }
 

  inline Jets discard(const Jets& jets, const Cut& c, Jets& out) {
    out = discard(jets, c);
    return out;
  }

 


  Jet& itrimJetsFrac(Jet &jet, const double frac);

  template <
    typename CONTAINER,
    typename = std::enable_if_t<
      is_citerable_v<CONTAINER>,
      Jet
    >
  >
  CONTAINER& itrimJetsFrac(CONTAINER &jets, const double frac, const JetSorter &sortFunc=cmpMomByPt){
    for ( Jet &jet : jets ) itrimJetsFrac(jet, frac);
    isortBy(jets, sortFunc);
    return jets;
  }
 
  template <typename T, typename U>
  std::map<T, U>& itrimJetsFrac(std::map<T, U> &jetMap, const double frac, const JetSorter &sortFunc=cmpMomByPt){
    for ( auto &item : jetMap ) itrimJetsFrac(item.second, frac, sortFunc);
    return jetMap;
  }

  inline Jet trimJetsFrac(const Jet &jet, double frac){
    Jet rtn = jet;
    return itrimJetsFrac(rtn, frac);
  }

  template <
    typename... Args, typename CONTAINER,
    typename = std::enable_if_t<
      is_citerable_v<CONTAINER>,
      Jet
    >
  >
  CONTAINER trimJetsFrac(const CONTAINER &jets, Args&&... args){
    CONTAINER rtn = jets;
    return itrimJetsFrac(rtn, std::forward<Args>(args)...);
  }
 
  template <typename T, typename U, typename... Args>
  std::map<T, U> trimJetsFrac(const std::map<T, U> &jetMap, Args&&... args){
    std::map<T, U> rtn = jetMap;
    return itrimJetsFrac(rtn, std::forward<Args>(args)...);
  }



  PseudoJet& ifilterPseudoJets(PseudoJet &pj, const fastjet::Filter &filter);

  template <
    typename CONTAINER,
    typename = std::enable_if_t<
      is_citerable_v<CONTAINER>,
      PseudoJet
    >
  >
  CONTAINER& ifilterPseudoJets(CONTAINER &pjs, const fastjet::Filter &filter, const JetSorter &sortFunc=cmpMomByPt){
    for ( PseudoJet& pj : pjs ) ifilterPseudoJets(pj, filter);
    isortBy(pjs, sortFunc);
    return pjs;
  }

  template <typename T, typename U, typename... Args>
  std::map<T, U>& ifilterPseudoJets(std::map<T, U> &pjMap, Args&&... args){
    for ( auto &item : pjMap ) ifilterPseudoJets(item.second, std::forward<Args>(args)...);
    return pjMap;
  }

  inline PseudoJet filterPseudoJets(const PseudoJet &pj, const fastjet::Filter &filter){
    PseudoJet rtn = pj;
    return ifilterPseudoJets(rtn, filter);
  }

  template <
    typename... Args, typename CONTAINER,
    typename = std::enable_if_t<
      is_citerable_v<CONTAINER>,
      PseudoJet
    >
  >
  CONTAINER filterPseudoJets(const CONTAINER &pjs, Args&&... args){
    CONTAINER rtn = pjs;
    return ifilterPseudoJets(rtn, std::forward<Args>(args)...);
  }

  template <typename T, typename U, typename... Args>
  std::map<T, U> filterPseudoJets(const std::map<T, U> &pjMap, Args&&... args){
    std::map<T, U> rtn = pjMap;
    return ifilterPseudoJets(rtn, std::forward<Args>(args)...);
  }

 
  // Import Kin namespace into Rivet
  using namespace Kin;
 

 
}
 
#endif