HICaloCompatibleTrackSelector.cc

Go to the documentation of this file.

/*
 
 Based on analytical track selector   
 
 - This track selector assigns a quality bit to tracks deemed compatible with matching calo info
 - The default mode is to use the matching provided by particle flow,
   but a delta R matching to calorimeter towers is also supported
 - No selection is done other then selecting calo-compatible tracks.
 - The code always keeps all tracks in the input collection (should make configurable) 
 - Note that matching by PF candidate only works on the same track collection used as input to PF
 - Tower code not up to data
 
   Authors:  Matthew Nguyen, Andre Yoon, Frank Ma (November 4th, 2011)

 */


// Basic inclusion
#include "RecoHI/HiTracking/interface/HICaloCompatibleTrackSelector.h"

#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/ESHandle.h"

#include "DataFormats/Math/interface/deltaR.h"
#include <Math/DistFunc.h>
#include "TMath.h"


using reco::modules::HICaloCompatibleTrackSelector;

HICaloCompatibleTrackSelector::HICaloCompatibleTrackSelector( const edm::ParameterSet & cfg ) :
  srcTracks_(consumes<TrackCollection>(cfg.getParameter<edm::InputTag>("srcTracks"))),
  srcPFCands_(consumes<PFCandidateCollection>(cfg.getParameter<edm::InputTag>("srcPFCands"))),
  srcTower_(consumes<CaloTowerCollection>(cfg.getParameter<edm::InputTag>("srcTower"))),
  usePFCandMatching_(cfg.getUntrackedParameter<bool>("usePFCandMatching", true)),
  trkMatchPtMin_(cfg.getUntrackedParameter<double>("trkMatchPtMin",10.0)),
  trkCompPtMin_(cfg.getUntrackedParameter<double>("trkCompPtMin",35.0)),
  trkEtaMax_(cfg.getUntrackedParameter<double>("trkEtaMax",2.4)),
  towerPtMin_(cfg.getUntrackedParameter<double>("towerPtMin",5.0)),
  matchConeRadius_(cfg.getUntrackedParameter<double>("matchConeRadius",0.087)),
  keepAllTracks_(cfg.getUntrackedParameter<bool>("keepAllTracks", true)),
  copyExtras_(cfg.getUntrackedParameter<bool>("copyExtras", true)),
  copyTrajectories_(cfg.getUntrackedParameter<bool>("copyTrajectories", true)),
  qualityToSet_(cfg.getParameter<std::string>("qualityToSet")),
  qualityToSkip_(cfg.getParameter<std::string>("qualityToSkip")),
  qualityToMatch_(cfg.getParameter<std::string>("qualityToMatch")),
  minimumQuality_(cfg.getParameter<std::string>("minimumQuality")),
  resetQuality_(cfg.getUntrackedParameter<bool>("resetQuality", true)),
  passMuons_(cfg.getUntrackedParameter<bool>("passMuons", true)),
  passElectrons_(cfg.getUntrackedParameter<bool>("passElectrons", false)),
  funcDeltaRTowerMatch_(cfg.getParameter<std::string>("funcDeltaRTowerMatch")),
  funcCaloComp_(cfg.getParameter<std::string>("funcCaloComp"))
{
  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");
    srcTrackTrajs_ = (consumes<std::vector<Trajectory> >(cfg.getParameter<edm::InputTag>("srcTracks")));
    srcTrackTrajAssoc_ = (consumes<TrajTrackAssociationCollection>(cfg.getParameter<edm::InputTag>("srcTracks")));
  }

  // pt dependence of delta R matching requirement
  fDeltaRTowerMatch = new TF1("fDeltaRTowerMatch",funcDeltaRTowerMatch_.c_str(),0,200); 
  // pt dependance of calo compatibility, i.e., minimum sum Calo Et vs. track pT
  fCaloComp = new TF1("fCaloComp",funcCaloComp_.c_str(),0,200); // a parameterization of pt dependent cut


}

HICaloCompatibleTrackSelector::~HICaloCompatibleTrackSelector() {
}


