PATIsolatedTrackProducer.cc

Go to the documentation of this file.

#include <string>

#include "FWCore/Framework/interface/stream/EDProducer.h"
#include "FWCore/Framework/interface/ESHandle.h"
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/EventSetup.h"
#include "FWCore/Framework/interface/Frameworkfwd.h"
#include "FWCore/ParameterSet/interface/ParameterSet.h"
#include "FWCore/Utilities/interface/Exception.h"

#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"

#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
#include "DataFormats/PatCandidates/interface/IsolatedTrack.h"
#include "DataFormats/PatCandidates/interface/PFIsolation.h"
#include "DataFormats/Candidate/interface/CandidateFwd.h"
#include "DataFormats/Common/interface/View.h"
#include "DataFormats/Common/interface/RefToPtr.h"
#include "DataFormats/Math/interface/deltaR.h"
#include "DataFormats/JetReco/interface/CaloJet.h"
#include "DataFormats/TrackReco/interface/Track.h"
#include "DataFormats/TrackReco/interface/TrackExtraFwd.h"
#include "DataFormats/TrackReco/interface/DeDxData.h"
#include "DataFormats/TrackReco/interface/DeDxHitInfo.h"
#include "DataFormats/VertexReco/interface/Vertex.h"
#include "DataFormats/VertexReco/interface/VertexFwd.h"
#include "RecoTracker/DeDx/interface/DeDxTools.h"
#include "CommonTools/Utils/interface/StringCutObjectSelector.h"

#include "TrackingTools/TrackAssociator/interface/TrackDetectorAssociator.h"
#include "TrackingTools/TrackAssociator/interface/TrackAssociatorParameters.h"
#include "TrackingTools/TrajectoryState/interface/TrajectoryStateTransform.h"

#include "CondFormats/HcalObjects/interface/HcalChannelQuality.h"
#include "CondFormats/HcalObjects/interface/HcalChannelStatus.h"
#include "CondFormats/EcalObjects/interface/EcalChannelStatus.h"
#include "CondFormats/DataRecord/interface/HcalChannelQualityRcd.h"
#include "CondFormats/DataRecord/interface/EcalChannelStatusRcd.h"

#include "MagneticField/Engine/interface/MagneticField.h"

namespace pat {

  class PATIsolatedTrackProducer : public edm::stream::EDProducer<> {
  public:
    typedef pat::IsolatedTrack::PolarLorentzVector PolarLorentzVector;
    typedef pat::IsolatedTrack::LorentzVector LorentzVector;

    explicit PATIsolatedTrackProducer(const edm::ParameterSet&);
    ~PATIsolatedTrackProducer() override;

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

    // compute iso/miniiso
    void getIsolation(const PolarLorentzVector& p4,
                      const pat::PackedCandidateCollection* pc,
                      int pc_idx,
                      pat::PFIsolation& iso,
                      pat::PFIsolation& miniiso) const;

    bool getPFLeptonOverlap(const PolarLorentzVector& p4, const pat::PackedCandidateCollection* pc) const;

    float getPFNeutralSum(const PolarLorentzVector& p4, const pat::PackedCandidateCollection* pc, int pc_idx) const;

    void getNearestPCRef(const PolarLorentzVector& p4,
                         const pat::PackedCandidateCollection* pc,
                         int pc_idx,
                         int& pc_ref_idx) const;

    float getDeDx(const reco::DeDxHitInfo* hitInfo, bool doPixel, bool doStrip) const;

    TrackDetMatchInfo getTrackDetMatchInfo(const edm::Event&, const edm::EventSetup&, const reco::Track&);

    void getCaloJetEnergy(const PolarLorentzVector&, const reco::CaloJetCollection*, float&, float&) const;

