IPProducer.h

Go to the documentation of this file.

#ifndef RecoBTag_IPProducer
#define RecoBTag_IPProducer

// system include files
#include <algorithm>
#include <cmath>
#include <functional>
#include <iostream>
#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/Event.h"
#include "DataFormats/Common/interface/View.h"
#include "FWCore/Framework/interface/MakerMacros.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
#include "FWCore/Utilities/interface/ESGetToken.h"

#include "DataFormats/TrackReco/interface/Track.h"
#include "DataFormats/JetReco/interface/JetTracksAssociation.h"
#include "DataFormats/BTauReco/interface/JetTag.h"
#include "DataFormats/BTauReco/interface/TrackIPTagInfo.h"
#include "DataFormats/Candidate/interface/Candidate.h"

#include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
#include "TrackingTools/IPTools/interface/IPTools.h"
#include "TrackingTools/Records/interface/TransientTrackRecord.h"

#include "RecoVertex/VertexPrimitives/interface/VertexState.h"
#include "RecoVertex/VertexPrimitives/interface/ConvertError.h"
#include "RecoVertex/VertexPrimitives/interface/ConvertToFromReco.h"
#include "RecoVertex/VertexTools/interface/VertexDistance3D.h"
#include "RecoVertex/GhostTrackFitter/interface/GhostTrack.h"
#include "RecoVertex/GhostTrackFitter/interface/GhostTrackState.h"
#include "RecoVertex/GhostTrackFitter/interface/GhostTrackPrediction.h"
#include "RecoVertex/GhostTrackFitter/interface/GhostTrackFitter.h"

#include "RecoBTag/TrackProbability/interface/HistogramProbabilityEstimator.h"
#include "DataFormats/BTauReco/interface/JTATagInfo.h"
#include "DataFormats/BTauReco/interface/JetTagInfo.h"
#include "CondFormats/DataRecord/interface/BTagTrackProbability2DRcd.h"
#include "CondFormats/DataRecord/interface/BTagTrackProbability3DRcd.h"

// system include files
#include <memory>

// user include files
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Utilities/interface/InputTag.h"
#include "FWCore/Framework/interface/ConsumesCollector.h"

class HistogramProbabilityEstimator;

namespace IPProducerHelpers {
  class FromJTA {
  public:
    FromJTA(const edm::ParameterSet& iConfig, edm::ConsumesCollector&& iC)
        : token_associator(
              iC.consumes<reco::JetTracksAssociationCollection>(iConfig.getParameter<edm::InputTag>("jetTracks"))) {}
    reco::TrackRefVector tracks(const reco::JTATagInfo& it) { return it.tracks(); }
    std::vector<reco::JTATagInfo> makeBaseVector(const edm::Event& iEvent) {
      edm::Handle<reco::JetTracksAssociationCollection> jetTracksAssociation;
      iEvent.getByToken(token_associator, jetTracksAssociation);
      std::vector<reco::JTATagInfo> bases;
      size_t i = 0;
      for (reco::JetTracksAssociationCollection::const_iterator it = jetTracksAssociation->begin();
           it != jetTracksAssociation->end();
           it++, i++) {
        edm::Ref<reco::JetTracksAssociationCollection> jtaRef(jetTracksAssociation, i);
        bases.push_back(reco::JTATagInfo(jtaRef));
      }
      return bases;
    }

    edm::EDGetTokenT<reco::JetTracksAssociationCollection> token_associator;
  };
  class FromJetAndCands {
  public:
    FromJetAndCands(const edm::ParameterSet& iConfig, edm::ConsumesCollector&& iC, const std::string& jets = "jets")
        : token_jets(iC.consumes<edm::View<reco::Jet> >(iConfig.getParameter<edm::InputTag>(jets))),
          token_cands(iC.consumes<edm::View<reco::Candidate> >(iConfig.getParameter<edm::InputTag>("candidates"))),
          maxDeltaR(iConfig.getParameter<double>("maxDeltaR")),
          explicitJTA(iConfig.existsAs<bool>("explicitJTA") ? iConfig.getParameter<bool>("explicitJTA") : false) {}

