BPHNanoAODs

DataFormats/PatCandidates/src/classes_def_objects.xml

@@ -630,6 +630,10 @@
   <class name="edm::Association<std::vector<pat::PackedGenParticle> >" />
   <class name="edm::Wrapper<edm::Association<std::vector<pat::PackedGenParticle> > >"/>
 
+  <class name="edm::RefProd<pat::CompositeCandidateCollection>" />
+  <class name="edm::Association<pat::CompositeCandidateCollection>" />
+  <class name="edm::Wrapper<edm::Association<pat::CompositeCandidateCollection> >"/>
+
   <!--- vectors of Ptrs -->
   <class name="std::vector<edm::Ptr<pat::Jet> >" />
   <class name="std::vector< std::vector<edm::Ptr<pat::Jet> > >" />

PhysicsTools/BPHNano/plugins/BPHTrackMerger.cc

@@ -0,0 +1,264 @@
+/////////////////////////// BPHTrackMerger ////////////////////////////////
+/// original authors: G Karathanasis (CERN),  G Melachroinos (NKUA)
+// Takes Lost tracks and packed candidates filters them removes overlap and
+// appl// -ies dz cut wrt to a dilepton vertex. Also applies selection cuts
+
+#include "CommonTools/Utils/interface/StringCutObjectSelector.h"
+#include "DataFormats/BeamSpot/interface/BeamSpot.h"
+#include "DataFormats/Common/interface/RefToPtr.h"
+#include "DataFormats/Common/interface/View.h"
+#include "DataFormats/GeometryVector/interface/GlobalPoint.h"
+#include "DataFormats/PatCandidates/interface/CompositeCandidate.h"
+#include "DataFormats/PatCandidates/interface/Electron.h"
+#include "DataFormats/PatCandidates/interface/Muon.h"
+#include "DataFormats/PatCandidates/interface/PackedCandidate.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/Framework/interface/global/EDProducer.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+#include "FWCore/Utilities/interface/InputTag.h"
+#include "MagneticField/Engine/interface/MagneticField.h"
+#include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
+#include "TrackingTools/TransientTrack/interface/TransientTrack.h"
+#include "helper.h"
+
+class BPHTrackMerger : public edm::global::EDProducer<> {
+public:
+  // would it be useful to give this a bit more standard structure?
+  explicit BPHTrackMerger(const edm::ParameterSet &cfg)
+      : bFieldToken_(esConsumes<MagneticField, IdealMagneticFieldRecord>()),
+        beamSpotSrc_(consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))),
+        tracksToken_(consumes<pat::PackedCandidateCollection>(cfg.getParameter<edm::InputTag>("tracks"))),
+        lostTracksToken_(consumes<pat::PackedCandidateCollection>(cfg.getParameter<edm::InputTag>("lostTracks"))),
+        dileptonToken_(consumes<pat::CompositeCandidateCollection>(cfg.getParameter<edm::InputTag>("dileptons"))),
+        muonToken_(consumes<pat::MuonCollection>(cfg.getParameter<edm::InputTag>("muons"))),
+        eleToken_(consumes<pat::ElectronCollection>(cfg.getParameter<edm::InputTag>("electrons"))),
+        maxDzDilep_(cfg.getParameter<double>("maxDzDilep")),
+        dcaSig_(cfg.getParameter<double>("dcaSig")),
+        track_selection_(cfg.getParameter<std::string>("trackSelection")) {
+    produces<pat::CompositeCandidateCollection>("SelectedTracks");
+    produces<TransientTrackCollection>("SelectedTransientTracks");
+    produces<edm::Association<pat::CompositeCandidateCollection>>("SelectedTracks");
+  }
+
+  ~BPHTrackMerger() override {}
+
+  void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
+
+private:
+  const edm::ESGetToken<MagneticField, IdealMagneticFieldRecord> bFieldToken_;
+  const edm::EDGetTokenT<reco::BeamSpot> beamSpotSrc_;
+  const edm::EDGetTokenT<pat::PackedCandidateCollection> tracksToken_;
+  const edm::EDGetTokenT<pat::PackedCandidateCollection> lostTracksToken_;
+  const edm::EDGetTokenT<pat::CompositeCandidateCollection> dileptonToken_;
+  const edm::EDGetTokenT<pat::MuonCollection> muonToken_;
+  const edm::EDGetTokenT<pat::ElectronCollection> eleToken_;
+
+  // selections
+  const double maxDzDilep_;
+  const double dcaSig_;
+  const StringCutObjectSelector<pat::PackedCandidate> track_selection_;
+};
+
+void BPHTrackMerger::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &stp) const {
+  // input
+  edm::Handle<reco::BeamSpot> beamSpotHandle;
+  evt.getByToken(beamSpotSrc_, beamSpotHandle);
+  if (!beamSpotHandle.isValid()) {
+    edm::LogError("BToKstllProducer") << "No beam spot available from Event";
+  }
+  const reco::BeamSpot &beamSpot = *beamSpotHandle;
+
+  const auto &bField = stp.getData(bFieldToken_);
+  edm::Handle<pat::PackedCandidateCollection> tracks;
+  evt.getByToken(tracksToken_, tracks);
+  edm::Handle<pat::PackedCandidateCollection> lostTracks;
+  evt.getByToken(lostTracksToken_, lostTracks);
+  edm::Handle<pat::CompositeCandidateCollection> dileptons;
+  evt.getByToken(dileptonToken_, dileptons);
+
+  edm::Handle<pat::MuonCollection> muons;
+  evt.getByToken(muonToken_, muons);
+  edm::Handle<pat::ElectronCollection> pfele;
+  evt.getByToken(eleToken_, pfele);
+
+  // for lost tracks / pf discrimination
+  unsigned int nTracks = tracks->size();
+  unsigned int totalTracks = nTracks + lostTracks->size();
+
+  // ok this was CompositeCandidateCollection
+  std::unique_ptr<pat::CompositeCandidateCollection> tracks_out(new pat::CompositeCandidateCollection);
+  std::unique_ptr<TransientTrackCollection> trans_tracks_out(new TransientTrackCollection);
+
+  std::vector<std::pair<pat::CompositeCandidate, reco::TransientTrack>> vectrk_ttrk;
+  // try topreserve same logic avoiding the copy of the full collection
+  /*
+  //correct logic but a bit convoluted -> changing to smthn simpler
+   std::vector<pat::PackedCandidate> totalTracks(*tracks);
+   totalTracks.insert(totalTracks.end(),lostTracks->begin(),lostTracks->end());
+  */
+
+  std::vector<int> match_indices(totalTracks, -1);
+  // for loop is better to be range based - especially for large ensembles
+  for (unsigned int iTrk = 0; iTrk < totalTracks; ++iTrk) {
+    const pat::PackedCandidate &trk = (iTrk < nTracks) ? (*tracks)[iTrk] : (*lostTracks)[iTrk - nTracks];
+    // arranging cuts for speed
+    if (!trk.hasTrackDetails())
+      continue;
+    if (abs(trk.pdgId()) != 211)
+      continue;  // do we want also to keep muons?
+    if (!track_selection_(trk))
+      continue;
+
+    bool skipTrack = true;
+    float dzTrg = 0.0;
+    for (const pat::CompositeCandidate &dilep : *dileptons) {
+      // if dz is negative it is deactivated
+      if (fabs(trk.vz() - dilep.vz()) > maxDzDilep_ && maxDzDilep_ > 0)
+        continue;
+      skipTrack = false;
+      dzTrg = trk.vz() - dilep.vz();
+      break;  // at least for one dilepton candidate to pass this cuts
+    }
+
+    // if track is far from all dilepton candidate
+    if (skipTrack)
+      continue;
+
+    // high purity requirment applied only in packedCands
+    if (iTrk < nTracks && !trk.trackHighPurity())
+      continue;
+    const reco::TransientTrack trackTT((*trk.bestTrack()), &bField);
+
+    // distance closest approach in x,y wrt beam spot
+    std::pair<double, double> DCA = bph::computeDCA(trackTT, beamSpot);
+    float DCABS = DCA.first;
+    float DCABSErr = DCA.second;
+    float DCASig = (DCABSErr != 0 && float(DCABSErr) == DCABSErr) ? fabs(DCABS / DCABSErr) : -1;
+    if (DCASig > dcaSig_ && dcaSig_ > 0)
+      continue;
+
+    // clean tracks wrt to all muons
+    int matchedToMuon = 0;
+    for (const pat::Muon &imutmp : *muons) {
+      for (unsigned int i = 0; i < imutmp.numberOfSourceCandidatePtrs(); ++i) {
+        if (!((imutmp.sourceCandidatePtr(i)).isNonnull() && (imutmp.sourceCandidatePtr(i)).isAvailable()))
+          continue;
+
+        const edm::Ptr<reco::Candidate> &source = imutmp.sourceCandidatePtr(i);
+        if (source.id() == tracks.id() && source.key() == iTrk) {
+          matchedToMuon = 1;
+          break;
+        }
+      }
+    }
+
+    // clean tracks wrt to all pf electrons
+    int matchedToEle = 0;
+    for (const pat::Electron &ietmp : *pfele) {
+      for (unsigned int i = 0; i < ietmp.numberOfSourceCandidatePtrs(); ++i) {
+        if (!((ietmp.sourceCandidatePtr(i)).isNonnull() && (ietmp.sourceCandidatePtr(i)).isAvailable()))
+          continue;
+        const edm::Ptr<reco::Candidate> &source = ietmp.sourceCandidatePtr(i);
+        if (source.id() == tracks.id() && source.key() == iTrk) {
+          matchedToEle = 1;
+          break;
+        }
+      }
+    }
+
+    // output
+    pat::CompositeCandidate pcand;
+    pcand.setP4(trk.p4());
+    pcand.setCharge(trk.charge());
+    pcand.setVertex(trk.vertex());
+    pcand.setPdgId(trk.pdgId());
+    pcand.addUserInt("isPacked", (iTrk < nTracks));
+    pcand.addUserInt("isLostTrk", (iTrk < nTracks) ? 0 : 1);
+    pcand.addUserFloat("dxy", trk.dxy());
+    pcand.addUserFloat("dxyS", trk.dxy() / trk.dxyError());
+    pcand.addUserFloat("dz", trk.dz());
+    pcand.addUserFloat("dzS", trk.dz() / trk.dzError());
+    pcand.addUserFloat("DCASig", DCASig);
+    pcand.addUserFloat("dzTrg", dzTrg);
+    pcand.addUserInt("isMatchedToMuon", matchedToMuon);
+    pcand.addUserInt("isMatchedToEle", matchedToEle);
+    pcand.addUserInt("nValidHits", trk.bestTrack()->found());
+    pcand.addUserInt("keyPacked", iTrk);
+
+    // Covariance matrix elements for helix parameters for decay time
+    // uncertainty
+    pcand.addUserFloat("covQopQop", trk.bestTrack()->covariance(0, 0));
+    pcand.addUserFloat("covLamLam", trk.bestTrack()->covariance(1, 1));
+    pcand.addUserFloat("covPhiPhi", trk.bestTrack()->covariance(2, 2));
+    pcand.addUserFloat("covQopLam", trk.bestTrack()->covariance(0, 1));
+    pcand.addUserFloat("covQopPhi", trk.bestTrack()->covariance(0, 2));
+    pcand.addUserFloat("covLamPhi", trk.bestTrack()->covariance(1, 2));
+
+    // Additional track parameters for tagging
+    pcand.addUserFloat("ptErr", trk.bestTrack()->ptError());
+    pcand.addUserFloat("normChi2", trk.bestTrack()->normalizedChi2());
+    pcand.addUserInt("nValidPixelHits", trk.numberOfPixelHits());
+
+    // adding the candidate in the composite stuff for fit (need to test)
+    if (iTrk < nTracks)
+      pcand.addUserCand("cand", edm::Ptr<pat::PackedCandidate>(tracks, iTrk));
+    else
+      pcand.addUserCand("cand", edm::Ptr<pat::PackedCandidate>(lostTracks, iTrk - nTracks));
+
+    // in order to avoid revoking the sxpensive ttrack builder many times and
+    // still have everything sorted, we add them to vector of pairs
+    match_indices[iTrk] = vectrk_ttrk.size();
+    vectrk_ttrk.emplace_back(std::make_pair(pcand, trackTT));
+  }
+
+  std::vector<int> sort_indices(vectrk_ttrk.size());
+  std::iota(sort_indices.begin(), sort_indices.end(), 0);
+
+  // sort to be uniform with leptons
+  // sort by index since we want to update the match too
+  std::sort(sort_indices.begin(), sort_indices.end(), [&vectrk_ttrk](auto &iTrk1, auto &iTrk2) -> bool {
+    return (vectrk_ttrk[iTrk1].first).pt() > (vectrk_ttrk[iTrk2].first).pt();
+  });
+  // std::sort( vectrk_ttrk.begin(), vectrk_ttrk.end(),
+  //            [] ( auto & trk1, auto & trk2) ->
+  //            bool {return (trk1.first).pt() > (trk2.first).pt();}
+  //          );
+
+  // finally save ttrks and trks to the correct _out vectors
+  // also fill the reverse matching
+  std::vector<int> reverse_sort_indices(vectrk_ttrk.size());
+  for (size_t iSort = 0; iSort < sort_indices.size(); iSort++) {
+    auto iUnsortedTrack = sort_indices[iSort];
+    auto &&trk = vectrk_ttrk[iUnsortedTrack];
+    tracks_out->emplace_back(trk.first);
+    trans_tracks_out->emplace_back(trk.second);
+    reverse_sort_indices[iUnsortedTrack] = iSort;
+  }
+
+  // Now point the match indices to the sorted output collection
+  std::transform(
+      match_indices.begin(), match_indices.end(), match_indices.begin(), [&reverse_sort_indices](int iUnsortedTrack) {
+        if (iUnsortedTrack < 0)
+          return -1;
+        return reverse_sort_indices[iUnsortedTrack];
+      });
+
+  auto tracks_orphan_handle = evt.put(std::move(tracks_out), "SelectedTracks");
+  evt.put(std::move(trans_tracks_out), "SelectedTransientTracks");
+
+  // Associate PackedCandidates to the merged Track collection
+  auto tracks_out_match = std::make_unique<edm::Association<pat::CompositeCandidateCollection>>(tracks_orphan_handle);
+  edm::Association<pat::CompositeCandidateCollection>::Filler filler(*tracks_out_match);
+
+  filler.insert(tracks, match_indices.begin(), match_indices.begin() + nTracks);
+  filler.insert(lostTracks, match_indices.begin() + nTracks, match_indices.end());
+  filler.fill();
+
+  evt.put(std::move(tracks_out_match), "SelectedTracks");
+}
+
+// define this as a plug-in
+DEFINE_FWK_MODULE(BPHTrackMerger);