  private:
    const edm::EDGetTokenT<pat::PackedCandidateCollection> pc_;
    const edm::EDGetTokenT<pat::PackedCandidateCollection> lt_;
    const edm::EDGetTokenT<reco::TrackCollection> gt_;
    const edm::EDGetTokenT<reco::VertexCollection> pv_;
    const edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>> gt2pc_;
    const edm::EDGetTokenT<edm::Association<pat::PackedCandidateCollection>> gt2lt_;
    const edm::EDGetTokenT<edm::Association<reco::PFCandidateCollection>> pc2pf_;
    const edm::EDGetTokenT<reco::CaloJetCollection> caloJets_;
    const edm::EDGetTokenT<edm::ValueMap<reco::DeDxData>> gt2dedxStrip_;
    const edm::EDGetTokenT<edm::ValueMap<reco::DeDxData>> gt2dedxPixel_;
    const edm::EDGetTokenT<reco::DeDxHitInfoAss> gt2dedxHitInfo_;
    const edm::ESGetToken<HcalChannelQuality, HcalChannelQualityRcd> hcalQToken_;
    const edm::ESGetToken<EcalChannelStatus, EcalChannelStatusRcd> ecalSToken_;
    const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> bFieldToken_;
    const bool addPrescaledDeDxTracks_;
    const edm::EDGetTokenT<edm::ValueMap<int>> gt2dedxHitInfoPrescale_;
    const bool usePrecomputedDeDxStrip_;
    const bool usePrecomputedDeDxPixel_;
    const float pT_cut_;          // only save cands with pT>pT_cut_
    const float pT_cut_noIso_;    // above this pT, don't apply any iso cut
    const float pfIsolation_DR_;  // isolation radius
    const float pfIsolation_DZ_;  // used in determining if pfcand is from PV or PU
    const float absIso_cut_;      // save if ANY of absIso, relIso, or miniRelIso pass the cuts
    const float relIso_cut_;
    const float miniRelIso_cut_;
    const float caloJet_DR_;          // save energy of nearest calojet within caloJet_DR_
    const float pflepoverlap_DR_;     // pf lepton overlap radius
    const float pflepoverlap_pTmin_;  // pf lepton overlap min pT (only look at PF candidates with pT>pflepoverlap_pTmin_)
    const float pcRefNearest_DR_;     // radius for nearest charged packed candidate
    const float pcRefNearest_pTmin_;  // min pT for nearest charged packed candidate
    const float pfneutralsum_DR_;     // pf lepton overlap radius
    const bool useHighPurity_;
    const bool saveDeDxHitInfo_;
    StringCutObjectSelector<pat::IsolatedTrack> saveDeDxHitInfoCut_;

    std::vector<double> miniIsoParams_;

    TrackDetectorAssociator trackAssociator_;
    TrackAssociatorParameters trackAssocParameters_;
  };
}  // namespace pat

pat::PATIsolatedTrackProducer::PATIsolatedTrackProducer(const edm::ParameterSet& iConfig)
    : pc_(consumes<pat::PackedCandidateCollection>(iConfig.getParameter<edm::InputTag>("packedPFCandidates"))),
      lt_(consumes<pat::PackedCandidateCollection>(iConfig.getParameter<edm::InputTag>("lostTracks"))),
      gt_(consumes<reco::TrackCollection>(iConfig.getParameter<edm::InputTag>("generalTracks"))),
      pv_(consumes<reco::VertexCollection>(iConfig.getParameter<edm::InputTag>("primaryVertices"))),
      gt2pc_(consumes<edm::Association<pat::PackedCandidateCollection>>(
          iConfig.getParameter<edm::InputTag>("packedPFCandidates"))),
      gt2lt_(consumes<edm::Association<pat::PackedCandidateCollection>>(
          iConfig.getParameter<edm::InputTag>("lostTracks"))),
      pc2pf_(consumes<edm::Association<reco::PFCandidateCollection>>(
          iConfig.getParameter<edm::InputTag>("packedPFCandidates"))),
      caloJets_(consumes<reco::CaloJetCollection>(iConfig.getParameter<edm::InputTag>("caloJets"))),
      gt2dedxStrip_(consumes<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<edm::InputTag>("dEdxDataStrip"))),
      gt2dedxPixel_(consumes<edm::ValueMap<reco::DeDxData>>(iConfig.getParameter<edm::InputTag>("dEdxDataPixel"))),
      gt2dedxHitInfo_(consumes<reco::DeDxHitInfoAss>(iConfig.getParameter<edm::InputTag>("dEdxHitInfo"))),
      hcalQToken_(esConsumes(edm::ESInputTag("", "withTopo"))),
      ecalSToken_(esConsumes()),
      bFieldToken_(esConsumes()),
      addPrescaledDeDxTracks_(iConfig.getParameter<bool>("addPrescaledDeDxTracks")),
      gt2dedxHitInfoPrescale_(addPrescaledDeDxTracks_ ? consumes<edm::ValueMap<int>>(
                                                            iConfig.getParameter<edm::InputTag>("dEdxHitInfoPrescale"))
                                                      : edm::EDGetTokenT<edm::ValueMap<int>>()),
      usePrecomputedDeDxStrip_(iConfig.getParameter<bool>("usePrecomputedDeDxStrip")),
      usePrecomputedDeDxPixel_(iConfig.getParameter<bool>("usePrecomputedDeDxPixel")),
      pT_cut_(iConfig.getParameter<double>("pT_cut")),
      pT_cut_noIso_(iConfig.getParameter<double>("pT_cut_noIso")),
      pfIsolation_DR_(iConfig.getParameter<double>("pfIsolation_DR")),
      pfIsolation_DZ_(iConfig.getParameter<double>("pfIsolation_DZ")),
      absIso_cut_(iConfig.getParameter<double>("absIso_cut")),
      relIso_cut_(iConfig.getParameter<double>("relIso_cut")),
      miniRelIso_cut_(iConfig.getParameter<double>("miniRelIso_cut")),
      caloJet_DR_(iConfig.getParameter<double>("caloJet_DR")),
      pflepoverlap_DR_(iConfig.getParameter<double>("pflepoverlap_DR")),
      pflepoverlap_pTmin_(iConfig.getParameter<double>("pflepoverlap_pTmin")),
      pcRefNearest_DR_(iConfig.getParameter<double>("pcRefNearest_DR")),
      pcRefNearest_pTmin_(iConfig.getParameter<double>("pcRefNearest_pTmin")),
      pfneutralsum_DR_(iConfig.getParameter<double>("pfneutralsum_DR")),
      useHighPurity_(iConfig.getParameter<bool>("useHighPurity")),
      saveDeDxHitInfo_(iConfig.getParameter<bool>("saveDeDxHitInfo")),
      saveDeDxHitInfoCut_(iConfig.getParameter<std::string>("saveDeDxHitInfoCut")) {
  // TrackAssociator parameters
  edm::ParameterSet parameters = iConfig.getParameter<edm::ParameterSet>("TrackAssociatorParameters");
  edm::ConsumesCollector iC = consumesCollector();
  trackAssocParameters_.loadParameters(parameters, iC);

  trackAssociator_.useDefaultPropagator();

  miniIsoParams_ = iConfig.getParameter<std::vector<double>>("miniIsoParams");
  if (miniIsoParams_.size() != 3)
    throw cms::Exception("ParameterError") << "miniIsoParams must have exactly 3 elements.\n";

  produces<pat::IsolatedTrackCollection>();

  if (saveDeDxHitInfo_) {
    produces<reco::DeDxHitInfoCollection>();
    produces<reco::DeDxHitInfoAss>();
  }
}