    const std::vector<reco::CandidatePtr>& tracks(const reco::JetTagInfo& it) { return m_map[it.jet().key()]; }
    std::vector<reco::JetTagInfo> makeBaseVector(const edm::Event& iEvent) {
      edm::Handle<edm::View<reco::Jet> > jets;
      iEvent.getByToken(token_jets, jets);
      std::vector<reco::JetTagInfo> bases;

      edm::Handle<edm::View<reco::Candidate> > cands;
      iEvent.getByToken(token_cands, cands);
      m_map.clear();
      m_map.resize(jets->size());
      double maxDeltaR2 = maxDeltaR * maxDeltaR;
      size_t i = 0;
      for (edm::View<reco::Jet>::const_iterator it = jets->begin(); it != jets->end(); it++, i++) {
        edm::RefToBase<reco::Jet> jRef(jets, i);
        bases.push_back(reco::JetTagInfo(jRef));
        if (explicitJTA) {
          for (size_t j = 0; j < it->numberOfDaughters(); ++j) {
            if (it->daughterPtr(j)->bestTrack() != nullptr && it->daughterPtr(j)->charge() != 0) {
              m_map[i].push_back(it->daughterPtr(j));
            }
          }
        } else {
          for (size_t j = 0; j < cands->size(); ++j) {
            if ((*cands)[j].bestTrack() != nullptr && (*cands)[j].charge() != 0 && (*cands)[j].pt() > 0 &&
                Geom::deltaR2((*cands)[j], (*jets)[i]) < maxDeltaR2) {
              m_map[i].push_back(cands->ptrAt(j));
            }
          }
        }
      }
      return bases;
    }
    std::vector<std::vector<reco::CandidatePtr> > m_map;
    edm::EDGetTokenT<edm::View<reco::Jet> > token_jets;
    edm::EDGetTokenT<edm::View<reco::Candidate> > token_cands;
    double maxDeltaR;
    bool explicitJTA;
  };
}  // namespace IPProducerHelpers
template <class Container, class Base, class Helper>
class IPProducer : public edm::stream::EDProducer<> {
public:
  typedef std::vector<reco::IPTagInfo<Container, Base> > Product;

  explicit IPProducer(const edm::ParameterSet&);
  ~IPProducer() override;
  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);

  void produce(edm::Event&, const edm::EventSetup&) override;

private:
  void checkEventSetup(const edm::EventSetup& iSetup);

  const edm::ParameterSet& m_config;
  edm::EDGetTokenT<reco::VertexCollection> token_primaryVertex;
  edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> token_trackBuilder;
  edm::ESGetToken<TrackProbabilityCalibration, BTagTrackProbability2DRcd> token_calib2D;
  edm::ESGetToken<TrackProbabilityCalibration, BTagTrackProbability3DRcd> token_calib3D;

  bool m_computeProbabilities;
  bool m_computeGhostTrack;
  double m_ghostTrackPriorDeltaR;
  std::unique_ptr<HistogramProbabilityEstimator> m_probabilityEstimator;
  unsigned long long m_calibrationCacheId2D;
  unsigned long long m_calibrationCacheId3D;
  bool m_useDB;

  int m_cutPixelHits;
  int m_cutTotalHits;
  double m_cutMaxTIP;
  double m_cutMinPt;
  double m_cutMaxChiSquared;
  double m_cutMaxLIP;
  bool m_directionWithTracks;
  bool m_directionWithGhostTrack;
  bool m_useTrackQuality;
  Helper m_helper;
};

// -*- C++ -*-
//
// Package:    IPProducer
// Class:      IPProducer
//
//
// Original Author:  Andrea Rizzi
//         Created:  Thu Apr  6 09:56:23 CEST 2006
//
//

