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
0028 class DiTrackBuilder : public edm::global::EDProducer<> {
0029 public:
0030   typedef std::vector<reco::TransientTrack> TransientTrackCollection;
0031
0032   explicit DiTrackBuilder(const edm::ParameterSet &cfg)
0033       : trk1_selection_{cfg.getParameter<std::string>("trk1Selection")},
0034         trk2_selection_{cfg.getParameter<std::string>("trk2Selection")},
0035         pre_vtx_selection_{cfg.getParameter<std::string>("preVtxSelection")},
0036         post_vtx_selection_{cfg.getParameter<std::string>("postVtxSelection")},
0037         pfcands_{consumes<pat::CompositeCandidateCollection>(cfg.getParameter<edm::InputTag>("tracks"))},
0038         ttracks_{consumes<TransientTrackCollection>(cfg.getParameter<edm::InputTag>("transientTracks"))},
0039         beamspot_{consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))},
0040         trk1_mass_{cfg.getParameter<double>("trk1Mass")},
0041         trk2_mass_{cfg.getParameter<double>("trk2Mass")} {
0042     // output
0043     produces<pat::CompositeCandidateCollection>();
0044   }
0045
0046   ~DiTrackBuilder() override {}
0047
0048   void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0049
0050 private:
0051   const StringCutObjectSelector<pat::CompositeCandidate> trk1_selection_;      // cuts on leading cand
0052   const StringCutObjectSelector<pat::CompositeCandidate> trk2_selection_;      // sub-leading cand
0053   const StringCutObjectSelector<pat::CompositeCandidate> pre_vtx_selection_;   // cut on the di-track before the SV fit
0054   const StringCutObjectSelector<pat::CompositeCandidate> post_vtx_selection_;  // cut on the di-track after the SV fit
0055   const edm::EDGetTokenT<pat::CompositeCandidateCollection>
0056       pfcands_;                                               // input PF cands this is sorted in pT in previous step
0057   const edm::EDGetTokenT<TransientTrackCollection> ttracks_;  // input TTracks of PF cands
0058   const edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0059   double trk1_mass_;
0060   double trk2_mass_;
0061 };
0062
0063 void DiTrackBuilder::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &) const {
0064   // inputs
0065   edm::Handle<pat::CompositeCandidateCollection> pfcands;
0066   evt.getByToken(pfcands_, pfcands);
0067   edm::Handle<TransientTrackCollection> ttracks;
0068   evt.getByToken(ttracks_, ttracks);
0069   edm::Handle<reco::BeamSpot> beamspot;
0070   evt.getByToken(beamspot_, beamspot);
0071
0072   // output
0073   std::unique_ptr<pat::CompositeCandidateCollection> ditrack_out(new pat::CompositeCandidateCollection());
0074
0075   // main loop
0076   for (size_t trk1_idx = 0; trk1_idx < pfcands->size(); ++trk1_idx) {
0077     edm::Ptr<pat::CompositeCandidate> trk1_ptr(pfcands, trk1_idx);
0078     if (!trk1_selection_(*trk1_ptr))
0079       continue;
0080
0081     for (size_t trk2_idx = trk1_idx + 1; trk2_idx < pfcands->size(); ++trk2_idx) {
0082       edm::Ptr<pat::CompositeCandidate> trk2_ptr(pfcands, trk2_idx);
0083       // if (trk1_ptr->charge() == trk2_ptr->charge()) continue;
0084       if (!trk2_selection_(*trk2_ptr))
0085         continue;
0086
0087       pat::CompositeCandidate ditrack_cand;
0088       auto trk1_p4 = trk1_ptr->polarP4();
0089       auto trk2_p4 = trk2_ptr->polarP4();
0090       trk1_p4.SetM(bph::K_MASS);
0091       trk2_p4.SetM(bph::K_MASS);
0092       ditrack_cand.setP4(trk1_p4 + trk2_p4);
0093       ditrack_cand.setCharge(trk1_ptr->charge() + trk2_ptr->charge());
0094       ditrack_cand.addUserFloat("trk_deltaR", reco::deltaR(*trk1_ptr, *trk2_ptr));
0095       // save indices
0096       ditrack_cand.addUserInt("trk1_idx", trk1_idx);
0097       ditrack_cand.addUserInt("trk2_idx", trk2_idx);
0098       // save cands
0099       ditrack_cand.addUserCand("trk1", trk1_ptr);
0100       ditrack_cand.addUserCand("trk2", trk2_ptr);
0101
0102       ditrack_cand.addUserFloat("trk_dz", trk1_ptr->vz() - trk2_ptr->vz());
0103       ditrack_cand.addUserFloat("unfitted_mass_KK", (trk1_p4 + trk2_p4).M());
0104       trk1_p4.SetM(bph::K_MASS);
0105       trk2_p4.SetM(bph::PI_MASS);
0106       ditrack_cand.addUserFloat("unfitted_mass_Kpi", (trk1_p4 + trk2_p4).M());
0107       trk2_p4.SetM(bph::K_MASS);
0108       trk1_p4.SetM(bph::PI_MASS);
0109       ditrack_cand.addUserFloat("unfitted_mass_piK", (trk1_p4 + trk2_p4).M());
0110       trk2_p4.SetM(bph::K_MASS);
0111       trk1_p4.SetM(bph::K_MASS);
0112
0113       if (!pre_vtx_selection_(ditrack_cand))
0114         continue;
0115
0116       KinVtxFitter fitter(
0117           {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::K_MASS, bph::K_MASS}, {bph::K_SIGMA, bph::K_SIGMA}
0118           // K and PI sigma equal...
0119       );
0120       if (!fitter.success())
0121         continue;
0122       ditrack_cand.addUserFloat("fitted_mass_KK", fitter.fitted_candidate().mass());
0123       ditrack_cand.addUserFloat("fitted_mass_KK_Err",
0124                                 sqrt(fitter.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0125       // fits required in order to calculate the error of the mass for each mass
0126       // hypothesis.
0127       KinVtxFitter fitter_Kpi(
0128           {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::K_MASS, bph::PI_MASS}, {bph::K_SIGMA, bph::K_SIGMA}
0129           // K and PI sigma equal...
0130       );
0131       if (!fitter_Kpi.success())
0132         continue;
0133       ditrack_cand.addUserFloat("fitted_mass_Kpi", fitter_Kpi.fitted_candidate().mass());
0134       ditrack_cand.addUserFloat("fitted_mass_Kpi_Err",
0135                                 sqrt(fitter_Kpi.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0136       KinVtxFitter fitter_piK(
0137           {ttracks->at(trk1_idx), ttracks->at(trk2_idx)}, {bph::PI_MASS, bph::K_MASS}, {bph::K_SIGMA, bph::K_SIGMA}
0138           // K and PI sigma equal...
0139       );
0140       if (!fitter_piK.success())
0141         continue;
0142       ditrack_cand.addUserFloat("fitted_mass_piK", fitter_piK.fitted_candidate().mass());
0143       ditrack_cand.addUserFloat("fitted_mass_piK_Err",
0144                                 sqrt(fitter_piK.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0145
0146       ditrack_cand.setVertex(
0147           reco::Candidate::Point(fitter.fitted_vtx().x(), fitter.fitted_vtx().y(), fitter.fitted_vtx().z()));
0148       // save quantities after fit
0149       auto lxy = bph::l_xy(fitter, *beamspot);
0150       ditrack_cand.addUserFloat("l_xy", lxy.value());
0151       ditrack_cand.addUserFloat("l_xy_unc", lxy.error());
0152       ditrack_cand.addUserInt("sv_ok", fitter.success() ? 1 : 0);
0153       auto fit_p4 = fitter.fitted_p4();
0154       ditrack_cand.addUserFloat("fitted_cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, fit_p4));
0155       // The following quantities do not independent on the mass hypothesis
0156       ditrack_cand.addUserFloat("sv_chi2", fitter.chi2());
0157       ditrack_cand.addUserFloat("sv_ndof", fitter.dof());
0158       ditrack_cand.addUserFloat("sv_prob", fitter.prob());
0159       ditrack_cand.addUserFloat("fitted_pt", fitter.fitted_candidate().globalMomentum().perp());
0160       ditrack_cand.addUserFloat("fitted_eta", fitter.fitted_candidate().globalMomentum().eta());
0161       ditrack_cand.addUserFloat("fitted_phi", fitter.fitted_candidate().globalMomentum().phi());
0162       ditrack_cand.addUserFloat("vtx_x", ditrack_cand.vx());
0163       ditrack_cand.addUserFloat("vtx_y", ditrack_cand.vy());
0164       ditrack_cand.addUserFloat("vtx_z", ditrack_cand.vz());
0165       const auto &covMatrix = fitter.fitted_vtx_uncertainty();
0166       ditrack_cand.addUserFloat("vtx_cxx", covMatrix.cxx());
0167       ditrack_cand.addUserFloat("vtx_cyy", covMatrix.cyy());
0168       ditrack_cand.addUserFloat("vtx_czz", covMatrix.czz());
0169       ditrack_cand.addUserFloat("vtx_cyx", covMatrix.cyx());
0170       ditrack_cand.addUserFloat("vtx_czx", covMatrix.czx());
0171       ditrack_cand.addUserFloat("vtx_czy", covMatrix.czy());
0172
0173       // after fit selection
0174       if (!post_vtx_selection_(ditrack_cand))
0175         continue;
0176       ditrack_out->emplace_back(ditrack_cand);
0177     }  // end for(size_t trk2_idx = trk1_idx + 1
0178   }  // for(size_t trk1_idx = 0
0179
0180   evt.put(std::move(ditrack_out));
0181 }
0182
0183 #include "FWCore/Framework/interface/MakerMacros.h"
0184 DEFINE_FWK_MODULE(DiTrackBuilder);