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