//
// constructors and destructor
//
template <class Container, class Base, class Helper>
IPProducer<Container, Base, Helper>::IPProducer(const edm::ParameterSet& iConfig)
    : m_config(iConfig), m_helper(iConfig, consumesCollector()) {
  m_calibrationCacheId3D = 0;
  m_calibrationCacheId2D = 0;

  token_primaryVertex = consumes<reco::VertexCollection>(m_config.getParameter<edm::InputTag>("primaryVertex"));
  token_trackBuilder =
      esConsumes<TransientTrackBuilder, TransientTrackRecord>(edm::ESInputTag("", "TransientTrackBuilder"));
  token_calib2D = esConsumes<TrackProbabilityCalibration, BTagTrackProbability2DRcd>();
  token_calib3D = esConsumes<TrackProbabilityCalibration, BTagTrackProbability3DRcd>();

  m_computeProbabilities = m_config.getParameter<bool>("computeProbabilities");
  m_computeGhostTrack = m_config.getParameter<bool>("computeGhostTrack");
  m_ghostTrackPriorDeltaR = m_config.getParameter<double>("ghostTrackPriorDeltaR");
  m_cutPixelHits = m_config.getParameter<int>("minimumNumberOfPixelHits");
  m_cutTotalHits = m_config.getParameter<int>("minimumNumberOfHits");
  m_cutMaxTIP = m_config.getParameter<double>("maximumTransverseImpactParameter");
  m_cutMinPt = m_config.getParameter<double>("minimumTransverseMomentum");
  m_cutMaxChiSquared = m_config.getParameter<double>("maximumChiSquared");
  m_cutMaxLIP = m_config.getParameter<double>("maximumLongitudinalImpactParameter");
  m_directionWithTracks = m_config.getParameter<bool>("jetDirectionUsingTracks");
  m_directionWithGhostTrack = m_config.getParameter<bool>("jetDirectionUsingGhostTrack");
  m_useTrackQuality = m_config.getParameter<bool>("useTrackQuality");

  if (m_computeGhostTrack)
    produces<reco::TrackCollection>("ghostTracks");
  produces<Product>();
}

template <class Container, class Base, class Helper>
IPProducer<Container, Base, Helper>::~IPProducer() {}

