FinalTrackSelectors/plugins/TrackSelectorByRegion.cc

0001 #include "FWCore/Framework/interface/Frameworkfwd.h"
0002 #include "FWCore/Framework/interface/global/EDProducer.h"
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0007 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0008
0009 #include "DataFormats/TrackReco/interface/TrackFwd.h"
0010 #include "DataFormats/TrackReco/interface/Track.h"
0011 #include "DataFormats/Common/interface/OwnVector.h"
0012 #include "RecoTracker/TkTrackingRegions/interface/TrackingRegion.h"
0013
0014 #include <vector>
0015 #include <memory>
0016
0017 class TrackSelectorByRegion final : public edm::global::EDProducer<> {
0018 public:
0019   explicit TrackSelectorByRegion(const edm::ParameterSet& conf)
0020       : produceCollection_(conf.getParameter<bool>("produceTrackCollection")),
0021         produceMask_(conf.getParameter<bool>("produceMask")),
0022         tracksToken_(consumes<reco::TrackCollection>(conf.getParameter<edm::InputTag>("tracks"))),
0023         inputTrkRegionToken_(consumes<edm::OwnVector<TrackingRegion>>(conf.getParameter<edm::InputTag>("regions"))),
0024         outputTracksToken_(produceCollection_ ? produces<reco::TrackCollection>()
0025                                               : edm::EDPutTokenT<reco::TrackCollection>{}),
0026         outputMaskToken_(produceMask_ ? produces<std::vector<bool>>() : edm::EDPutTokenT<std::vector<bool>>{}) {}
0027
0028   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0029     edm::ParameterSetDescription desc;
0030     desc.add<edm::InputTag>("tracks", edm::InputTag("hltPixelTracks"));
0031     desc.add<edm::InputTag>("regions", edm::InputTag(""));
0032     desc.add<bool>("produceTrackCollection", true);
0033     desc.add<bool>("produceMask", true);
0034     descriptions.add("trackSelectorByRegion", desc);
0035   }
0036
0037 private:
0038   using MaskCollection = std::vector<bool>;
0039
0040   void produce(edm::StreamID, edm::Event& iEvent, const edm::EventSetup&) const override {
0041     if (not produceCollection_ and not produceMask_)
0042       return;
0043
0044     auto const& regions = iEvent.get(inputTrkRegionToken_);
0045     auto const& tracks = iEvent.get(tracksToken_);
0046     MaskCollection mask(tracks.size(), false);  // output mask
0047
0048     for (auto const& region : regions) {
0049       region.checkTracks(tracks, mask);
0050     }
0051
0052     if (produceCollection_) {
0053       reco::TrackCollection output_tracks;  // output collection with a (shallow) copy of the selected tracks
0054       size_t size = 0;
0055       // count the number of selected tracks
0056       for (size_t i = 0; i < mask.size(); i++) {
0057         size += mask[i];
0058       }
0059       output_tracks.reserve(size);
0060       for (size_t i = 0; i < mask.size(); i++) {
0061         if (mask[i])
0062           output_tracks.push_back(tracks[i]);
0063       }
0064       iEvent.emplace(outputTracksToken_, std::move(output_tracks));
0065     }
0066     if (produceMask_) {
0067       iEvent.emplace(outputMaskToken_, std::move(mask));
0068     }
0069   }
0070
0071   const bool produceCollection_;
0072   const bool produceMask_;
0073   const edm::EDGetTokenT<reco::TrackCollection> tracksToken_;
0074   const edm::EDGetTokenT<edm::OwnVector<TrackingRegion>> inputTrkRegionToken_;
0075   const edm::EDPutTokenT<reco::TrackCollection> outputTracksToken_;
0076   const edm::EDPutTokenT<std::vector<bool>> outputMaskToken_;
0077 };
0078
0079 #include "FWCore/PluginManager/interface/ModuleDef.h"
0080 #include "FWCore/Framework/interface/MakerMacros.h"
0081 DEFINE_FWK_MODULE(TrackSelectorByRegion);