BPHNano/plugins/DiTrackBuilder.cc

0001 /////////////////////////////// DiTrackBuilder ///////////////////////////////
0002 /// original authors: G Karathanasis (CERN),  G Melachroinos (NKUA)
0003 // takes selected track collection and a mass hypothesis and produces ditrack ca
0004 // -ndidates
0005
0006 #include <algorithm>
0007 #include <limits>
0008 #include <map>
0009 #include <memory>
0010 #include <string>
0011 #include <vector>
0012
0013 #include "CommonTools/Statistics/interface/ChiSquaredProbability.h"
0014 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0015 #include "DataFormats/Math/interface/deltaR.h"
0016 #include "DataFormats/PatCandidates/interface/CompositeCandidate.h"
0017 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0018 #include "FWCore/Framework/interface/Event.h"
0019 #include "FWCore/Framework/interface/global/EDProducer.h"
0020 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0021 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0022 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0023 #include "FWCore/Utilities/interface/InputTag.h"
0024 #include "KinVtxFitter.h"
0025 #include "TrackingTools/TransientTrack/interface/TransientTrack.h"
0026 #include "helper.h"
0027 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0028
0029 class DiTrackBuilder : public edm::global::EDProducer<> {
0030 public:
0031   typedef std::vector<reco::TransientTrack> TransientTrackCollection;
0032
0033   explicit DiTrackBuilder(const edm::ParameterSet &cfg)
0034       : trk1_selection_{cfg.getParameter<std::string>("trk1Selection")},
0035         trk2_selection_{cfg.getParameter<std::string>("trk2Selection")},
0036         pre_vtx_selection_{cfg.getParameter<std::string>("preVtxSelection")},
0037         post_vtx_selection_{cfg.getParameter<std::string>("postVtxSelection")},
0038         pfcands_{consumes<pat::CompositeCandidateCollection>(cfg.getParameter<edm::InputTag>("tracks"))},
0039         ttracks_{consumes<TransientTrackCollection>(cfg.getParameter<edm::InputTag>("transientTracks"))},
0040         beamspot_{consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))},
0041         trk1_mass_{cfg.getParameter<double>("trk1Mass")},
0042         trk2_mass_{cfg.getParameter<double>("trk2Mass")} {
0043     // output
0044     produces<pat::CompositeCandidateCollection>();
0045   }
0046
0047   ~DiTrackBuilder() override {}
0048
0049   void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0050
0051 private:
0052   const StringCutObjectSelector<pat::CompositeCandidate> trk1_selection_;      // cuts on leading cand
0053   const StringCutObjectSelector<pat::CompositeCandidate> trk2_selection_;      // sub-leading cand
0054   const StringCutObjectSelector<pat::CompositeCandidate> pre_vtx_selection_;   // cut on the di-track before the SV fit
0055   const StringCutObjectSelector<pat::CompositeCandidate> post_vtx_selection_;  // cut on the di-track after the SV fit
0056   const edm::EDGetTokenT<pat::CompositeCandidateCollection>
0057       pfcands_;                                               // input PF cands this is sorted in pT in previous step
0058   const edm::EDGetTokenT<TransientTrackCollection> ttracks_;  // input TTracks of PF cands
0059   const edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0060   double trk1_mass_;
0061   double trk2_mass_;
0062 };
0063
0064 void DiTrackBuilder::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &) const {
0065   // inputs
0066   edm::Handle<pat::CompositeCandidateCollection> pfcands;
0067   evt.getByToken(pfcands_, pfcands);
0068   edm::Handle<TransientTrackCollection> ttracks;
0069   evt.getByToken(ttracks_, ttracks);
0070   edm::Handle<reco::BeamSpot> beamspot;
0071   evt.getByToken(beamspot_, beamspot);
0072
0073   // output
0074   std::unique_ptr<pat::CompositeCandidateCollection> ditrack_out(new pat::CompositeCandidateCollection());
0075
0076   // main loop
0077   for (size_t trk1_idx = 0; trk1_idx < pfcands->size(); ++trk1_idx) {
0078     edm::Ptr<pat::CompositeCandidate> trk1_ptr(pfcands, trk1_idx);
0079     if (!trk1_selection_(*trk1_ptr))
0080       continue;
0081
0082     for (size_t trk2_idx = trk1_idx + 1; trk2_idx < pfcands->size(); ++trk2_idx) {
0083       edm::Ptr<pat::CompositeCandidate> trk2_ptr(pfcands, trk2_idx);
0084       // if (trk1_ptr->charge() == trk2_ptr->charge()) continue;
0085       if (!trk2_selection_(*trk2_ptr))
0086         continue;
0087
0088       pat::CompositeCandidate ditrack_cand;
0089       auto trk1_p4 = trk1_ptr->polarP4();
0090       auto trk2_p4 = trk2_ptr->polarP4();
0091       trk1_p4.SetM(bph::K_MASS);
0092       trk2_p4.SetM(bph::K_MASS);
0093       ditrack_cand.setP4(trk1_p4 + trk2_p4);
0094       ditrack_cand.setCharge(trk1_ptr->charge() + trk2_ptr->charge());
0095       ditrack_cand.addUserFloat("trk_deltaR", reco::deltaR(*trk1_ptr, *trk2_ptr));
0096       // save indices
0097       ditrack_cand.addUserInt("trk1_idx", trk1_idx);
0098       ditrack_cand.addUserInt("trk2_idx", trk2_idx);
0099       // save cands
0100       ditrack_cand.addUserCand("trk1", trk1_ptr);
0101       ditrack_cand.addUserCand("trk2", trk2_ptr);
0102
0103       ditrack_cand.addUserFloat("trk_dz", trk1_ptr->vz() - trk2_ptr->vz());
0104       ditrack_cand.addUserFloat("unfitted_mass_KK", (trk1_p4 + trk2_p4).M());
0105       trk1_p4.SetM(bph::K_MASS);
0106       trk2_p4.SetM(bph::PI_MASS);
0107       ditrack_cand.addUserFloat("unfitted_mass_Kpi", (trk1_p4 + trk2_p4).M());
0108       trk2_p4.SetM(bph::K_MASS);
0109       trk1_p4.SetM(bph::PI_MASS);
0110       ditrack_cand.addUserFloat("unfitted_mass_piK", (trk1_p4 + trk2_p4).M());
0111       trk2_p4.SetM(bph::K_MASS);
0112       trk1_p4.SetM(bph::K_MASS);
0113
0114       if (!pre_vtx_selection_(ditrack_cand))
0115         continue;
0116
0117       KinVtxFitter fitter;
0118       try {
0119         fitter = KinVtxFitter(
0120             {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::K_MASS, bph::K_MASS}, {bph::K_SIGMA, bph::K_SIGMA});
0121       } catch (const VertexException &e) {
0122         edm::LogWarning("KinematicFit")
0123             << "DiTrack Rebuilder - KK mass hypothesis: Skipping candidate due to fit failure: " << e.what();
0124         continue;
0125       }
0126       if (!fitter.success())
0127         continue;
0128       ditrack_cand.addUserFloat("fitted_mass_KK", fitter.fitted_candidate().mass());
0129       ditrack_cand.addUserFloat("fitted_mass_KK_Err",
0130                                 sqrt(fitter.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0131       // fits required in order to calculate the error of the mass for each mass
0132       // hypothesis.
0133       KinVtxFitter fitter_Kpi;
0134       try {
0135         fitter_Kpi = KinVtxFitter(
0136             {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::K_MASS, bph::PI_MASS}, {bph::K_SIGMA, bph::K_SIGMA});
0137       } catch (const VertexException &e) {
0138         edm::LogWarning("KinematicFit")
0139             << "DiTrack Rebuilder - Kpi mass hypothesis: Skipping candidate due to fit failure: " << e.what();
0140         continue;
0141       }
0142       if (!fitter_Kpi.success())
0143         continue;
0144       ditrack_cand.addUserFloat("fitted_mass_Kpi", fitter_Kpi.fitted_candidate().mass());
0145       ditrack_cand.addUserFloat("fitted_mass_Kpi_Err",
0146                                 sqrt(fitter_Kpi.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0147       KinVtxFitter fitter_piK;
0148       try {
0149         fitter_piK = KinVtxFitter(
0150             {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::PI_MASS, bph::K_MASS}, {bph::K_SIGMA, bph::K_SIGMA});
0151       } catch (const VertexException &e) {
0152         edm::LogWarning("KinematicFit")
0153             << "DiTrack Rebuilder - piK mass hypothesis: Skipping candidate due to fit failure: " << e.what();
0154         continue;
0155       }
0156       if (!fitter_piK.success())
0157         continue;
0158       ditrack_cand.addUserFloat("fitted_mass_piK", fitter_piK.fitted_candidate().mass());
0159       ditrack_cand.addUserFloat("fitted_mass_piK_Err",
0160                                 sqrt(fitter_piK.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0161       ditrack_cand.addUserFloat("fitted_mass_piK", fitter_piK.fitted_candidate().mass());
0162       ditrack_cand.addUserFloat("fitted_mass_piK_Err",
0163                                 sqrt(fitter_piK.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0164
0165       ditrack_cand.setVertex(
0166           reco::Candidate::Point(fitter.fitted_vtx().x(), fitter.fitted_vtx().y(), fitter.fitted_vtx().z()));
0167       // save quantities after fit
0168       auto lxy = bph::l_xy(fitter, *beamspot);
0169       ditrack_cand.addUserFloat("l_xy", lxy.value());
0170       ditrack_cand.addUserFloat("l_xy_unc", lxy.error());
0171       ditrack_cand.addUserInt("sv_ok", fitter.success() ? 1 : 0);
0172       auto fit_p4 = fitter.fitted_p4();
0173       ditrack_cand.addUserFloat("fitted_cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, fit_p4));
0174       // The following quantities do not independent on the mass hypothesis
0175       ditrack_cand.addUserFloat("sv_chi2", fitter.chi2());
0176       ditrack_cand.addUserFloat("sv_ndof", fitter.dof());
0177       ditrack_cand.addUserFloat("sv_prob", fitter.prob());
0178       ditrack_cand.addUserFloat("fitted_pt", fitter.fitted_candidate().globalMomentum().perp());
0179       ditrack_cand.addUserFloat("fitted_eta", fitter.fitted_candidate().globalMomentum().eta());
0180       ditrack_cand.addUserFloat("fitted_phi", fitter.fitted_candidate().globalMomentum().phi());
0181       ditrack_cand.addUserFloat("vtx_x", ditrack_cand.vx());
0182       ditrack_cand.addUserFloat("vtx_y", ditrack_cand.vy());
0183       ditrack_cand.addUserFloat("vtx_z", ditrack_cand.vz());
0184       const auto &covMatrix = fitter.fitted_vtx_uncertainty();
0185       ditrack_cand.addUserFloat("vtx_cxx", covMatrix.cxx());
0186       ditrack_cand.addUserFloat("vtx_cyy", covMatrix.cyy());
0187       ditrack_cand.addUserFloat("vtx_czz", covMatrix.czz());
0188       ditrack_cand.addUserFloat("vtx_cyx", covMatrix.cyx());
0189       ditrack_cand.addUserFloat("vtx_czx", covMatrix.czx());
0190       ditrack_cand.addUserFloat("vtx_czy", covMatrix.czy());
0191
0192       // after fit selection
0193       if (!post_vtx_selection_(ditrack_cand))
0194         continue;
0195       ditrack_out->emplace_back(ditrack_cand);
0196     }  // end for(size_t trk2_idx = trk1_idx + 1
0197   }  // for(size_t trk1_idx = 0
0198
0199   evt.put(std::move(ditrack_out));
0200 }
0201
0202 #include "FWCore/Framework/interface/MakerMacros.h"
0203 DEFINE_FWK_MODULE(DiTrackBuilder);