pat::PATIsolatedTrackProducer::~PATIsolatedTrackProducer() {}

void pat::PATIsolatedTrackProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
  // packedPFCandidate collection
  edm::Handle<pat::PackedCandidateCollection> pc_h;
  iEvent.getByToken(pc_, pc_h);
  const pat::PackedCandidateCollection* pc = pc_h.product();

  // lostTracks collection
  edm::Handle<pat::PackedCandidateCollection> lt_h;
  iEvent.getByToken(lt_, lt_h);
  const pat::PackedCandidateCollection* lt = lt_h.product();

  // generalTracks collection
  edm::Handle<reco::TrackCollection> gt_h;
  iEvent.getByToken(gt_, gt_h);
  const reco::TrackCollection* generalTracks = gt_h.product();

  // get the primary vertex
  edm::Handle<reco::VertexCollection> pvs;
  iEvent.getByToken(pv_, pvs);
  const reco::Vertex& pv = (*pvs)[0];

  // generalTracks-->packedPFCandidate association
  edm::Handle<edm::Association<pat::PackedCandidateCollection>> gt2pc;
  iEvent.getByToken(gt2pc_, gt2pc);

  // generalTracks-->lostTracks association
  edm::Handle<edm::Association<pat::PackedCandidateCollection>> gt2lt;
  iEvent.getByToken(gt2lt_, gt2lt);

  // packedPFCandidates-->particleFlow(reco::PFCandidate) association
  edm::Handle<edm::Association<reco::PFCandidateCollection>> pc2pf;
  iEvent.getByToken(pc2pf_, pc2pf);

  edm::Handle<reco::CaloJetCollection> caloJets;
  iEvent.getByToken(caloJets_, caloJets);

  // associate generalTracks with their DeDx data (estimator for strip dE/dx)
  edm::Handle<edm::ValueMap<reco::DeDxData>> gt2dedxStrip;
  iEvent.getByToken(gt2dedxStrip_, gt2dedxStrip);

  // associate generalTracks with their DeDx data (estimator for pixel dE/dx)
  edm::Handle<edm::ValueMap<reco::DeDxData>> gt2dedxPixel;
  iEvent.getByToken(gt2dedxPixel_, gt2dedxPixel);

  // associate generalTracks with their DeDx hit info (used to estimate pixel dE/dx)
  edm::Handle<reco::DeDxHitInfoAss> gt2dedxHitInfo;
  iEvent.getByToken(gt2dedxHitInfo_, gt2dedxHitInfo);
  edm::Handle<edm::ValueMap<int>> gt2dedxHitInfoPrescale;
  if (addPrescaledDeDxTracks_) {
    iEvent.getByToken(gt2dedxHitInfoPrescale_, gt2dedxHitInfoPrescale);
  }

  const HcalChannelQuality* hcalQ = &iSetup.getData(hcalQToken_);

  const EcalChannelStatus* ecalS = &iSetup.getData(ecalSToken_);

  auto outDeDxC = std::make_unique<reco::DeDxHitInfoCollection>();
  std::vector<int> dEdXass;

  auto outPtrP = std::make_unique<std::vector<pat::IsolatedTrack>>();

  //add general tracks
  for (unsigned int igt = 0; igt < generalTracks->size(); igt++) {
    const reco::Track& gentk = (*gt_h)[igt];
    if (useHighPurity_)
      if (!(gentk.quality(reco::TrackBase::qualityByName("highPurity"))))
        continue;
    reco::TrackRef tkref = reco::TrackRef(gt_h, igt);
    pat::PackedCandidateRef pcref = (*gt2pc)[tkref];
    pat::PackedCandidateRef ltref = (*gt2lt)[tkref];
    const pat::PackedCandidate* pfCand = pcref.get();
    const pat::PackedCandidate* lostTrack = ltref.get();

    // Determine if this general track is associated with anything in packedPFCandidates or lostTracks
    // Sometimes, a track gets associated w/ a neutral pfCand.
    // In this case, ignore the pfCand and take from lostTracks
    bool isInPackedCands = (pcref.isNonnull() && pcref.id() == pc_h.id() && pfCand->charge() != 0);
    bool isInLostTracks = (ltref.isNonnull() && ltref.id() == lt_h.id());

    PolarLorentzVector polarP4;
    LorentzVector p4;
    pat::PackedCandidateRef refToCand;
    int pdgId, charge, fromPV;
    float dz, dxy, dzError, dxyError;
    int pfCandInd;  //to avoid counting packedPFCands in their own isolation
    int ltCandInd;  //to avoid pointing lost track to itself when looking for closest

    // get the four-momentum and charge
    if (isInPackedCands) {
      p4 = pfCand->p4();
      polarP4 = pfCand->p4();
      charge = pfCand->charge();
      pfCandInd = pcref.key();
      ltCandInd = -1;
    } else if (isInLostTracks) {
      p4 = lostTrack->p4();
      polarP4 = lostTrack->p4();
      charge = lostTrack->charge();
      pfCandInd = -1;
      ltCandInd = ltref.key();
    } else {
      double m = 0.13957018;  //assume pion mass
      double E = sqrt(m * m + gentk.p() * gentk.p());
      p4.SetPxPyPzE(gentk.px(), gentk.py(), gentk.pz(), E);
      polarP4.SetCoordinates(gentk.pt(), gentk.eta(), gentk.phi(), m);
      charge = gentk.charge();
      pfCandInd = -1;
      ltCandInd = -1;
    }

    int prescaled = 0;
    if (addPrescaledDeDxTracks_) {
      const auto& dedxRef = (*gt2dedxHitInfo)[tkref];
      if (dedxRef.isNonnull()) {
        prescaled = (*gt2dedxHitInfoPrescale)[dedxRef];
      }
    }

    if (polarP4.pt() < pT_cut_ && prescaled <= 1)
      continue;
    if (charge == 0)
      continue;

    // get the isolation of the track
    pat::PFIsolation isolationDR03;
    pat::PFIsolation miniIso;
    getIsolation(polarP4, pc, pfCandInd, isolationDR03, miniIso);

    // isolation cut
    if (polarP4.pt() < pT_cut_noIso_ && prescaled <= 1 &&
        !(isolationDR03.chargedHadronIso() < absIso_cut_ ||
          isolationDR03.chargedHadronIso() / polarP4.pt() < relIso_cut_ ||
          miniIso.chargedHadronIso() / polarP4.pt() < miniRelIso_cut_))
      continue;

    // get the rest after the pt/iso cuts. Saves some runtime
    if (isInPackedCands) {
      pdgId = pfCand->pdgId();
      dz = pfCand->dz();
      dxy = pfCand->dxy();
      dzError = pfCand->hasTrackDetails() ? pfCand->dzError() : gentk.dzError();
      dxyError = pfCand->hasTrackDetails() ? pfCand->dxyError() : gentk.dxyError();
      fromPV = pfCand->fromPV();
      refToCand = pcref;
    } else if (isInLostTracks) {
      pdgId = lostTrack->pdgId();
      dz = lostTrack->dz();
      dxy = lostTrack->dxy();
      dzError = lostTrack->hasTrackDetails() ? lostTrack->dzError() : gentk.dzError();
      dxyError = lostTrack->hasTrackDetails() ? lostTrack->dxyError() : gentk.dxyError();
      fromPV = lostTrack->fromPV();
      refToCand = ltref;
    } else {
      pdgId = 0;
      dz = gentk.dz(pv.position());
      dxy = gentk.dxy(pv.position());
      dzError = gentk.dzError();
      dxyError = gentk.dxyError();
      fromPV = -1;
      refToCand = pat::PackedCandidateRef();  //NULL reference
    }

    float caloJetEm, caloJetHad;
    getCaloJetEnergy(polarP4, caloJets.product(), caloJetEm, caloJetHad);

    bool pfLepOverlap = getPFLeptonOverlap(polarP4, pc);
    float pfNeutralSum = getPFNeutralSum(polarP4, pc, pfCandInd);

    pat::PackedCandidateRef refToNearestPF = pat::PackedCandidateRef();
    int refToNearestPF_idx = -1;
    getNearestPCRef(polarP4, pc, pfCandInd, refToNearestPF_idx);
    if (refToNearestPF_idx != -1)
      refToNearestPF = pat::PackedCandidateRef(pc_h, refToNearestPF_idx);

    pat::PackedCandidateRef refToNearestLostTrack = pat::PackedCandidateRef();
    int refToNearestLostTrack_idx = -1;
    getNearestPCRef(polarP4, lt, ltCandInd, refToNearestLostTrack_idx);
    if (refToNearestLostTrack_idx != -1)
      refToNearestLostTrack = pat::PackedCandidateRef(lt_h, refToNearestLostTrack_idx);

    // if no dEdx info exists, just store -1
    float dEdxPixel = -1, dEdxStrip = -1;
    if (usePrecomputedDeDxStrip_ && gt2dedxStrip.isValid() && gt2dedxStrip->contains(tkref.id())) {
      dEdxStrip = (*gt2dedxStrip)[tkref].dEdx();
    } else if (gt2dedxHitInfo.isValid() && gt2dedxHitInfo->contains(tkref.id())) {
      const reco::DeDxHitInfo* hitInfo = (*gt2dedxHitInfo)[tkref].get();
      dEdxStrip = getDeDx(hitInfo, false, true);
    }
    if (usePrecomputedDeDxPixel_ && gt2dedxPixel.isValid() && gt2dedxPixel->contains(tkref.id())) {
      dEdxPixel = (*gt2dedxPixel)[tkref].dEdx();
    } else if (gt2dedxHitInfo.isValid() && gt2dedxHitInfo->contains(tkref.id())) {
      const reco::DeDxHitInfo* hitInfo = (*gt2dedxHitInfo)[tkref].get();
      dEdxPixel = getDeDx(hitInfo, true, false);
    }

    int trackQuality = gentk.qualityMask();

    // get the associated ecal/hcal detectors
    TrackDetMatchInfo trackDetInfo = getTrackDetMatchInfo(iEvent, iSetup, gentk);

    // fill ecal/hcal status vectors
    std::vector<uint32_t> crossedHcalStatus;
    for (auto const& did : trackDetInfo.crossedHcalIds) {
      crossedHcalStatus.push_back(hcalQ->getValues(did.rawId())->getValue());
    }
    std::vector<uint16_t> crossedEcalStatus;
    for (auto const& did : trackDetInfo.crossedEcalIds) {
      crossedEcalStatus.push_back(ecalS->find(did.rawId())->getStatusCode());
    }

    int deltaEta = int((trackDetInfo.trkGlobPosAtEcal.eta() - gentk.eta()) / 0.5 * 250);
    int deltaPhi = int((trackDetInfo.trkGlobPosAtEcal.phi() - gentk.phi()) / 0.5 * 250);
    if (deltaEta < -250)
      deltaEta = -250;
    if (deltaEta > 250)
      deltaEta = 250;
    if (deltaPhi < -250)
      deltaPhi = -250;
    if (deltaPhi > 250)
      deltaPhi = 250;

    outPtrP->push_back(pat::IsolatedTrack(isolationDR03,
                                          miniIso,
                                          caloJetEm,
                                          caloJetHad,
                                          pfLepOverlap,
                                          pfNeutralSum,
                                          p4,
                                          charge,
                                          pdgId,
                                          dz,
                                          dxy,
                                          dzError,
                                          dxyError,
                                          gentk.hitPattern(),
                                          dEdxStrip,
                                          dEdxPixel,
                                          fromPV,
                                          trackQuality,
                                          crossedEcalStatus,
                                          crossedHcalStatus,
                                          deltaEta,
                                          deltaPhi,
                                          refToCand,
                                          refToNearestPF,
                                          refToNearestLostTrack));
    outPtrP->back().setStatus(prescaled);

    if (saveDeDxHitInfo_) {
      const auto& dedxRef = (*gt2dedxHitInfo)[tkref];
      if (saveDeDxHitInfoCut_(outPtrP->back()) && dedxRef.isNonnull()) {
        outDeDxC->push_back(*dedxRef);
        dEdXass.push_back(outDeDxC->size() - 1);
      } else {
        dEdXass.push_back(-1);
      }
    }
  }

  // there are some number of pfcandidates with no associated track
  // (mostly electrons, with a handful of muons)
  // here we find these and store. Track-specific variables get some default values
  for (unsigned int ipc = 0; ipc < pc->size(); ipc++) {
    const pat::PackedCandidate& pfCand = pc->at(ipc);
    pat::PackedCandidateRef pcref = pat::PackedCandidateRef(pc_h, ipc);
    reco::PFCandidateRef pfref = (*pc2pf)[pcref];

    // already counted if it has a track reference in the generalTracks collection
    if (pfref.get()->trackRef().isNonnull() && pfref.get()->trackRef().id() == gt_h.id())
      continue;

    PolarLorentzVector polarP4;
    pat::PackedCandidateRef refToCand;
    int pdgId, charge, fromPV;
    float dz, dxy, dzError, dxyError;

    polarP4 = pfCand.polarP4();
    charge = pfCand.charge();

    if (polarP4.pt() < pT_cut_)
      continue;
    if (charge == 0)
      continue;

    // get the isolation of the track
    pat::PFIsolation isolationDR03;
    pat::PFIsolation miniIso;
    getIsolation(polarP4, pc, ipc, isolationDR03, miniIso);

    // isolation cut
    if (polarP4.pt() < pT_cut_noIso_ && !(isolationDR03.chargedHadronIso() < absIso_cut_ ||
                                          isolationDR03.chargedHadronIso() / polarP4.pt() < relIso_cut_ ||
                                          miniIso.chargedHadronIso() / polarP4.pt() < miniRelIso_cut_))
      continue;

    pdgId = pfCand.pdgId();
    dz = pfCand.dz();
    dxy = pfCand.dxy();
    if (pfCand.hasTrackDetails()) {
      dzError = pfCand.dzError();
      dxyError = pfCand.dxyError();
    } else {
      dzError = 0;
      dxyError = 0;
    }
    fromPV = pfCand.fromPV();
    refToCand = pcref;

    float caloJetEm, caloJetHad;
    getCaloJetEnergy(polarP4, caloJets.product(), caloJetEm, caloJetHad);

    bool pfLepOverlap = getPFLeptonOverlap(polarP4, pc);
    float pfNeutralSum = getPFNeutralSum(polarP4, pc, ipc);

    pat::PackedCandidateRef refToNearestPF = pat::PackedCandidateRef();
    int refToNearestPF_idx = -1;
    getNearestPCRef(polarP4, pc, ipc, refToNearestPF_idx);
    if (refToNearestPF_idx != -1)
      refToNearestPF = pat::PackedCandidateRef(pc_h, refToNearestPF_idx);

    pat::PackedCandidateRef refToNearestLostTrack = pat::PackedCandidateRef();
    int refToNearestLostTrack_idx = -1;
    getNearestPCRef(polarP4, lt, -1, refToNearestLostTrack_idx);
    if (refToNearestLostTrack_idx != -1)
      refToNearestLostTrack = pat::PackedCandidateRef(lt_h, refToNearestLostTrack_idx);

    // fill with default values
    reco::HitPattern hp;
    float dEdxPixel = -1, dEdxStrip = -1;
    int trackQuality = 0;
    std::vector<uint16_t> ecalStatus;
    std::vector<uint32_t> hcalStatus;
    int deltaEta = 0;
    int deltaPhi = 0;

    outPtrP->push_back(pat::IsolatedTrack(isolationDR03,
                                          miniIso,
                                          caloJetEm,
                                          caloJetHad,
                                          pfLepOverlap,
                                          pfNeutralSum,
                                          pfCand.p4(),
                                          charge,
                                          pdgId,
                                          dz,
                                          dxy,
                                          dzError,
                                          dxyError,
                                          hp,
                                          dEdxStrip,
                                          dEdxPixel,
                                          fromPV,
                                          trackQuality,
                                          ecalStatus,
                                          hcalStatus,
                                          deltaEta,
                                          deltaPhi,
                                          refToCand,
                                          refToNearestPF,
                                          refToNearestLostTrack));

    dEdXass.push_back(-1);  // these never have dE/dx hit info, as there's no track
  }

  auto orphHandle = iEvent.put(std::move(outPtrP));
  if (saveDeDxHitInfo_) {
    auto dedxOH = iEvent.put(std::move(outDeDxC));
    auto dedxMatch = std::make_unique<reco::DeDxHitInfoAss>(dedxOH);
    reco::DeDxHitInfoAss::Filler filler(*dedxMatch);
    filler.insert(orphHandle, dEdXass.begin(), dEdXass.end());
    filler.fill();
    iEvent.put(std::move(dedxMatch));
  }
}