//
// member functions
//
// ------------ method called to produce the data  ------------
template <class Container, class Base, class Helper>
void IPProducer<Container, Base, Helper>::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
  // Update probability estimator if event setup is changed
  if (m_computeProbabilities)
    checkEventSetup(iSetup);

  edm::Handle<reco::VertexCollection> primaryVertex;
  iEvent.getByToken(token_primaryVertex, primaryVertex);

  edm::ESHandle<TransientTrackBuilder> builder = iSetup.getHandle(token_trackBuilder);
  // m_algo.setTransientTrackBuilder(builder.product());

  // output collections
  auto result = std::make_unique<Product>();

  std::unique_ptr<reco::TrackCollection> ghostTracks;
  reco::TrackRefProd ghostTrackRefProd;
  if (m_computeGhostTrack) {
    ghostTracks = std::make_unique<reco::TrackCollection>();
    ghostTrackRefProd = iEvent.getRefBeforePut<reco::TrackCollection>("ghostTracks");
  }

  // use first pv of the collection
  reco::Vertex dummy;
  const reco::Vertex* pv = &dummy;
  edm::Ref<reco::VertexCollection> pvRef;
  if (!primaryVertex->empty()) {
    pv = &*primaryVertex->begin();
    // we always use the first vertex (at the moment)
    pvRef = edm::Ref<reco::VertexCollection>(primaryVertex, 0);
  } else {  // create a dummy PV
    reco::Vertex::Error e;
    e(0, 0) = 0.0015 * 0.0015;
    e(1, 1) = 0.0015 * 0.0015;
    e(2, 2) = 15. * 15.;
    reco::Vertex::Point p(0, 0, 0);
    dummy = reco::Vertex(p, e, 0, 0, 0);
  }

  std::vector<Base> baseTagInfos = m_helper.makeBaseVector(iEvent);
  for (typename std::vector<Base>::const_iterator it = baseTagInfos.begin(); it != baseTagInfos.end(); it++) {
    Container tracks = m_helper.tracks(*it);
    math::XYZVector jetMomentum = it->jet()->momentum();

    if (m_directionWithTracks) {
      jetMomentum *= 0.5;
      for (typename Container::const_iterator itTrack = tracks.begin(); itTrack != tracks.end(); ++itTrack)
        if (reco::btag::toTrack(*itTrack)->numberOfValidHits() >= m_cutTotalHits)  //minimal quality cuts
          jetMomentum += (*itTrack)->momentum();
    }

    Container selectedTracks;
    std::vector<reco::TransientTrack> transientTracks;

    for (typename Container::const_iterator itTrack = tracks.begin(); itTrack != tracks.end(); ++itTrack) {
      reco::TransientTrack transientTrack = builder->build(*itTrack);
      const reco::Track& track = transientTrack.track();  //**itTrack;
                                                          /*    cout << " pt " <<  track.pt() <<
               " d0 " <<  fabs(track.d0()) <<
               " #hit " <<    track.hitPattern().numberOfValidHits()<<
               " ipZ " <<   fabs(track.dz()-pv->z())<<
               " chi2 " <<  track.normalizedChi2()<<
               " #pixel " <<    track.hitPattern().numberOfValidPixelHits()<< endl;
*/
      if (track.pt() > m_cutMinPt &&
          track.hitPattern().numberOfValidHits() >= m_cutTotalHits &&  // min num tracker hits
          track.hitPattern().numberOfValidPixelHits() >= m_cutPixelHits &&
          track.normalizedChi2() < m_cutMaxChiSquared && std::abs(track.dxy(pv->position())) < m_cutMaxTIP &&
          std::abs(track.dz(pv->position())) < m_cutMaxLIP) {
        //   std::cout << "selected" << std::endl;
        selectedTracks.push_back(*itTrack);
        transientTracks.push_back(transientTrack);
      }
    }
    //  std::cout <<"SIZE: " << transientTracks.size() << std::endl;
    GlobalVector direction(jetMomentum.x(), jetMomentum.y(), jetMomentum.z());

    std::unique_ptr<reco::GhostTrack> ghostTrack;
    reco::TrackRef ghostTrackRef;
    if (m_computeGhostTrack) {
      reco::GhostTrackFitter fitter;
      GlobalPoint origin = RecoVertex::convertPos(pv->position());
      GlobalError error = RecoVertex::convertError(pv->error());
      ghostTrack = std::make_unique<reco::GhostTrack>(
          fitter.fit(origin, error, direction, m_ghostTrackPriorDeltaR, transientTracks));

      /*
    if (std::sqrt(jetMomentum.Perp2()) > 30) {
        double offset = ghostTrack->prediction().lambda(origin);
        std::cout << "------------------ jet pt " << std::sqrt(jetMomentum.Perp2()) << std::endl;
        const std::vector<GhostTrackState> *states = &ghostTrack->states();
        for(std::vector<GhostTrackState>::const_iterator state = states->begin();
            state != states->end(); ++state) {
            double dist = state->lambda() - offset;
            double err = state->lambdaError(ghostTrack->prediction(), error);
            double ipSig = IPTools::signedImpactParameter3D(state->track(), direction, *pv).second.significance();
            double axisDist = state->axisDistance(ghostTrack->prediction());
            std::cout << state->track().impactPointState().freeState()->momentum().perp()
                      << ": " << dist << "/" << err << " [" << (dist / err) << "], ipsig = " << ipSig << ", dist = " << axisDist << ", w = " << state->weight() << std::endl;
        }
    }
*/
      ghostTrackRef = reco::TrackRef(ghostTrackRefProd, ghostTracks->size());
      ghostTracks->push_back(*ghostTrack);

      if (m_directionWithGhostTrack) {
        const reco::GhostTrackPrediction& pred = ghostTrack->prediction();
        double lambda = pred.lambda(origin);
        dummy = reco::Vertex(RecoVertex::convertPos(pred.position(lambda)),
                             RecoVertex::convertError(pred.positionError(lambda)),
                             0,
                             0,
                             0);
        pv = &dummy;
        direction = pred.direction();
      }
    }

    std::vector<float> prob2D, prob3D;
    std::vector<reco::btag::TrackIPData> ipData;

    for (unsigned int ind = 0; ind < transientTracks.size(); ind++) {
      const reco::TransientTrack& transientTrack = transientTracks[ind];
      const reco::Track& track = transientTrack.track();

      reco::btag::TrackIPData trackIP;
      trackIP.ip3d = IPTools::signedImpactParameter3D(transientTrack, direction, *pv).second;
      trackIP.ip2d = IPTools::signedTransverseImpactParameter(transientTrack, direction, *pv).second;

      TrajectoryStateOnSurface closest =
          IPTools::closestApproachToJet(transientTrack.impactPointState(), *pv, direction, transientTrack.field());
      if (closest.isValid())
        trackIP.closestToJetAxis = closest.globalPosition();

      // TODO: cross check if it is the same using other methods
      trackIP.distanceToJetAxis = IPTools::jetTrackDistance(transientTrack, direction, *pv).second;

      if (ghostTrack.get()) {
        const std::vector<reco::GhostTrackState>& states = ghostTrack->states();
        std::vector<reco::GhostTrackState>::const_iterator pos =
            std::find_if(states.begin(), states.end(), [&](auto& arg) { return arg.track() == transientTrack; });

        if (pos != states.end() && pos->isValid()) {
          VertexDistance3D dist;
          const reco::GhostTrackPrediction& pred = ghostTrack->prediction();
          GlobalPoint p1 = pos->tsos().globalPosition();
          GlobalError e1 = pos->tsos().cartesianError().position();
          GlobalPoint p2 = pred.position(pos->lambda());
          GlobalError e2 = pred.positionError(pos->lambda());
          trackIP.closestToGhostTrack = p1;
          trackIP.distanceToGhostTrack = dist.distance(VertexState(p1, e1), VertexState(p2, e2));
          trackIP.ghostTrackWeight = pos->weight();
        } else {
          trackIP.distanceToGhostTrack = Measurement1D(0., -1.);
          trackIP.ghostTrackWeight = 0.;
        }
      } else {
        trackIP.distanceToGhostTrack = Measurement1D(0., -1.);
        trackIP.ghostTrackWeight = 1.;
      }

      ipData.push_back(trackIP);

      if (m_computeProbabilities) {
        //probability with 3D ip
        std::pair<bool, double> probability = m_probabilityEstimator->probability(
            m_useTrackQuality, 0, ipData.back().ip3d.significance(), track, *(it->jet()), *pv);
        prob3D.push_back(probability.first ? probability.second : -1.);

        //probability with 2D ip
        probability = m_probabilityEstimator->probability(
            m_useTrackQuality, 1, ipData.back().ip2d.significance(), track, *(it->jet()), *pv);
        prob2D.push_back(probability.first ? probability.second : -1.);
      }
    }

    result->push_back(
        typename Product::value_type(ipData, prob2D, prob3D, selectedTracks, *it, pvRef, direction, ghostTrackRef));
  }

  if (m_computeGhostTrack)
    iEvent.put(std::move(ghostTracks), "ghostTracks");
  iEvent.put(std::move(result));
}

