ProbeTreeProducer.cc

Go to the documentation of this file.

// -*- C++ -*-
//
// Package:    ProbeTreeProducer
// Class:      ProbeTreeProducer
//
#include <memory>
#include <cctype>
#include "boost/bind.hpp"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/EDFilter.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "DataFormats/Candidate/interface/CandidateFwd.h"
#include "DataFormats/Candidate/interface/Candidate.h"
#include "PhysicsTools/TagAndProbe/interface/BaseTreeFiller.h"
#include <set>

class ProbeTreeProducer : public edm::EDFilter {
  public:
    explicit ProbeTreeProducer(const edm::ParameterSet&);
    ~ProbeTreeProducer() override;

  private:
    bool filter(edm::Event&, const edm::EventSetup&) override;
    void endJob() override;

    edm::EDGetTokenT<reco::CandidateView> probesToken_;

    StringCutObjectSelector<reco::Candidate, true> cut_;

    bool filter_;

    std::string sortDescendingBy_;

    StringObjectFunction<reco::Candidate, true> sortFunction_;

    int32_t maxProbes_;

    std::unique_ptr<tnp::BaseTreeFiller> probeFiller_;
};

ProbeTreeProducer::ProbeTreeProducer(const edm::ParameterSet& iConfig) :
  probesToken_(consumes<reco::CandidateView>(iConfig.getParameter<edm::InputTag>("src"))),
  cut_(iConfig.existsAs<std::string>("cut") ? iConfig.getParameter<std::string>("cut") : ""),
  filter_(iConfig.existsAs<bool>("filter") ? iConfig.getParameter<bool>("filter") : false),
  sortDescendingBy_(iConfig.existsAs<std::string>("sortDescendingBy") ? iConfig.getParameter<std::string>("sortDescendingBy") : ""),
  sortFunction_(!sortDescendingBy_.empty() ? sortDescendingBy_ : "pt"), //need to pass a valid default
  maxProbes_(iConfig.existsAs<int32_t>("maxProbes") ? iConfig.getParameter<int32_t>("maxProbes") : -1),
  probeFiller_(new tnp::BaseTreeFiller("probe_tree", iConfig, consumesCollector()))
{
}

ProbeTreeProducer::~ProbeTreeProducer(){
}

bool ProbeTreeProducer::filter(edm::Event& iEvent, const edm::EventSetup& iSetup){
  bool result = !filter_;
  edm::Handle<reco::CandidateView> probes;
  iEvent.getByToken(probesToken_, probes);
  if(!probes.isValid()) return result;
  probeFiller_->init(iEvent);
  // select probes and calculate the sorting value
  typedef std::pair<reco::CandidateBaseRef, double> Pair;
  std::vector<Pair> selectedProbes;
  for (size_t i = 0; i < probes->size(); ++i){
    const reco::CandidateBaseRef &probe = probes->refAt(i);
    if(cut_(*probe)){
      selectedProbes.push_back(Pair(probe, sortFunction_(*probe)));
    }
  }
  // sort only if a function was provided
  if(!sortDescendingBy_.empty()) sort(selectedProbes.begin(), selectedProbes.end(), boost::bind(&Pair::second, _1) > boost::bind(&Pair::second, _2));
  // fill the first maxProbes_ into the tree
  for (size_t i = 0; i < (maxProbes_<0 ? selectedProbes.size() : std::min((size_t)maxProbes_, selectedProbes.size())); ++i){
    probeFiller_->fill(selectedProbes[i].first);
    result = true;
  }
  return result;
}

void ProbeTreeProducer::endJob(){
    // ask to write the current PSet info into the TTree header
    probeFiller_->writeProvenance(edm::getProcessParameterSetContainingModule(moduleDescription()));
}

//define this as a plug-in
DEFINE_FWK_MODULE(ProbeTreeProducer);