void pat::PATIsolatedTrackProducer::getIsolation(const PolarLorentzVector& p4,
                                                 const pat::PackedCandidateCollection* pc,
                                                 int pc_idx,
                                                 pat::PFIsolation& iso,
                                                 pat::PFIsolation& miniiso) const {
  float chiso = 0, nhiso = 0, phiso = 0, puiso = 0;      // standard isolation
  float chmiso = 0, nhmiso = 0, phmiso = 0, pumiso = 0;  // mini isolation
  float miniDR = std::max(miniIsoParams_[0], std::min(miniIsoParams_[1], miniIsoParams_[2] / p4.pt()));
  for (pat::PackedCandidateCollection::const_iterator pf_it = pc->begin(); pf_it != pc->end(); pf_it++) {
    if (int(pf_it - pc->begin()) == pc_idx)  //don't count itself
      continue;
    int id = std::abs(pf_it->pdgId());
    bool fromPV = (pf_it->fromPV() > 1 || fabs(pf_it->dz()) < pfIsolation_DZ_);
    float pt = pf_it->p4().pt();
    float dr = deltaR(p4, *pf_it);

    if (dr < pfIsolation_DR_) {
      // charged cands from PV get added to trackIso
      if (id == 211 && fromPV)
        chiso += pt;
      // charged cands not from PV get added to pileup iso
      else if (id == 211)
        puiso += pt;
      // neutral hadron iso
      if (id == 130)
        nhiso += pt;
      // photon iso
      if (id == 22)
        phiso += pt;
    }
    // same for mini isolation
    if (dr < miniDR) {
      if (id == 211 && fromPV)
        chmiso += pt;
      else if (id == 211)
        pumiso += pt;
      if (id == 130)
        nhmiso += pt;
      if (id == 22)
        phmiso += pt;
    }
  }

  iso = pat::PFIsolation(chiso, nhiso, phiso, puiso);
  miniiso = pat::PFIsolation(chmiso, nhmiso, phmiso, pumiso);
}

