TrackFullCloneSelectorBase.h

Go to the documentation of this file.

#ifndef RecoAlgos_TrackFullCloneSelectorBase_h
#define RecoAlgos_TrackFullCloneSelectorBase_h

#include <utility>
#include <vector>
#include <memory>
#include <algorithm>
#include <map>
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/ConsumesCollector.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/Utilities/interface/InputTag.h"

#include "DataFormats/TrackReco/interface/TrackFwd.h"
#include "DataFormats/TrackReco/interface/Track.h"
#include "DataFormats/TrackReco/interface/TrackExtra.h"
#include "TrackingTools/PatternTools/interface/Trajectory.h"
#include "TrackingTools/PatternTools/interface/TrajTrackAssociation.h"


namespace reco { namespace modules {

template<typename Selector>
class TrackFullCloneSelectorBase : public edm::stream::EDProducer<> {
public:
  explicit TrackFullCloneSelectorBase( const edm::ParameterSet & cfg ) :
    hSrcTrackToken_( consumes<reco::TrackCollection>( cfg.template getParameter<edm::InputTag>( "src" ) ) ),
    hTrajToken_( mayConsume< std::vector<Trajectory> >( cfg.template getParameter<edm::InputTag>( "src" ) ) ),
    hTTAssToken_( mayConsume< TrajTrackAssociationCollection >( cfg.template getParameter<edm::InputTag>( "src" ) ) ),
    copyExtras_(cfg.template getUntrackedParameter<bool>("copyExtras", false)),
    copyTrajectories_(cfg.template getUntrackedParameter<bool>("copyTrajectories", false)),
    selector_( cfg, consumesCollector() ) {
      std::string alias( cfg.getParameter<std::string>( "@module_label" ) );
      produces<reco::TrackCollection>().setBranchAlias( alias + "Tracks" );
      if (copyExtras_) {
          produces<reco::TrackExtraCollection>().setBranchAlias( alias + "TrackExtras" );
          produces<TrackingRecHitCollection>().setBranchAlias( alias + "RecHits" );
          if (copyTrajectories_) {
              produces< std::vector<Trajectory> >().setBranchAlias( alias + "Trajectories" );
              produces< TrajTrackAssociationCollection >().setBranchAlias( alias + "TrajectoryTrackAssociations" );
          }
      }
   }
  ~TrackFullCloneSelectorBase() override { }

private:
  void produce( edm::Event& evt, const edm::EventSetup& es) override {
      edm::Handle<reco::TrackCollection> hSrcTrack;
      evt.getByToken( hSrcTrackToken_, hSrcTrack );

      selTracks_ = std::unique_ptr<reco::TrackCollection>(new reco::TrackCollection());
      if (copyExtras_) {
          selTrackExtras_ = std::unique_ptr<reco::TrackExtraCollection>(new reco::TrackExtraCollection());
          selHits_ = std::unique_ptr<TrackingRecHitCollection>(new TrackingRecHitCollection());
      }

      TrackRefProd rTracks = evt.template getRefBeforePut<TrackCollection>();

      TrackingRecHitRefProd rHits;
      TrackExtraRefProd rTrackExtras;
      if (copyExtras_) {
          rHits = evt.template getRefBeforePut<TrackingRecHitCollection>();
          rTrackExtras = evt.template getRefBeforePut<TrackExtraCollection>();
      }

      typedef reco::TrackRef::key_type TrackRefKey;
      std::map<TrackRefKey, reco::TrackRef  > goodTracks;
      TrackRefKey current = 0;

      selector_.init(evt,es);
      auto tkBegin = hSrcTrack->begin();
      for (reco::TrackCollection::const_iterator it = tkBegin, ed = hSrcTrack->end(); it != ed; ++it, ++current) {
      const reco::Track & trk = * it;
      const reco::TrackRef tkref(hSrcTrack,std::distance(tkBegin,it));
          if (!selector_(tkref)) continue;

          selTracks_->push_back( Track( trk ) ); // clone and store
          if (!copyExtras_) continue;

          // TrackExtras
          selTrackExtras_->push_back( TrackExtra( trk.outerPosition(), trk.outerMomentum(), trk.outerOk(),
                      trk.innerPosition(), trk.innerMomentum(), trk.innerOk(),
                      trk.outerStateCovariance(), trk.outerDetId(),
                      trk.innerStateCovariance(), trk.innerDetId(),
                      trk.seedDirection() ) );
          selTracks_->back().setExtra( TrackExtraRef( rTrackExtras, selTrackExtras_->size() - 1) );
          TrackExtra & tx = selTrackExtras_->back();
          // TrackingRecHits
          auto const firstHitIndex = selHits_->size();
          for( trackingRecHit_iterator hit = trk.recHitsBegin(); hit != trk.recHitsEnd(); ++ hit ) {
              selHits_->push_back( (*hit)->clone() );
          }
          tx.setHits( rHits, firstHitIndex, selHits_->size() - firstHitIndex );
          tx.setTrajParams(trk.extra()->trajParams(),trk.extra()->chi2sX5());
          assert(tx.trajParams().size()==tx.recHitsSize());
          if (copyTrajectories_) {
              goodTracks[current] = reco::TrackRef(rTracks, selTracks_->size() - 1);
          }
      }
      if ( copyTrajectories_ ) {
          edm::Handle< std::vector<Trajectory> > hTraj;
          edm::Handle< TrajTrackAssociationCollection > hTTAss;
          evt.getByToken(hTTAssToken_, hTTAss);
          evt.getByToken(hTrajToken_, hTraj);
          edm::RefProd< std::vector<Trajectory> > TrajRefProd = evt.template getRefBeforePut< std::vector<Trajectory> >();
          selTrajs_ = std::unique_ptr< std::vector<Trajectory> >(new std::vector<Trajectory>());
          selTTAss_ = std::unique_ptr< TrajTrackAssociationCollection >(new TrajTrackAssociationCollection());
          for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
              edm::Ref< std::vector<Trajectory> > trajRef(hTraj, i);
              TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
              if (match != hTTAss->end()) {
                  const edm::Ref<reco::TrackCollection> &trkRef = match->val;
                  TrackRefKey oldKey = trkRef.key();
                  std::map<TrackRefKey, reco::TrackRef>::iterator getref = goodTracks.find(oldKey);
                  if (getref != goodTracks.end()) {
                      // do the clone
                      selTrajs_->push_back( Trajectory(*trajRef) );
                      selTTAss_->insert ( edm::Ref< std::vector<Trajectory> >(TrajRefProd, selTrajs_->size() - 1),
                                          getref->second );
                  }
              }
          }
      }

      evt.put(std::move(selTracks_));
      if (copyExtras_) {
            evt.put(std::move(selTrackExtras_));
            evt.put(std::move(selHits_));
            if ( copyTrajectories_ ) {
                evt.put(std::move(selTrajs_));
                evt.put(std::move(selTTAss_));
            }
      }
  }
  edm::EDGetTokenT<reco::TrackCollection> hSrcTrackToken_;
  edm::EDGetTokenT< std::vector<Trajectory> > hTrajToken_;
  edm::EDGetTokenT< TrajTrackAssociationCollection > hTTAssToken_;
  bool copyExtras_;
  bool copyTrajectories_;
  Selector selector_;
  // some space
  std::unique_ptr<reco::TrackCollection> selTracks_;
  std::unique_ptr<reco::TrackExtraCollection> selTrackExtras_;
  std::unique_ptr<TrackingRecHitCollection> selHits_;
  std::unique_ptr< std::vector<Trajectory> > selTrajs_;
  std::unique_ptr< TrajTrackAssociationCollection > selTTAss_;
};

} }
#endif