RecoTracker/PixelLowPtUtilities/plugins/TrackListCombiner.cc - master - githubmirror/cmssw

RecoTracker/PixelLowPtUtilities/plugins/TrackListCombiner.cc
160 строк6 KB
mmusich
move RecoPixelVertexing subsystem to RecoTracker
15 мар 2023, 10:40
15 мар 2023, 10:40ac584e8
#include "TrackListCombiner.h"

#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"

#include "FWCore/MessageLogger/interface/MessageLogger.h"

#include "DataFormats/TrackReco/interface/TrackFwd.h"
#include "DataFormats/TrackReco/interface/Track.h"
#include "DataFormats/TrackReco/interface/TrackExtra.h"

using namespace std;

/*****************************************************************************/
TrackListCombiner::TrackListCombiner(const edm::ParameterSet& ps) {
  for (const std::string& prod : ps.getParameter<vector<string>>("trackProducers")) {
    trackProducers.emplace_back(consumes<vector<Trajectory>>(prod), consumes<TrajTrackAssociationCollection>(prod));
  }

  produces<reco::TrackCollection>();
  produces<reco::TrackExtraCollection>();
  produces<TrackingRecHitCollection>();
  produces<vector<Trajectory>>();
  produces<TrajTrackAssociationCollection>();
}

/*****************************************************************************/
TrackListCombiner::~TrackListCombiner() {}

/*****************************************************************************/
void TrackListCombiner::produce(edm::StreamID, edm::Event& ev, const edm::EventSetup& es) const {
  auto recoTracks = std::make_unique<reco::TrackCollection>();
  auto recoTrackExtras = std::make_unique<reco::TrackExtraCollection>();
  auto recoHits = std::make_unique<TrackingRecHitCollection>();
  auto recoTrajectories = std::make_unique<vector<Trajectory>>();
  auto recoTrajTrackMap = std::make_unique<TrajTrackAssociationCollection>();

  LogTrace("MinBiasTracking") << "[TrackListCombiner]";

  // Go through all track producers
  int i = 1;
  for (auto trackProducer = trackProducers.begin(); trackProducer != trackProducers.end(); trackProducer++, i++) {
    reco::TrackBase::TrackAlgorithm algo;
    switch (i) {
      case 1:
        algo = reco::TrackBase::lowPtTripletStep;
        break;
      case 2:
        algo = reco::TrackBase::pixelPairStep;
        break;
      case 3:
        algo = reco::TrackBase::detachedTripletStep;
        break;
      default:
        algo = reco::TrackBase::undefAlgorithm;
    }

    edm::Handle<vector<Trajectory>> theTrajectoryCollection;
    edm::Handle<TrajTrackAssociationCollection> theAssoMap;

    ev.getByToken(trackProducer->trajectory, theTrajectoryCollection);
    ev.getByToken(trackProducer->assoMap, theAssoMap);

#ifdef EDM_ML_DEBUG
    edm::EDConsumerBase::Labels labels;
    labelsForToken(trackProducer->trajectory, labels);

    LogTrace("MinBiasTracking") << " [TrackListCombiner] " << labels.module << " : " << theAssoMap->size();
#endif

    // The track collection iterators
    TrajTrackAssociationCollection::const_iterator anAssociation;
    TrajTrackAssociationCollection::const_iterator lastAssociation;
    anAssociation = theAssoMap->begin();
    lastAssociation = theAssoMap->end();

    // Build the map of correspondance between reco tracks and sim tracks
    for (; anAssociation != lastAssociation; ++anAssociation) {
      edm::Ref<vector<Trajectory>> aTrajectoryRef = anAssociation->key;
      reco::TrackRef aTrackRef = anAssociation->val;

      // A copy of the track
      reco::Track aRecoTrack(*aTrackRef);

      // Set algorithm
      aRecoTrack.setAlgorithm(algo);

      recoTracks->push_back(aRecoTrack);

      // A copy of the hits
      unsigned nh = aRecoTrack.recHitsSize();
      for (unsigned ih = 0; ih < nh; ++ih) {
        TrackingRecHit* hit = aRecoTrack.recHit(ih)->clone();
        recoHits->push_back(hit);
      }

      // A copy of the trajectories
      recoTrajectories->push_back(*aTrajectoryRef);
    }
  }

  LogTrace("MinBiasTracking") << " [TrackListCombiner] allTracks : " << recoTracks->size() << "|"
                              << recoTrajectories->size();

  // Save the tracking recHits
  edm::OrphanHandle<TrackingRecHitCollection> theRecoHits = ev.put(std::move(recoHits));

  edm::RefProd<TrackingRecHitCollection> theRecoHitsProd(theRecoHits);
  // Create the track extras and add the references to the rechits
  unsigned hits = 0;
  unsigned nTracks = recoTracks->size();
  recoTrackExtras->reserve(nTracks);  // To save some time at push_back
  for (unsigned index = 0; index < nTracks; ++index) {
    reco::Track& aTrack = recoTracks->at(index);
    reco::TrackExtra aTrackExtra(aTrack.outerPosition(),
                                 aTrack.outerMomentum(),
                                 aTrack.outerOk(),
                                 aTrack.innerPosition(),
                                 aTrack.innerMomentum(),
                                 aTrack.innerOk(),
                                 aTrack.outerStateCovariance(),
                                 aTrack.outerDetId(),
                                 aTrack.innerStateCovariance(),
                                 aTrack.innerDetId(),
                                 aTrack.seedDirection(),
                                 aTrack.seedRef());

    unsigned nHits = aTrack.recHitsSize();
    aTrackExtra.setHits(theRecoHitsProd, hits, nHits);
    hits += nHits;

    recoTrackExtras->push_back(aTrackExtra);
  }

  // Save the track extras
  edm::OrphanHandle<reco::TrackExtraCollection> theRecoTrackExtras = ev.put(std::move(recoTrackExtras));

  // Add the reference to the track extra in the tracks
  for (unsigned index = 0; index < nTracks; ++index) {
    const reco::TrackExtraRef theTrackExtraRef(theRecoTrackExtras, index);
    (recoTracks->at(index)).setExtra(theTrackExtraRef);
  }

  // Save the tracks
  edm::OrphanHandle<reco::TrackCollection> theRecoTracks = ev.put(std::move(recoTracks));

  // Save the trajectories
  edm::OrphanHandle<vector<Trajectory>> theRecoTrajectories = ev.put(std::move(recoTrajectories));

  // Create and set the trajectory/track association map
  for (unsigned index = 0; index < nTracks; ++index) {
    edm::Ref<vector<Trajectory>> trajRef(theRecoTrajectories, index);
    edm::Ref<reco::TrackCollection> tkRef(theRecoTracks, index);
    recoTrajTrackMap->insert(trajRef, tkRef);
  }

  // Save the association map
  ev.put(std::move(recoTrajTrackMap));
}