//get overlap of isolated track with a PF lepton
bool pat::PATIsolatedTrackProducer::getPFLeptonOverlap(const PolarLorentzVector& p4,
                                                       const pat::PackedCandidateCollection* pc) const {
  bool isOverlap = false;
  float dr_min = pflepoverlap_DR_;
  int id_drmin = 0;
  for (const auto& pf : *pc) {
    int id = std::abs(pf.pdgId());
    int charge = std::abs(pf.charge());
    bool fromPV = (pf.fromPV() > 1 || std::abs(pf.dz()) < pfIsolation_DZ_);
    float pt = pf.pt();
    if (charge == 0)  // exclude neutral candidates
      continue;
    if (!(fromPV))  // exclude candidates not from PV
      continue;
    if (pt < pflepoverlap_pTmin_)  // exclude pf candidates w/ pT below threshold
      continue;

    float dr = deltaR(p4, pf);
    if (dr > pflepoverlap_DR_)  // exclude pf candidates far from isolated track
      continue;

    if (dr < dr_min) {
      dr_min = dr;
      id_drmin = id;
    }
  }

  if (dr_min < pflepoverlap_DR_ && (id_drmin == 11 || id_drmin == 13))
    isOverlap = true;

  return isOverlap;
}

//get ref to nearest pf packed candidate
void pat::PATIsolatedTrackProducer::getNearestPCRef(const PolarLorentzVector& p4,
                                                    const pat::PackedCandidateCollection* pc,
                                                    int pc_idx,
                                                    int& pc_ref_idx) const {
  float dr_min = pcRefNearest_DR_;
  float dr_min_pu = pcRefNearest_DR_;
  int pc_ref_idx_pu = -1;
  for (pat::PackedCandidateCollection::const_iterator pf_it = pc->begin(); pf_it != pc->end(); pf_it++) {
    if (int(pf_it - pc->begin()) == pc_idx)  //don't count itself
      continue;
    int charge = std::abs(pf_it->charge());
    bool fromPV = (pf_it->fromPV() > 1 || fabs(pf_it->dz()) < pfIsolation_DZ_);
    float pt = pf_it->p4().pt();
    float dr = deltaR(p4, *pf_it);
    if (charge == 0)  // exclude neutral candidates
      continue;
    if (pt < pcRefNearest_pTmin_)  // exclude candidates w/ pT below threshold
      continue;
    if (dr > dr_min && dr > dr_min_pu)  // exclude too far candidates
      continue;

    if (fromPV) {  // Priority to candidates from PV
      if (dr < dr_min) {
        dr_min = dr;
        pc_ref_idx = int(pf_it - pc->begin());
      }
    } else {  // Otherwise, store candidate from non-PV if no candidate from PV found
      if (dr < dr_min_pu) {
        dr_min_pu = dr;
        pc_ref_idx_pu = int(pf_it - pc->begin());
      }
    }
  }

  if (pc_ref_idx == -1 &&
      pc_ref_idx_pu != -1)  // If no candidate from PV was found, store candidate from non-PV (if found)
    pc_ref_idx = pc_ref_idx_pu;
}

