PatAlgos/plugins/PATTauIDEmbedder.cc

0001 #include "FWCore/Framework/interface/stream/EDProducer.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/Framework/interface/EventSetup.h"
0004 #include "FWCore/Utilities/interface/InputTag.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006 #include "DataFormats/PatCandidates/interface/PATTauDiscriminator.h"
0007 #include "DataFormats/TauReco/interface/TauDiscriminatorContainer.h"
0008 #include "DataFormats/PatCandidates/interface/Tau.h"
0009 #include "FWCore/Utilities/interface/transform.h"
0010
0011 class PATTauIDEmbedder : public edm::stream::EDProducer<> {
0012 public:
0013   explicit PATTauIDEmbedder(const edm::ParameterSet&);
0014   ~PATTauIDEmbedder() override {}
0015
0016   void produce(edm::Event&, const edm::EventSetup&) override;
0017
0018 private:
0019   //--- configuration parameters
0020   edm::EDGetTokenT<pat::TauCollection> src_;
0021   typedef std::pair<std::string, edm::InputTag> NameTag;
0022   typedef std::pair<std::string, int> NameWPIdx;
0023   //to save input module tag and corresponding pairs <working point name for the output tree, WP index in the input ID container>
0024   typedef std::pair<edm::InputTag, std::vector<NameWPIdx> > IDContainerData;
0025   std::vector<NameTag> tauIDSrcs_;
0026   std::vector<std::vector<NameWPIdx> > tauIDSrcContainers_;
0027   std::vector<edm::EDGetTokenT<pat::PATTauDiscriminator> > patTauIDTokens_;
0028   std::vector<edm::EDGetTokenT<reco::TauDiscriminatorContainer> > patTauIDContainerTokens_;
0029   size_t nNewPlainTauIds_;
0030   size_t nNewTauIds_;
0031 };
0032
0033 PATTauIDEmbedder::PATTauIDEmbedder(const edm::ParameterSet& cfg) {
0034   src_ = consumes<pat::TauCollection>(cfg.getParameter<edm::InputTag>("src"));
0035   // read the different tau ID names
0036   edm::ParameterSet idps = cfg.getParameter<edm::ParameterSet>("tauIDSources");
0037   std::vector<std::string> names = idps.getParameterNamesForType<edm::ParameterSet>();
0038   std::map<std::string, IDContainerData> idContainerMap;
0039   for (auto const& name : names) {
0040     edm::ParameterSet idp = idps.getParameter<edm::ParameterSet>(name);
0041     int wpidx = idp.getParameter<int>("workingPointIndex");
0042     edm::InputTag tag = idp.getParameter<edm::InputTag>("inputTag");
0043     if (wpidx == -99) {
0044       tauIDSrcs_.push_back(NameTag(name, tag));
0045     } else {
0046       std::map<std::string, IDContainerData>::iterator it;
0047       it = idContainerMap.insert({tag.label() + tag.instance(), {tag, std::vector<NameWPIdx>()}}).first;
0048       it->second.second.push_back(NameWPIdx(name, wpidx));
0049     }
0050   }
0051   // but in any case at least once
0052   if (tauIDSrcs_.empty() && idContainerMap.empty())
0053     throw cms::Exception("Configuration") << "PATTauProducer: id addTauID is true, you must specify:\n"
0054                                           << "\tPSet tauIDSources = { \n"
0055                                           << "\t\tInputTag <someName> = <someTag>   // as many as you want \n "
0056                                           << "\t}\n";
0057   for (auto const& mapEntry : idContainerMap) {
0058     tauIDSrcContainers_.push_back(mapEntry.second.second);
0059     patTauIDContainerTokens_.push_back(mayConsume<reco::TauDiscriminatorContainer>(mapEntry.second.first));
0060   }
0061   patTauIDTokens_ = edm::vector_transform(
0062       tauIDSrcs_, [this](NameTag const& tag) { return mayConsume<pat::PATTauDiscriminator>(tag.second); });
0063   nNewPlainTauIds_ = tauIDSrcs_.size();
0064   nNewTauIds_ = nNewPlainTauIds_;
0065   for (auto const& cont : tauIDSrcContainers_) {
0066     nNewTauIds_ += cont.size();
0067   }
0068
0069   produces<std::vector<pat::Tau> >();
0070 }
0071
0072 void PATTauIDEmbedder::produce(edm::Event& evt, const edm::EventSetup& es) {
0073   edm::Handle<pat::TauCollection> inputTaus;
0074   evt.getByToken(src_, inputTaus);
0075
0076   auto outputTaus = std::make_unique<std::vector<pat::Tau> >();
0077   outputTaus->reserve(inputTaus->size());
0078
0079   int tau_idx = 0;
0080   for (pat::TauCollection::const_iterator inputTau = inputTaus->begin(); inputTau != inputTaus->end();
0081        ++inputTau, ++tau_idx) {
0082     pat::Tau outputTau(*inputTau);
0083     pat::TauRef inputTauRef(inputTaus, tau_idx);
0084     size_t nTauIds = inputTau->tauIDs().size();
0085     std::vector<pat::Tau::IdPair> tauIds(nTauIds + nNewTauIds_);
0086
0087     // copy IDs that are already stored in PAT taus
0088     for (size_t i = 0; i < nTauIds; ++i) {
0089       tauIds[i] = inputTau->tauIDs().at(i);
0090     }
0091
0092     // store IDs that were produced in PATTauDiscriminator format
0093     edm::Handle<pat::PATTauDiscriminator> tauDiscr;
0094     for (size_t i = 0; i < nNewPlainTauIds_; ++i) {
0095       evt.getByToken(patTauIDTokens_[i], tauDiscr);
0096       tauIds[nTauIds + i].first = tauIDSrcs_[i].first;
0097       tauIds[nTauIds + i].second = (*tauDiscr)[inputTauRef];
0098     }
0099
0100     // store IDs that were produced in PATTauDiscriminatorContainer format
0101     size_t nEmbeddedIDs = nTauIds + nNewPlainTauIds_;
0102     for (size_t i = 0; i < tauIDSrcContainers_.size(); ++i) {
0103       auto const& tauDiscrCont = evt.get(patTauIDContainerTokens_[i]);
0104       for (size_t j = 0; j < tauIDSrcContainers_[i].size(); ++j) {
0105         tauIds[nEmbeddedIDs + j].first = tauIDSrcContainers_[i][j].first;
0106         int wpIdx = tauIDSrcContainers_[i][j].second;
0107         if (wpIdx < 0) {
0108           if (tauDiscrCont[inputTauRef].rawValues.size() == 1)
0109             //Only 0th component filled with default value if prediscriminor in PatTauDiscriminator failed.
0110             tauIds[nEmbeddedIDs + j].second = tauDiscrCont[inputTauRef].rawValues.at(0);
0111           else
0112             //uses negative indices to access rawValues. In most cases only one rawValue at WPIdx=-1 exists.
0113             tauIds[nEmbeddedIDs + j].second = tauDiscrCont[inputTauRef].rawValues.at(-1 - wpIdx);
0114         } else {
0115           if (tauDiscrCont[inputTauRef].workingPoints.empty())
0116             //WP vector not filled if prediscriminor in PatTauDiscriminator failed. Set PAT output to false in this case
0117             tauIds[nEmbeddedIDs + j].second = 0.0;
0118           else
0119             tauIds[nEmbeddedIDs + j].second = tauDiscrCont[inputTauRef].workingPoints.at(wpIdx);
0120         }
0121       }
0122       nEmbeddedIDs += tauIDSrcContainers_[i].size();
0123     }
0124
0125     outputTau.setTauIDs(tauIds);
0126     outputTaus->push_back(outputTau);
0127   }
0128
0129   evt.put(std::move(outputTaus));
0130 }
0131
0132 #include "FWCore/Framework/interface/MakerMacros.h"
0133
0134 DEFINE_FWK_MODULE(PATTauIDEmbedder);