TagAndProbe/plugins/ProbeTreeProducer.cc

0001 // -*- C++ -*-
0002 //
0003 // Package:    ProbeTreeProducer
0004 // Class:      ProbeTreeProducer
0005 //
0006 /**\class ProbeTreeProducer ProbeTreeProducer.cc
0007
0008  Description: TTree producer based on input probe parameters
0009
0010  Implementation:
0011      <Notes on implementation>
0012 */
0013
0014 #include <memory>
0015 #include <cctype>
0016 #include "FWCore/Framework/interface/Frameworkfwd.h"
0017 #include "FWCore/Framework/interface/one/EDFilter.h"
0018 #include "FWCore/Framework/interface/Event.h"
0019 #include "FWCore/Framework/interface/MakerMacros.h"
0020 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0021 #include "FWCore/Utilities/interface/InputTag.h"
0022 #include "DataFormats/Candidate/interface/CandidateFwd.h"
0023 #include "DataFormats/Candidate/interface/Candidate.h"
0024 #include "PhysicsTools/TagAndProbe/interface/BaseTreeFiller.h"
0025 #include <set>
0026
0027 class ProbeTreeProducer : public edm::one::EDFilter<> {
0028 public:
0029   explicit ProbeTreeProducer(const edm::ParameterSet&);
0030   ~ProbeTreeProducer() override;
0031
0032 private:
0033   bool filter(edm::Event&, const edm::EventSetup&) override;
0034   void endJob() override;
0035
0036   /// InputTag to the collection of all probes
0037   edm::EDGetTokenT<reco::CandidateView> probesToken_;
0038
0039   /// The selector object
0040   StringCutObjectSelector<reco::Candidate, true> cut_;
0041
0042   /// Specifies whether this module should filter
0043   bool filter_;
0044
0045   /// Name of the reco::Candidate function used for sorting
0046   std::string sortDescendingBy_;
0047
0048   /// The StringObjectFunction itself
0049   StringObjectFunction<reco::Candidate, true> sortFunction_;
0050
0051   /// The number of first probes used to fill the tree
0052   int32_t maxProbes_;
0053
0054   /// The object that actually computes variables and fills the tree for the probe
0055   std::unique_ptr<tnp::BaseTreeFiller> probeFiller_;
0056 };
0057
0058 ProbeTreeProducer::ProbeTreeProducer(const edm::ParameterSet& iConfig)
0059     : probesToken_(consumes<reco::CandidateView>(iConfig.getParameter<edm::InputTag>("src"))),
0060       cut_(iConfig.existsAs<std::string>("cut") ? iConfig.getParameter<std::string>("cut") : ""),
0061       filter_(iConfig.existsAs<bool>("filter") ? iConfig.getParameter<bool>("filter") : false),
0062       sortDescendingBy_(iConfig.existsAs<std::string>("sortDescendingBy")
0063                             ? iConfig.getParameter<std::string>("sortDescendingBy")
0064                             : ""),
0065       sortFunction_(!sortDescendingBy_.empty() ? sortDescendingBy_ : "pt"),  //need to pass a valid default
0066       maxProbes_(iConfig.existsAs<int32_t>("maxProbes") ? iConfig.getParameter<int32_t>("maxProbes") : -1),
0067       probeFiller_(new tnp::BaseTreeFiller("probe_tree", iConfig, consumesCollector())) {}
0068
0069 ProbeTreeProducer::~ProbeTreeProducer() {}
0070
0071 bool ProbeTreeProducer::filter(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0072   bool result = !filter_;
0073   edm::Handle<reco::CandidateView> probes;
0074   iEvent.getByToken(probesToken_, probes);
0075   if (!probes.isValid())
0076     return result;
0077   probeFiller_->init(iEvent);
0078   // select probes and calculate the sorting value
0079   typedef std::pair<reco::CandidateBaseRef, double> Pair;
0080   std::vector<Pair> selectedProbes;
0081   for (size_t i = 0; i < probes->size(); ++i) {
0082     const reco::CandidateBaseRef& probe = probes->refAt(i);
0083     if (cut_(*probe)) {
0084       selectedProbes.push_back(Pair(probe, sortFunction_(*probe)));
0085     }
0086   }
0087   // sort only if a function was provided
0088   if (!sortDescendingBy_.empty())
0089     sort(selectedProbes.begin(), selectedProbes.end(), [&](auto& arg1, auto& arg2) {
0090       return arg1.second > arg2.second;
0091     });
0092   // fill the first maxProbes_ into the tree
0093   for (size_t i = 0; i < (maxProbes_ < 0 ? selectedProbes.size() : std::min((size_t)maxProbes_, selectedProbes.size()));
0094        ++i) {
0095     probeFiller_->fill(selectedProbes[i].first);
0096     result = true;
0097   }
0098   return result;
0099 }
0100
0101 void ProbeTreeProducer::endJob() {
0102   // ask to write the current PSet info into the TTree header
0103   probeFiller_->writeProvenance(edm::getProcessParameterSetContainingModule(moduleDescription()));
0104 }
0105
0106 //define this as a plug-in
0107 DEFINE_FWK_MODULE(ProbeTreeProducer);