// get PF neutral pT sum around isolated track
float pat::PATIsolatedTrackProducer::getPFNeutralSum(const PolarLorentzVector& p4,
                                                     const pat::PackedCandidateCollection* pc,
                                                     int pc_idx) const {
  float nsum = 0;
  float nhsum = 0, phsum = 0;
  for (pat::PackedCandidateCollection::const_iterator pf_it = pc->begin(); pf_it != pc->end(); pf_it++) {
    if (int(pf_it - pc->begin()) == pc_idx)  //don't count itself
      continue;
    int id = std::abs(pf_it->pdgId());
    float pt = pf_it->p4().pt();
    float dr = deltaR(p4, *pf_it);

    if (dr < pfneutralsum_DR_) {
      // neutral hadron sum
      if (id == 130)
        nhsum += pt;
      // photon iso
      if (id == 22)
        phsum += pt;
    }
  }

  nsum = nhsum + phsum;

  return nsum;
}

// get the estimated DeDx in either the pixels or strips (or both)
float pat::PATIsolatedTrackProducer::getDeDx(const reco::DeDxHitInfo* hitInfo, bool doPixel, bool doStrip) const {
  if (hitInfo == nullptr) {
    return -1;
  }

  std::vector<float> charge_vec;
  for (unsigned int ih = 0; ih < hitInfo->size(); ih++) {
    bool isPixel = (hitInfo->pixelCluster(ih) != nullptr);
    bool isStrip = (hitInfo->stripCluster(ih) != nullptr);

    if (isPixel && !doPixel)
      continue;
    if (isStrip && !doStrip)
      continue;

    // probably shouldn't happen
    if (!isPixel && !isStrip)
      continue;

    // shape selection for strips
    if (isStrip && !deDxTools::shapeSelection(*(hitInfo->stripCluster(ih))))
      continue;

    float Norm = 0;
    if (isPixel)
      Norm = 3.61e-06;  //compute the normalization factor to get the energy in MeV/mm
    if (isStrip)
      Norm = 3.61e-06 * 265;

    //save the dE/dx in MeV/mm to a vector.
    charge_vec.push_back(Norm * hitInfo->charge(ih) / hitInfo->pathlength(ih));
  }

  int size = charge_vec.size();
  float result = 0.0;

  //build the harmonic 2 dE/dx estimator
  float expo = -2;
  for (int i = 0; i < size; i++) {
    result += pow(charge_vec[i], expo);
  }
  result = (size > 0) ? pow(result / size, 1. / expo) : 0.0;

  return result;
}

