RecoAlgos/plugins/ConstrainedFitCandProducer.h

0001 #ifndef RecoAlgos_ConstrainedFitCandProducer_h
0002 #define RecoAlgos_ConstrainedFitCandProducer_h
0003 /* \class ConstrainedFitCandProducer

0004  *

0005  * \author Luca Lista, INFN

0006  *

0007  */
0008 #include "FWCore/Framework/interface/stream/EDProducer.h"
0009 #include "FWCore/Utilities/interface/InputTag.h"
0010 #include "CommonTools/UtilAlgos/interface/ParameterAdapter.h"
0011 #include "CommonTools/UtilAlgos/interface/EventSetupInitTrait.h"
0012 #include "DataFormats/Candidate/interface/CandidateFwd.h"
0013 #include "SimGeneral/HepPDTRecord/interface/PdtEntry.h"
0014 #include <vector>
0015
0016 template <typename Fitter,
0017           typename InputCollection = reco::CandidateCollection,
0018           typename OutputCollection = InputCollection,
0019           typename Init = typename ::reco::modules::EventSetupInit<Fitter>::type>
0020 class ConstrainedFitCandProducer : public edm::stream::EDProducer<> {
0021 public:
0022   explicit ConstrainedFitCandProducer(const edm::ParameterSet&);
0023
0024 private:
0025   edm::EDGetTokenT<InputCollection> srcToken_;
0026   bool setLongLived_;
0027   bool setMassConstraint_;
0028   bool setPdgId_;
0029   int pdgId_;
0030   Init fitterInit_;
0031   Fitter fitter_;
0032   void produce(edm::Event&, const edm::EventSetup&) override;
0033 };
0034
0035 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0036 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0037 #include "DataFormats/Common/interface/Handle.h"
0038 #include "FWCore/Framework/interface/ESHandle.h"
0039 #include "FWCore/Framework/interface/Event.h"
0040 #include "DataFormats/Candidate/interface/VertexCompositeCandidate.h"
0041 #include "FWCore/Framework/interface/EventSetup.h"
0042 #include <algorithm>
0043
0044 template <typename Fitter, typename InputCollection, typename OutputCollection, typename Init>
0045 ConstrainedFitCandProducer<Fitter, InputCollection, OutputCollection, Init>::ConstrainedFitCandProducer(
0046     const edm::ParameterSet& cfg)
0047     : srcToken_(consumes<InputCollection>(cfg.template getParameter<edm::InputTag>("src"))),
0048       setLongLived_(false),
0049       setMassConstraint_(false),
0050       setPdgId_(false),
0051       fitterInit_(consumesCollector()),
0052       fitter_(reco::modules::make<Fitter>(cfg)) {
0053   produces<OutputCollection>();
0054   std::string alias(cfg.getParameter<std::string>("@module_label"));
0055   const std::string setLongLived("setLongLived");
0056   std::vector<std::string> vBoolParams = cfg.template getParameterNamesForType<bool>();
0057   bool found = find(vBoolParams.begin(), vBoolParams.end(), setLongLived) != vBoolParams.end();
0058   if (found)
0059     setLongLived_ = cfg.template getParameter<bool>("setLongLived");
0060   const std::string setMassConstraint("setMassConstraint");
0061   found = find(vBoolParams.begin(), vBoolParams.end(), setMassConstraint) != vBoolParams.end();
0062   if (found)
0063     setMassConstraint_ = cfg.template getParameter<bool>("setMassConstraint");
0064   const std::string setPdgId("setPdgId");
0065   std::vector<std::string> vIntParams = cfg.getParameterNamesForType<int>();
0066   found = find(vIntParams.begin(), vIntParams.end(), setPdgId) != vIntParams.end();
0067   if (found) {
0068     setPdgId_ = true;
0069     pdgId_ = cfg.getParameter<int>("setPdgId");
0070   }
0071 }
0072
0073 namespace reco {
0074   namespace fitHelper {
0075     template <typename C>
0076     struct Adder {
0077       static void add(C* c, std::unique_ptr<reco::VertexCompositeCandidate> t) { c->push_back(*t); }
0078     };
0079
0080     template <typename T>
0081     struct Adder<edm::OwnVector<T> > {
0082       static void add(edm::OwnVector<T>* c, std::unique_ptr<reco::VertexCompositeCandidate> t) {
0083         c->push_back(std::move(t));
0084       }
0085     };
0086
0087     template <typename C>
0088     inline void add(C* c, std::unique_ptr<reco::VertexCompositeCandidate> t) {
0089       Adder<C>::add(c, std::move(t));
0090     }
0091   }  // namespace fitHelper

0092 }  // namespace reco

0093
0094 template <typename Fitter, typename InputCollection, typename OutputCollection, typename Init>
0095 void ConstrainedFitCandProducer<Fitter, InputCollection, OutputCollection, Init>::produce(edm::Event& evt,
0096                                                                                           const edm::EventSetup& es) {
0097   fitterInit_.init(fitter_, evt, es);
0098   edm::Handle<InputCollection> cands;
0099   evt.getByToken(srcToken_, cands);
0100   auto fitted = std::make_unique<OutputCollection>();
0101   fitted->reserve(cands->size());
0102   for (typename InputCollection::const_iterator c = cands->begin(); c != cands->end(); ++c) {
0103     auto clone = std::make_unique<reco::VertexCompositeCandidate>(*c);
0104     fitter_.set(*clone);
0105     if (setLongLived_)
0106       clone->setLongLived();
0107     if (setMassConstraint_)
0108       clone->setMassConstraint();
0109     if (setPdgId_)
0110       clone->setPdgId(pdgId_);
0111     reco::fitHelper::add(fitted.get(), std::move(clone));
0112   }
0113   evt.put(std::move(fitted));
0114 }
0115
0116 #endif