#include "CondFormats/BTauObjects/interface/TrackProbabilityCalibration.h"
#include "CondFormats/DataRecord/interface/BTagTrackProbability2DRcd.h"
#include "CondFormats/DataRecord/interface/BTagTrackProbability3DRcd.h"
#include "FWCore/Framework/interface/EventSetupRecord.h"
#include "FWCore/Framework/interface/EventSetupRecordImplementation.h"
#include "FWCore/Framework/interface/EventSetupRecordKey.h"

template <class Container, class Base, class Helper>
void IPProducer<Container, Base, Helper>::checkEventSetup(const edm::EventSetup& iSetup) {
  const edm::eventsetup::EventSetupRecord& re2D = iSetup.get<BTagTrackProbability2DRcd>();
  const edm::eventsetup::EventSetupRecord& re3D = iSetup.get<BTagTrackProbability3DRcd>();
  unsigned long long cacheId2D = re2D.cacheIdentifier();
  unsigned long long cacheId3D = re3D.cacheIdentifier();

  if (cacheId2D != m_calibrationCacheId2D || cacheId3D != m_calibrationCacheId3D)  //Calibration changed
  {
    //iSetup.get<BTagTrackProbabilityRcd>().get(calib);
    edm::ESHandle<TrackProbabilityCalibration> calib2DHandle = iSetup.getHandle(token_calib2D);
    edm::ESHandle<TrackProbabilityCalibration> calib3DHandle = iSetup.getHandle(token_calib3D);

    const TrackProbabilityCalibration* ca2D = calib2DHandle.product();
    const TrackProbabilityCalibration* ca3D = calib3DHandle.product();

    m_probabilityEstimator = std::make_unique<HistogramProbabilityEstimator>(ca3D, ca2D);
  }
  m_calibrationCacheId3D = cacheId3D;
  m_calibrationCacheId2D = cacheId2D;
}

