TimingIDTools/plugins/MTDTrackQualityMVAProducer.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/stream/EDProducer.h"
0003 #include "FWCore/Framework/interface/Event.h"
0004 #include "FWCore/Framework/interface/EventSetup.h"
0005 #include "FWCore/Framework/interface/ESHandle.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0007 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0008
0009 #include "DataFormats/VertexReco/interface/VertexFwd.h"
0010 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0011 #include "DataFormats/Common/interface/ValueMap.h"
0012 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0013 #include "DataFormats/TrackReco/interface/Track.h"
0014 #include "DataFormats/VertexReco/interface/Vertex.h"
0015
0016 #include "RecoMTD/TimingIDTools/interface/MTDTrackQualityMVA.h"
0017
0018 using namespace std;
0019 using namespace edm;
0020
0021 class MTDTrackQualityMVAProducer : public edm::stream::EDProducer<> {
0022 public:
0023   MTDTrackQualityMVAProducer(const ParameterSet& pset);
0024
0025   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0026
0027   template <class H, class T>
0028   void fillValueMap(edm::Event& iEvent,
0029                     const edm::Handle<H>& handle,
0030                     const std::vector<T>& vec,
0031                     const std::string& name) const;
0032
0033   void produce(edm::Event& ev, const edm::EventSetup& es) final;
0034
0035 private:
0036   static constexpr char mvaName[] = "mtdQualMVA";
0037
0038   edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
0039   edm::EDGetTokenT<reco::TrackCollection> tracksMTDToken_;
0040
0041   edm::EDGetTokenT<edm::ValueMap<float>> btlMatchChi2Token_;
0042   edm::EDGetTokenT<edm::ValueMap<float>> btlMatchTimeChi2Token_;
0043   edm::EDGetTokenT<edm::ValueMap<float>> etlMatchChi2Token_;
0044   edm::EDGetTokenT<edm::ValueMap<float>> etlMatchTimeChi2Token_;
0045   edm::EDGetTokenT<edm::ValueMap<float>> mtdTimeToken_;
0046   edm::EDGetTokenT<edm::ValueMap<float>> pathLengthToken_;
0047   edm::EDGetTokenT<edm::ValueMap<int>> npixBarrelToken_;
0048   edm::EDGetTokenT<edm::ValueMap<int>> npixEndcapToken_;
0049
0050   MTDTrackQualityMVA mva_;
0051 };
0052
0053 MTDTrackQualityMVAProducer::MTDTrackQualityMVAProducer(const ParameterSet& iConfig)
0054     : tracksToken_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("tracksSrc"))),
0055       btlMatchChi2Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("btlMatchChi2Src"))),
0056       btlMatchTimeChi2Token_(
0057           consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("btlMatchTimeChi2Src"))),
0058       etlMatchChi2Token_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("etlMatchChi2Src"))),
0059       etlMatchTimeChi2Token_(
0060           consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("etlMatchTimeChi2Src"))),
0061       mtdTimeToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("mtdTimeSrc"))),
0062       pathLengthToken_(consumes<edm::ValueMap<float>>(iConfig.getParameter<edm::InputTag>("pathLengthSrc"))),
0063       npixBarrelToken_(consumes<edm::ValueMap<int>>(iConfig.getParameter<edm::InputTag>("npixBarrelSrc"))),
0064       npixEndcapToken_(consumes<edm::ValueMap<int>>(iConfig.getParameter<edm::InputTag>("npixEndcapSrc"))),
0065       mva_(iConfig.getParameter<edm::FileInPath>("qualityBDT_weights_file").fullPath()) {
0066   produces<edm::ValueMap<float>>(mvaName);
0067 }
0068
0069 // Configuration descriptions
0070 void MTDTrackQualityMVAProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0071   edm::ParameterSetDescription desc;
0072   desc.add<edm::InputTag>("tracksSrc", edm::InputTag("generalTracks"))->setComment("Input tracks collection");
0073   desc.add<edm::InputTag>("btlMatchChi2Src", edm::InputTag("trackExtenderWithMTD", "btlMatchChi2"))
0074       ->setComment("BTL Chi2 Matching value Map");
0075   desc.add<edm::InputTag>("btlMatchTimeChi2Src", edm::InputTag("trackExtenderWithMTD", "btlMatchTimeChi2"))
0076       ->setComment("BTL Chi2 Matching value Map");
0077   desc.add<edm::InputTag>("etlMatchChi2Src", edm::InputTag("trackExtenderWithMTD", "etlMatchChi2"))
0078       ->setComment("ETL Chi2 Matching value Map");
0079   desc.add<edm::InputTag>("etlMatchTimeChi2Src", edm::InputTag("trackExtenderWithMTD", "etlMatchTimeChi2"))
0080       ->setComment("ETL Chi2 Matching value Map");
0081   desc.add<edm::InputTag>("mtdTimeSrc", edm::InputTag("trackExtenderWithMTD", "generalTracktmtd"))
0082       ->setComment("MTD TIme value Map");
0083   desc.add<edm::InputTag>("pathLengthSrc", edm::InputTag("trackExtenderWithMTD", "generalTrackPathLength"))
0084       ->setComment("MTD PathLength value Map");
0085   desc.add<edm::InputTag>("npixBarrelSrc", edm::InputTag("trackExtenderWithMTD", "npixBarrel"))
0086       ->setComment("# of Barrel pixel associated to refitted tracks");
0087   desc.add<edm::InputTag>("npixEndcapSrc", edm::InputTag("trackExtenderWithMTD", "npixEndcap"))
0088       ->setComment("# of Endcap pixel associated to refitted tracks");
0089   desc.add<edm::FileInPath>("qualityBDT_weights_file",
0090                             edm::FileInPath("RecoMTD/TimingIDTools/data/clf4D_MTDquality_bo.xml"))
0091       ->setComment("Track MTD quality BDT weights");
0092   descriptions.add("mtdTrackQualityMVAProducer", desc);
0093 }
0094
0095 template <class H, class T>
0096 void MTDTrackQualityMVAProducer::fillValueMap(edm::Event& iEvent,
0097                                               const edm::Handle<H>& handle,
0098                                               const std::vector<T>& vec,
0099                                               const std::string& name) const {
0100   auto out = std::make_unique<edm::ValueMap<T>>();
0101   typename edm::ValueMap<T>::Filler filler(*out);
0102   filler.insert(handle, vec.begin(), vec.end());
0103   filler.fill();
0104   iEvent.put(std::move(out), name);
0105 }
0106
0107 void MTDTrackQualityMVAProducer::produce(edm::Event& ev, const edm::EventSetup& es) {
0108   edm::Handle<reco::TrackCollection> tracksH;
0109   ev.getByToken(tracksToken_, tracksH);
0110   const auto& tracks = *tracksH;
0111
0112   const auto& btlMatchChi2 = ev.get(btlMatchChi2Token_);
0113   const auto& btlMatchTimeChi2 = ev.get(btlMatchTimeChi2Token_);
0114   const auto& etlMatchChi2 = ev.get(etlMatchChi2Token_);
0115   const auto& etlMatchTimeChi2 = ev.get(etlMatchTimeChi2Token_);
0116   const auto& pathLength = ev.get(pathLengthToken_);
0117   const auto& npixBarrel = ev.get(npixBarrelToken_);
0118   const auto& npixEndcap = ev.get(npixEndcapToken_);
0119   const auto& mtdTime = ev.get(mtdTimeToken_);
0120
0121   std::vector<float> mvaOutRaw;
0122
0123   //Loop over tracks collection
0124   for (unsigned int itrack = 0; itrack < tracks.size(); ++itrack) {
0125     const reco::TrackRef trackref(tracksH, itrack);
0126     if (pathLength[trackref] == -1.)
0127       mvaOutRaw.push_back(-1.);
0128     else {
0129       mvaOutRaw.push_back(mva_(trackref,
0130                                npixBarrel,
0131                                npixEndcap,
0132                                btlMatchChi2,
0133                                btlMatchTimeChi2,
0134                                etlMatchChi2,
0135                                etlMatchTimeChi2,
0136                                mtdTime,
0137                                pathLength));
0138     }
0139   }
0140   fillValueMap(ev, tracksH, mvaOutRaw, mvaName);
0141 }
0142
0143 //define this as a plug-in
0144 #include <FWCore/Framework/interface/MakerMacros.h>
0145 DEFINE_FWK_MODULE(MTDTrackQualityMVAProducer);