void HICaloCompatibleTrackSelector::produce( edm::Event& evt, const edm::EventSetup& es ) 
{
  using namespace std; 
  using namespace edm;
  using namespace reco;
  
  LogDebug("HICaloCompatibleTrackSelector")<<"min pt for selection = "<< trkMatchPtMin_<<endl;
  
  
  Handle<TrackCollection> hSrcTrack;
  Handle< vector<Trajectory> > hTraj;
  Handle< vector<Trajectory> > hTrajP;
  Handle< TrajTrackAssociationCollection > hTTAss;

  evt.getByToken(srcTracks_,hSrcTrack);
  
  selTracks_ = std::make_unique<TrackCollection>();
  rTracks_ = evt.getRefBeforePut<TrackCollection>();      
  if (copyExtras_) {
    selTrackExtras_ = std::make_unique<TrackExtraCollection>();
    selHits_ = std::make_unique<TrackingRecHitCollection>();
    rHits_ = evt.getRefBeforePut<TrackingRecHitCollection>();
    rTrackExtras_ = evt.getRefBeforePut<TrackExtraCollection>();
  }


  if (copyTrajectories_) trackRefs_.resize(hSrcTrack->size());


  Handle<PFCandidateCollection> pfCandidates;
  Handle<CaloTowerCollection> towers;

  bool isPFThere = false;
  bool isTowerThere = false;
  
  if(usePFCandMatching_) isPFThere = evt.getByToken(srcPFCands_, pfCandidates);  
  else isTowerThere = evt.getByToken(srcTower_, towers);
  
  size_t current = 0;
  for (TI ti = hSrcTrack->begin(), ed = hSrcTrack->end(); ti != ed; ++ti, ++current) {
    

    const reco::Track& trk = *ti;
    
    bool isSelected;
    if(usePFCandMatching_) isSelected = selectByPFCands(ti, hSrcTrack, pfCandidates, isPFThere);
    else isSelected = selectByTowers(ti, hSrcTrack, towers, isTowerThere);
    
    if(!keepAllTracks_ && !isSelected) continue; 

    // Add all tracks to output collection, the rest of the code only sets the quality bit
    selTracks_->push_back( reco::Track( trk ) ); // clone and store
    
    
    if(isSelected) selTracks_->back().setQuality(reco::TrackBase::qualityByName(qualityToSet_));
    
        
    if (copyExtras_) {
      // 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(), trk.seedRef() ) );
      selTracks_->back().setExtra( TrackExtraRef( rTrackExtras_, selTrackExtras_->size() - 1) );
      TrackExtra & tx = selTrackExtras_->back();
      tx.setResiduals(trk.residuals());
      // 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 );
    }
    if (copyTrajectories_) {
      trackRefs_[current] = TrackRef(rTracks_, selTracks_->size() - 1);
    }

    
  }  // close track loop
  
  if ( copyTrajectories_ ) {
    Handle< vector<Trajectory> > hTraj;
    Handle< TrajTrackAssociationCollection > hTTAss;
    evt.getByToken(srcTrackTrajs_, hTTAss);
    evt.getByToken(srcTrackTrajAssoc_, hTraj);
    selTrajs_ = std::make_unique<std::vector<Trajectory>>(); 
    rTrajectories_ = evt.getRefBeforePut< vector<Trajectory> >();
    selTTAss_ = std::make_unique<TrajTrackAssociationCollection>();
    for (size_t i = 0, n = hTraj->size(); i < n; ++i) {
      Ref< vector<Trajectory> > trajRef(hTraj, i);
      TrajTrackAssociationCollection::const_iterator match = hTTAss->find(trajRef);
      if (match != hTTAss->end()) {
    const Ref<TrackCollection> &trkRef = match->val; 
    short oldKey = static_cast<short>(trkRef.key());
    if (trackRefs_[oldKey].isNonnull()) {
      selTrajs_->push_back( Trajectory(*trajRef) );
      selTTAss_->insert ( Ref< vector<Trajectory> >(rTrajectories_, selTrajs_->size() - 1), trackRefs_[oldKey] );
    }
      }
    }
  }
  
  static const string emptyString;
  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_));
  }
}



bool 
HICaloCompatibleTrackSelector::selectByPFCands(TI ti, edm::Handle<TrackCollection> hSrcTrack, edm::Handle<PFCandidateCollection> pfCandidates, bool isPFThere)
{

  const reco::Track& trk = *ti;

  // If it passes this quality threshold or is under the minimum match pT, automatically save it
  if(trk.quality(reco::TrackBase::qualityByName(qualityToSkip_))){
    return true;
  }
  else if(!trk.quality(reco::TrackBase::qualityByName(minimumQuality_))){
    return false;
  }
  else
    {      
      
      double trkPt = trk.pt();
      //if(trkPt < trkMatchPtMin_ ) return false;
      
      double caloEt = 0.0;
      if(usePFCandMatching_){
    if(isPFThere){
      unsigned int trackKey = ti - hSrcTrack->begin();
      caloEt = matchPFCandToTrack(pfCandidates, trackKey, trkPt);      
    }
      }
      
      // Set quality bit based on calo matching
      if(!(caloEt>0.)) return false;
      
      
      if(trkPt<=trkCompPtMin_){
    if(trk.quality(reco::TrackBase::qualityByName(qualityToMatch_))) return true;
    else return false;
      }
      else{
    // loose cuts are implied in selectors, make configurable?
    float compPt = (fCaloComp->Eval(trkPt)!=fCaloComp->Eval(trkPt)) ? 0 : fCaloComp->Eval(trkPt); // protect agains NaN
    if(caloEt>compPt) return true;
    else return false;
      }            
    } // else above trkMatchPtMin_

  throw cms::Exception("Undefined case in HICaloCompatibleTrackSelector") << "Undefined case in HICaloCompatibleTrackSelector";
}