// Specialized templates used to fill 'descriptions'
// ------------ method fills 'descriptions' with the allowed parameters for the module ------------
template <>
inline void IPProducer<reco::TrackRefVector, reco::JTATagInfo, IPProducerHelpers::FromJTA>::fillDescriptions(
    edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  desc.add<double>("maximumTransverseImpactParameter", 0.2);
  desc.add<int>("minimumNumberOfHits", 8);
  desc.add<double>("minimumTransverseMomentum", 1.0);
  desc.add<edm::InputTag>("primaryVertex", edm::InputTag("offlinePrimaryVertices"));
  desc.add<double>("maximumLongitudinalImpactParameter", 17.0);
  desc.add<bool>("computeGhostTrack", true);
  desc.add<double>("ghostTrackPriorDeltaR", 0.03);
  desc.add<edm::InputTag>("jetTracks", edm::InputTag("ak4JetTracksAssociatorAtVertexPF"));
  desc.add<bool>("jetDirectionUsingGhostTrack", false);
  desc.add<int>("minimumNumberOfPixelHits", 2);
  desc.add<bool>("jetDirectionUsingTracks", false);
  desc.add<bool>("computeProbabilities", true);
  desc.add<bool>("useTrackQuality", false);
  desc.add<double>("maximumChiSquared", 5.0);
  descriptions.addDefault(desc);
}

template <>
inline void
IPProducer<std::vector<reco::CandidatePtr>, reco::JetTagInfo, IPProducerHelpers::FromJetAndCands>::fillDescriptions(
    edm::ConfigurationDescriptions& descriptions) {
  edm::ParameterSetDescription desc;
  desc.add<double>("maximumTransverseImpactParameter", 0.2);
  desc.add<int>("minimumNumberOfHits", 8);
  desc.add<double>("minimumTransverseMomentum", 1.0);
  desc.add<edm::InputTag>("primaryVertex", edm::InputTag("offlinePrimaryVertices"));
  desc.add<double>("maximumLongitudinalImpactParameter", 17.0);
  desc.add<bool>("computeGhostTrack", true);
  desc.add<double>("maxDeltaR", 0.4);
  desc.add<edm::InputTag>("candidates", edm::InputTag("particleFlow"));
  desc.add<bool>("jetDirectionUsingGhostTrack", false);
  desc.add<int>("minimumNumberOfPixelHits", 2);
  desc.add<bool>("jetDirectionUsingTracks", false);
  desc.add<bool>("computeProbabilities", true);
  desc.add<bool>("useTrackQuality", false);
  desc.add<edm::InputTag>("jets", edm::InputTag("ak4PFJetsCHS"));
  desc.add<double>("ghostTrackPriorDeltaR", 0.03);
  desc.add<double>("maximumChiSquared", 5.0);
  desc.addOptional<bool>("explicitJTA", false);
  descriptions.addDefault(desc);
}

#endif