TrackDetMatchInfo pat::PATIsolatedTrackProducer::getTrackDetMatchInfo(const edm::Event& iEvent,
                                                                      const edm::EventSetup& iSetup,
                                                                      const reco::Track& track) {
  auto const& bField = iSetup.getData(bFieldToken_);
  FreeTrajectoryState initialState = trajectoryStateTransform::initialFreeState(track, &bField);

  // can't use the associate() using reco::Track directly, since
  // track->extra() is non-null but segfaults when trying to use it
  return trackAssociator_.associate(iEvent, iSetup, trackAssocParameters_, &initialState);
}

void pat::PATIsolatedTrackProducer::getCaloJetEnergy(const PolarLorentzVector& p4,
                                                     const reco::CaloJetCollection* cJets,
                                                     float& caloJetEm,
                                                     float& caloJetHad) const {
  float nearestDR = 999;
  int ind = -1;
  for (unsigned int i = 0; i < cJets->size(); i++) {
    float dR = deltaR(cJets->at(i), p4);
    if (dR < caloJet_DR_ && dR < nearestDR) {
      nearestDR = dR;
      ind = i;
    }
  }

  if (ind == -1) {
    caloJetEm = 0;
    caloJetHad = 0;
  } else {
    const reco::CaloJet& cJet = cJets->at(ind);
    caloJetEm = cJet.emEnergyInEB() + cJet.emEnergyInEE() + cJet.emEnergyInHF();
    caloJetHad = cJet.hadEnergyInHB() + cJet.hadEnergyInHE() + cJet.hadEnergyInHF();
  }
}

using pat::PATIsolatedTrackProducer;
#include "FWCore/Framework/interface/MakerMacros.h"
DEFINE_FWK_MODULE(PATIsolatedTrackProducer);