BPHNano/plugins/V0ReBuilder.cc

0001 /////////////////////////////// BToV0LLBuilder ///////////////////////////////
0002 /// original authors: G Karathanasis (CERN),  G Melachroinos (NKUA)
0003 // takes V0 cabds from CMSSW and creates useful V0
0004
0005 #include <algorithm>
0006 #include <limits>
0007 #include <map>
0008 #include <memory>
0009 #include <string>
0010 #include <vector>
0011
0012 #include "CommonTools/Statistics/interface/ChiSquaredProbability.h"
0013 #include "CommonTools/Utils/interface/StringCutObjectSelector.h"
0014 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0015 #include "DataFormats/Candidate/interface/VertexCompositeCandidate.h"
0016 #include "DataFormats/Candidate/interface/VertexCompositePtrCandidate.h"
0017 #include "DataFormats/Candidate/interface/VertexCompositePtrCandidateFwd.h"
0018 #include "DataFormats/Math/interface/deltaR.h"
0019 #include "DataFormats/PatCandidates/interface/CompositeCandidate.h"
0020 #include "DataFormats/PatCandidates/interface/PackedCandidate.h"
0021 #include "FWCore/Framework/interface/Event.h"
0022 #include "FWCore/Framework/interface/global/EDProducer.h"
0023 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0024 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0025 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0026 #include "FWCore/Utilities/interface/InputTag.h"
0027 #include "KinVtxFitter.h"
0028 #include "MagneticField/Engine/interface/MagneticField.h"
0029 #include "MagneticField/Records/interface/IdealMagneticFieldRecord.h"
0030 #include "RecoVertex/KinematicFit/interface/KinematicConstrainedVertexFitter.h"
0031 #include "RecoVertex/KinematicFit/interface/TwoTrackMassKinematicConstraint.h"
0032 #include "RecoVertex/KinematicFitPrimitives/interface/KinematicParticleFactoryFromTransientTrack.h"
0033 #include "RecoVertex/KinematicFitPrimitives/interface/MultiTrackKinematicConstraint.h"
0034 #include "TrackingTools/Records/interface/TransientTrackRecord.h"
0035 #include "TrackingTools/TransientTrack/interface/TransientTrack.h"
0036 #include "TrackingTools/TransientTrack/interface/TransientTrackBuilder.h"
0037 #include "helper.h"
0038
0039 class V0ReBuilder : public edm::global::EDProducer<> {
0040   // perhaps we need better structure here (begin run etc)
0041 public:
0042   typedef std::vector<reco::TransientTrack> TransientTrackCollection;
0043   typedef std::vector<reco::VertexCompositePtrCandidate> V0Collection;
0044
0045   explicit V0ReBuilder(const edm::ParameterSet &cfg)
0046       : theB_(esConsumes(edm::ESInputTag{"", "TransientTrackBuilder"})),
0047         trk_selection_{cfg.getParameter<std::string>("trkSelection")},
0048         pre_vtx_selection_{cfg.getParameter<std::string>("V0Selection")},
0049         post_vtx_selection_{cfg.getParameter<std::string>("postVtxSelection")},
0050         v0s_{consumes<V0Collection>(cfg.getParameter<edm::InputTag>("V0s"))},
0051         beamspot_{consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))},
0052         track_match_{consumes<edm::Association<pat::CompositeCandidateCollection>>(
0053             cfg.getParameter<edm::InputTag>("track_match"))},
0054         isLambda_{cfg.getParameter<bool>("isLambda")} {
0055     produces<pat::CompositeCandidateCollection>("SelectedV0Collection");
0056     produces<TransientTrackCollection>("SelectedV0TransientCollection");
0057   }
0058
0059   ~V0ReBuilder() override {}
0060
0061   void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0062
0063 private:
0064   const edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> theB_;
0065   const StringCutObjectSelector<pat::PackedCandidate> trk_selection_;
0066   const StringCutObjectSelector<reco::VertexCompositePtrCandidate> pre_vtx_selection_;
0067   const StringCutObjectSelector<pat::CompositeCandidate> post_vtx_selection_;
0068   const edm::EDGetTokenT<V0Collection> v0s_;
0069   const edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0070   const edm::EDGetTokenT<edm::Association<pat::CompositeCandidateCollection>> track_match_;
0071   const bool isLambda_;
0072 };
0073
0074 void V0ReBuilder::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &iSetup) const {
0075   // input
0076   auto const theB = &iSetup.getData(theB_);
0077   edm::Handle<V0Collection> V0s;
0078   evt.getByToken(v0s_, V0s);
0079   edm::Handle<reco::BeamSpot> beamspot;
0080   evt.getByToken(beamspot_, beamspot);
0081
0082   auto &track_match = evt.get(track_match_);
0083
0084   // output
0085   std::unique_ptr<pat::CompositeCandidateCollection> ret_val(new pat::CompositeCandidateCollection());
0086   std::unique_ptr<TransientTrackCollection> trans_out(new TransientTrackCollection);
0087
0088   for (reco::VertexCompositePtrCandidateCollection::const_iterator v0 = V0s->begin(); v0 != V0s->end(); v0++) {
0089     // selection on V0s
0090     if (v0->numberOfDaughters() != 2)
0091       continue;
0092     if (!pre_vtx_selection_(*v0))
0093       continue;
0094
0095     pat::PackedCandidate v0daughter1 = *(dynamic_cast<const pat::PackedCandidate *>(v0->daughter(0)));
0096     pat::PackedCandidate v0daughter2 = *(dynamic_cast<const pat::PackedCandidate *>(v0->daughter(1)));
0097
0098     if (!v0daughter1.hasTrackDetails())
0099       continue;
0100     if (!v0daughter2.hasTrackDetails())
0101       continue;
0102
0103     if (abs(v0daughter1.pdgId()) != 211)
0104       continue;  // This cut do not affect the Lambda->proton pion candidates
0105     if (abs(v0daughter2.pdgId()) != 211)
0106       continue;  // This cut do not affect the Lambda->proton pion candidates
0107
0108     if (!trk_selection_(v0daughter1) || !trk_selection_(v0daughter2))
0109       continue;
0110
0111     reco::TransientTrack v0daughter1_ttrack;  // 1st daughter, leading daughter to be assigned.
0112                                               // Proton mass will be assigned for the
0113                                               // Lambda->Proton Pion mode, Pion mass will be
0114                                               // assigned for the Kshort->PionPion mode.
0115     reco::TransientTrack v0daughter2_ttrack;  // 2nd daughter, subleading daughter to be
0116                                               // assigned. It hass always the pion mass
0117
0118     if (v0daughter1.p() > v0daughter2.p()) {
0119       v0daughter1_ttrack = theB->build(v0daughter1.bestTrack());
0120       v0daughter2_ttrack = theB->build(v0daughter2.bestTrack());
0121     } else {
0122       v0daughter1_ttrack = theB->build(v0daughter2.bestTrack());
0123       v0daughter2_ttrack = theB->build(v0daughter1.bestTrack());
0124     }
0125
0126     float Track1_mass = (isLambda_) ? bph::PROT_MASS : bph::PI_MASS;
0127     float Track1_sigma = bph::PI_SIGMA;
0128     float Track2_mass = bph::PI_MASS;
0129     float Track2_sigma = bph::PI_SIGMA;
0130     // create V0 vertex
0131     KinVtxFitter fitter(
0132         {v0daughter1_ttrack, v0daughter2_ttrack}, {Track1_mass, Track2_mass}, {Track1_sigma, Track2_sigma});
0133
0134     if (!fitter.success())
0135       continue;
0136
0137     pat::CompositeCandidate cand;
0138     cand.setVertex(reco::Candidate::Point(fitter.fitted_vtx().x(), fitter.fitted_vtx().y(), fitter.fitted_vtx().z()));
0139     auto fit_p4 = fitter.fitted_p4();
0140     cand.setP4(fit_p4);
0141
0142     cand.setCharge(v0daughter1.charge() + v0daughter2.charge());
0143     cand.addUserFloat("sv_chi2", fitter.chi2());
0144     cand.addUserFloat("sv_prob", fitter.prob());
0145     cand.addUserFloat("fitted_mass", fitter.fitted_candidate().mass());
0146     cand.addUserFloat("massErr", sqrt(fitter.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0147     cand.addUserFloat("cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, cand.p4()));
0148     cand.addUserFloat("fitted_cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, fit_p4));
0149     auto lxy = bph::l_xy(fitter, *beamspot);
0150     cand.addUserFloat("l_xy", lxy.value());
0151     cand.addUserFloat("l_xy_unc", lxy.error());
0152
0153     if (!post_vtx_selection_(cand))
0154       continue;
0155
0156     cand.addUserFloat("vtx_x", cand.vx());
0157     cand.addUserFloat("vtx_y", cand.vy());
0158     cand.addUserFloat("vtx_z", cand.vz());
0159
0160     const auto &covMatrix = fitter.fitted_vtx_uncertainty();
0161     cand.addUserFloat("vtx_cxx", covMatrix.cxx());
0162     cand.addUserFloat("vtx_cyy", covMatrix.cyy());
0163     cand.addUserFloat("vtx_czz", covMatrix.czz());
0164     cand.addUserFloat("vtx_cyx", covMatrix.cyx());
0165     cand.addUserFloat("vtx_czx", covMatrix.czx());
0166     cand.addUserFloat("vtx_czy", covMatrix.czy());
0167
0168     cand.addUserFloat("prefit_mass", v0->mass());
0169     int trk1 = 0;
0170     int trk2 = 1;
0171     if (fitter.daughter_p4(0).pt() < fitter.daughter_p4(1).pt()) {
0172       trk1 = 1;
0173       trk2 = 0;
0174     }
0175     cand.addUserFloat("trk1_pt", fitter.daughter_p4(trk1).pt());
0176     cand.addUserFloat("trk1_eta", fitter.daughter_p4(trk1).eta());
0177     cand.addUserFloat("trk1_phi", fitter.daughter_p4(trk1).phi());
0178     cand.addUserFloat("trk2_pt", fitter.daughter_p4(trk2).pt());
0179     cand.addUserFloat("trk2_eta", fitter.daughter_p4(trk2).eta());
0180     cand.addUserFloat("trk2_phi", fitter.daughter_p4(trk2).phi());
0181
0182     // track match
0183     auto trk1_ptr = v0->daughterPtr(trk1);
0184     auto trk1_matched_ref = track_match.get(trk1_ptr.id(), trk1_ptr.key());
0185     auto trk2_ptr = v0->daughterPtr(trk2);
0186     auto trk2_matched_ref = track_match.get(trk2_ptr.id(), trk2_ptr.key());
0187
0188     cand.addUserInt("trk1_idx", trk1_matched_ref.key());
0189     cand.addUserInt("trk2_idx", trk2_matched_ref.key());
0190
0191     // save
0192     ret_val->push_back(cand);
0193     auto V0TT = fitter.fitted_candidate_ttrk();
0194     trans_out->emplace_back(V0TT);
0195   }
0196
0197   evt.put(std::move(ret_val), "SelectedV0Collection");
0198   evt.put(std::move(trans_out), "SelectedV0TransientCollection");
0199 }
0200
0201 #include "FWCore/Framework/interface/MakerMacros.h"
0202 DEFINE_FWK_MODULE(V0ReBuilder);