bool 
HICaloCompatibleTrackSelector::selectByTowers(TI ti, edm::Handle<TrackCollection> hSrcTrack, edm::Handle<CaloTowerCollection> towers, bool isTowerThere)
{

  // Not up to date! use PF towers instead

  const reco::Track& trk = *ti;

  // If it passes the high purity cuts, then consider it confirmed
  if(trk.quality(reco::TrackBase::qualityByName(qualityToSkip_))){
    return true;
  }
  else{
    if(trk.pt() <  trkMatchPtMin_ || !trk.quality(reco::TrackBase::qualityByName(qualityToMatch_))) return false;
    
    double caloEt = 0.0;
    if(isTowerThere){      
      double matchDr;
      matchByDrAllowReuse(trk,towers,matchDr,caloEt);
      float matchConeRadius_pt = (fDeltaRTowerMatch->Eval(trk.pt())!=fDeltaRTowerMatch->Eval(trk.pt())) ? 0 : fDeltaRTowerMatch->Eval(trk.pt()); // protect agains NaN
      if (caloEt>0 && matchDr>matchConeRadius_pt) caloEt=0.;      
    }
    
    if(trk.pt()<trkCompPtMin_||caloEt>0.75*(trk.pt()-trkCompPtMin_)) return true;
    else return false;
  }
}

double
HICaloCompatibleTrackSelector::matchPFCandToTrack( const edm::Handle<reco::PFCandidateCollection> & pfCandidates, unsigned it, double trkPt)
{

  // This function returns the sum of the calo energy in the block containing the track
  // There is another piece of information which could be useful which is the pT assigned to the charged hadron by PF
  

  double sum_ecal=0.0, sum_hcal=0.0;
    
  int candType = 0;
  reco::PFCandidate matchedCand; 

  // loop over the PFCandidates until you find the one whose trackRef points to the track

  for( CI ci  = pfCandidates->begin(); ci!=pfCandidates->end(); ++ci)  {

    const reco::PFCandidate& cand = *ci;

    int type = cand.particleId();
    
    // only charged hadrons and leptons can be asscociated with a track
    if(!(type == PFCandidate::h ||     //type1
     type == PFCandidate::e ||     //type2
     type == PFCandidate::mu      //type3
     )
       ) continue;
    
    
    unsigned candTrackRefKey = cand.trackRef().key();      
    
    if(it==candTrackRefKey) {
      matchedCand = cand;
      candType = type;
      break;      
    }
  }
  // take all muons as compatible, extend to electrons when validataed
  if(passMuons_ && candType==3) return 9999.;
  if(passElectrons_ && candType==2) return 9999.;

  if(trkPt < trkMatchPtMin_ ) return 0.;

  if(candType>0){
    
    // Now that we found the matched PF candidate, loop over the elements in the block summing the calo Et
    for(unsigned ib=0; ib<matchedCand.elementsInBlocks().size(); ib++) {
      
      PFBlockRef blockRef = matchedCand.elementsInBlocks()[ib].first;
            
      unsigned indexInBlock = matchedCand.elementsInBlocks()[ib].second;
      const edm::OwnVector<  reco::PFBlockElement>&  elements = (*blockRef).elements();
      
      //This tells you what type of element it is:
      //cout<<" block type"<<elements[indexInBlock].type()<<endl;
      
      switch (elements[indexInBlock].type()) {
    
      case PFBlockElement::ECAL: {
    reco::PFClusterRef clusterRef = elements[indexInBlock].clusterRef();
    sum_ecal += clusterRef->energy()/cosh(clusterRef->eta());
    break;
      }
    
      case PFBlockElement::HCAL: {
    reco::PFClusterRef clusterRef = elements[indexInBlock].clusterRef();
    sum_hcal += clusterRef->energy()/cosh(clusterRef->eta());
    break; 
      }       
      case PFBlockElement::TRACK: {
    //Do nothing since track are not normally linked to other tracks
    break; 
      }       
      default:
    break;
      }
      
      // We could also add in the PS, HO, ..

    } // end of elementsInBlocks()
  }  // end of isCandFound  
  
  

  return sum_ecal+sum_hcal;
  
}

void HICaloCompatibleTrackSelector::matchByDrAllowReuse(const reco::Track & trk, const edm::Handle<CaloTowerCollection> & towers, double & bestdr, double & bestpt)
{
  // loop over towers
  bestdr=1e10;
  bestpt=0.;
  for(unsigned int i = 0; i < towers->size(); ++i){
    const CaloTower & tower= (*towers)[i];
    double pt = tower.pt();
    if (pt<towerPtMin_) continue;
    if (fabs(tower.eta())>trkEtaMax_) continue;
    double dr = reco::deltaR(tower,trk);
    if (dr<bestdr) {
      bestdr = dr;
      bestpt = pt;
    }
  }
}