TrackingAnalysis/plugins/PreMixingTrackingParticleWorker.cc

0001 #include "FWCore/Framework/interface/ConsumesCollector.h"
0002 #include "FWCore/Framework/interface/Event.h"
0003 #include "FWCore/Framework/interface/ProducesCollector.h"
0004 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0005 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0006 #include "SimGeneral/MixingModule/interface/PileUpEventPrincipal.h"
0007
0008 #include "DataFormats/Common/interface/Handle.h"
0009 #include "SimDataFormats/TrackingAnalysis/interface/TrackingParticle.h"
0010 #include "SimDataFormats/TrackingAnalysis/interface/TrackingVertex.h"
0011
0012 #include "SimGeneral/PreMixingModule/interface/PreMixingWorker.h"
0013 #include "SimGeneral/PreMixingModule/interface/PreMixingWorkerFactory.h"
0014
0015 class PreMixingTrackingParticleWorker : public PreMixingWorker {
0016 public:
0017   PreMixingTrackingParticleWorker(const edm::ParameterSet &ps, edm::ProducesCollector, edm::ConsumesCollector &&iC);
0018   ~PreMixingTrackingParticleWorker() override = default;
0019
0020   void initializeEvent(edm::Event const &iEvent, edm::EventSetup const &iSetup) override;
0021   void addSignals(edm::Event const &iEvent, edm::EventSetup const &iSetup) override;
0022   void addPileups(PileUpEventPrincipal const &pep, edm::EventSetup const &iSetup) override;
0023   void put(edm::Event &iEvent,
0024            edm::EventSetup const &iSetup,
0025            std::vector<PileupSummaryInfo> const &ps,
0026            int bunchSpacing) override;
0027
0028 private:
0029   void add(const std::vector<TrackingParticle> &particles, const std::vector<TrackingVertex> &vertices);
0030
0031   edm::EDGetTokenT<std::vector<TrackingParticle>> TrackSigToken_;  // Token to retrieve information
0032   edm::EDGetTokenT<std::vector<TrackingVertex>> VtxSigToken_;      // Token to retrieve information
0033
0034   edm::InputTag TrackingParticlePileInputTag_;  // InputTag for pileup tracks
0035
0036   std::string TrackingParticleCollectionDM_;  // secondary name to be given to
0037                                               // new TrackingParticle
0038
0039   std::unique_ptr<std::vector<TrackingParticle>> NewTrackList_;
0040   std::unique_ptr<std::vector<TrackingVertex>> NewVertexList_;
0041   TrackingParticleRefProd TrackListRef_;
0042   TrackingVertexRefProd VertexListRef_;
0043 };
0044
0045 PreMixingTrackingParticleWorker::PreMixingTrackingParticleWorker(const edm::ParameterSet &ps,
0046                                                                  edm::ProducesCollector producesCollector,
0047                                                                  edm::ConsumesCollector &&iC)
0048     : TrackSigToken_(iC.consumes<std::vector<TrackingParticle>>(ps.getParameter<edm::InputTag>("labelSig"))),
0049       VtxSigToken_(iC.consumes<std::vector<TrackingVertex>>(ps.getParameter<edm::InputTag>("labelSig"))),
0050       TrackingParticlePileInputTag_(ps.getParameter<edm::InputTag>("pileInputTag")),
0051       TrackingParticleCollectionDM_(ps.getParameter<std::string>("collectionDM")) {
0052   producesCollector.produces<std::vector<TrackingParticle>>(TrackingParticleCollectionDM_);
0053   producesCollector.produces<std::vector<TrackingVertex>>(TrackingParticleCollectionDM_);
0054 }
0055
0056 void PreMixingTrackingParticleWorker::initializeEvent(edm::Event const &iEvent, edm::EventSetup const &iSetup) {
0057   NewTrackList_ = std::make_unique<std::vector<TrackingParticle>>();
0058   NewVertexList_ = std::make_unique<std::vector<TrackingVertex>>();
0059
0060   // need RefProds in order to re-key the particle<->vertex refs
0061   // TODO: try to remove const_cast, requires making Event non-const in
0062   // BMixingModule::initializeEvent
0063   TrackListRef_ =
0064       const_cast<edm::Event &>(iEvent).getRefBeforePut<std::vector<TrackingParticle>>(TrackingParticleCollectionDM_);
0065   VertexListRef_ =
0066       const_cast<edm::Event &>(iEvent).getRefBeforePut<std::vector<TrackingVertex>>(TrackingParticleCollectionDM_);
0067 }
0068
0069 void PreMixingTrackingParticleWorker::addSignals(edm::Event const &iEvent, edm::EventSetup const &iSetup) {
0070   edm::Handle<std::vector<TrackingParticle>> tracks;
0071   iEvent.getByToken(TrackSigToken_, tracks);
0072
0073   edm::Handle<std::vector<TrackingVertex>> vtxs;
0074   iEvent.getByToken(VtxSigToken_, vtxs);
0075
0076   if (tracks.isValid() && vtxs.isValid()) {
0077     add(*tracks, *vtxs);
0078   }
0079 }
0080
0081 void PreMixingTrackingParticleWorker::addPileups(PileUpEventPrincipal const &pep, edm::EventSetup const &iSetup) {
0082   LogDebug("PreMixingTrackingParticleWorker") << "\n===============> adding pileups from event  "
0083                                               << pep.principal().id() << " for bunchcrossing " << pep.bunchCrossing();
0084
0085   edm::Handle<std::vector<TrackingParticle>> inputHandle;
0086   pep.getByLabel(TrackingParticlePileInputTag_, inputHandle);
0087
0088   edm::Handle<std::vector<TrackingVertex>> inputVHandle;
0089   pep.getByLabel(TrackingParticlePileInputTag_, inputVHandle);
0090
0091   if (inputHandle.isValid() && inputVHandle.isValid()) {
0092     add(*inputHandle, *inputVHandle);
0093   }
0094 }
0095
0096 void PreMixingTrackingParticleWorker::add(const std::vector<TrackingParticle> &particles,
0097                                           const std::vector<TrackingVertex> &vertices) {
0098   const size_t StartingIndexV = NewVertexList_->size();
0099   const size_t StartingIndexT = NewTrackList_->size();
0100
0101   // grab Vertices, store copy, preserving indices.  Easier to loop over
0102   // vertices first - fewer links
0103   for (const auto &vtx : vertices) {
0104     NewVertexList_->push_back(vtx);
0105   }
0106
0107   // grab tracks, store copy
0108   for (const auto &track : particles) {
0109     const auto &oldRef = track.parentVertex();
0110     auto newRef = TrackingVertexRef(VertexListRef_, oldRef.index() + StartingIndexV);
0111     NewTrackList_->push_back(track);
0112
0113     auto &Ntrack = NewTrackList_->back();  // modify copy
0114
0115     Ntrack.setParentVertex(newRef);
0116     Ntrack.clearDecayVertices();
0117
0118     // next, loop over daughter vertices, same strategy
0119     for (auto const &vertexRef : track.decayVertices()) {
0120       auto newRef = TrackingVertexRef(VertexListRef_, vertexRef.index() + StartingIndexV);
0121       Ntrack.addDecayVertex(newRef);
0122     }
0123   }
0124
0125   // Now that tracks are handled, go back and put correct Refs in vertices
0126   // Operate only on the added pileup vertices, and leave the already-existing
0127   // vertices untouched
0128   std::vector<decltype(TrackingParticleRef().index())> sourceTrackIndices;
0129   std::vector<decltype(TrackingParticleRef().index())> daughterTrackIndices;
0130   for (size_t iVertex = StartingIndexV; iVertex != NewVertexList_->size(); ++iVertex) {
0131     auto &vertex = (*NewVertexList_)[iVertex];
0132
0133     // Need to copy the indices before clearing the vectors
0134     sourceTrackIndices.reserve(vertex.sourceTracks().size());
0135     daughterTrackIndices.reserve(vertex.daughterTracks().size());
0136     for (auto const &ref : vertex.sourceTracks())
0137       sourceTrackIndices.push_back(ref.index());
0138     for (auto const &ref : vertex.daughterTracks())
0139       daughterTrackIndices.push_back(ref.index());
0140
0141     vertex.clearParentTracks();
0142     vertex.clearDaughterTracks();
0143
0144     for (auto index : sourceTrackIndices) {
0145       auto newRef = TrackingParticleRef(TrackListRef_, index + StartingIndexT);
0146       vertex.addParentTrack(newRef);
0147     }
0148
0149     // next, loop over daughter tracks, same strategy
0150     for (auto index : daughterTrackIndices) {
0151       auto newRef = TrackingParticleRef(TrackListRef_, index + StartingIndexT);
0152       vertex.addDaughterTrack(newRef);
0153     }
0154
0155     sourceTrackIndices.clear();
0156     daughterTrackIndices.clear();
0157   }
0158 }
0159
0160 void PreMixingTrackingParticleWorker::put(edm::Event &iEvent,
0161                                           edm::EventSetup const &iSetup,
0162                                           std::vector<PileupSummaryInfo> const &ps,
0163                                           int bunchSpacing) {
0164   edm::LogInfo("PreMixingTrackingParticleWorker") << "total # Merged Tracks: " << NewTrackList_->size();
0165   iEvent.put(std::move(NewTrackList_), TrackingParticleCollectionDM_);
0166   iEvent.put(std::move(NewVertexList_), TrackingParticleCollectionDM_);
0167 }
0168
0169 DEFINE_PREMIXING_WORKER(PreMixingTrackingParticleWorker);