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 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0039
0040 class V0ReBuilder : public edm::global::EDProducer<> {
0041   // perhaps we need better structure here (begin run etc)
0042 public:
0043   typedef std::vector<reco::TransientTrack> TransientTrackCollection;
0044   typedef std::vector<reco::VertexCompositePtrCandidate> V0Collection;
0045
0046   explicit V0ReBuilder(const edm::ParameterSet &cfg)
0047       : theB_(esConsumes(edm::ESInputTag{"", "TransientTrackBuilder"})),
0048         trk_selection_{cfg.getParameter<std::string>("trkSelection")},
0049         pre_vtx_selection_{cfg.getParameter<std::string>("V0Selection")},
0050         post_vtx_selection_{cfg.getParameter<std::string>("postVtxSelection")},
0051         v0s_{consumes<V0Collection>(cfg.getParameter<edm::InputTag>("V0s"))},
0052         beamspot_{consumes<reco::BeamSpot>(cfg.getParameter<edm::InputTag>("beamSpot"))},
0053         track_match_{consumes<edm::Association<pat::CompositeCandidateCollection>>(
0054             cfg.getParameter<edm::InputTag>("track_match"))},
0055         isLambda_{cfg.getParameter<bool>("isLambda")} {
0056     produces<pat::CompositeCandidateCollection>("SelectedV0Collection");
0057     produces<TransientTrackCollection>("SelectedV0TransientCollection");
0058   }
0059
0060   ~V0ReBuilder() override {}
0061
0062   void produce(edm::StreamID, edm::Event &, const edm::EventSetup &) const override;
0063
0064 private:
0065   const edm::ESGetToken<TransientTrackBuilder, TransientTrackRecord> theB_;
0066   const StringCutObjectSelector<pat::PackedCandidate> trk_selection_;
0067   const StringCutObjectSelector<reco::VertexCompositePtrCandidate> pre_vtx_selection_;
0068   const StringCutObjectSelector<pat::CompositeCandidate> post_vtx_selection_;
0069   const edm::EDGetTokenT<V0Collection> v0s_;
0070   const edm::EDGetTokenT<reco::BeamSpot> beamspot_;
0071   const edm::EDGetTokenT<edm::Association<pat::CompositeCandidateCollection>> track_match_;
0072   const bool isLambda_;
0073 };
0074
0075 void V0ReBuilder::produce(edm::StreamID, edm::Event &evt, edm::EventSetup const &iSetup) const {
0076   // input
0077   auto const theB = &iSetup.getData(theB_);
0078   edm::Handle<V0Collection> V0s;
0079   evt.getByToken(v0s_, V0s);
0080   edm::Handle<reco::BeamSpot> beamspot;
0081   evt.getByToken(beamspot_, beamspot);
0082
0083   auto &track_match = evt.get(track_match_);
0084
0085   // output
0086   std::unique_ptr<pat::CompositeCandidateCollection> ret_val(new pat::CompositeCandidateCollection());
0087   std::unique_ptr<TransientTrackCollection> trans_out(new TransientTrackCollection);
0088
0089   for (reco::VertexCompositePtrCandidateCollection::const_iterator v0 = V0s->begin(); v0 != V0s->end(); v0++) {
0090     // selection on V0s
0091     if (v0->numberOfDaughters() != 2)
0092       continue;
0093     if (!pre_vtx_selection_(*v0))
0094       continue;
0095
0096     pat::PackedCandidate v0daughter1 = *(dynamic_cast<const pat::PackedCandidate *>(v0->daughter(0)));
0097     pat::PackedCandidate v0daughter2 = *(dynamic_cast<const pat::PackedCandidate *>(v0->daughter(1)));
0098
0099     if (!v0daughter1.hasTrackDetails())
0100       continue;
0101     if (!v0daughter2.hasTrackDetails())
0102       continue;
0103
0104     if (abs(v0daughter1.pdgId()) != 211)
0105       continue;  // This cut do not affect the Lambda->proton pion candidates
0106     if (abs(v0daughter2.pdgId()) != 211)
0107       continue;  // This cut do not affect the Lambda->proton pion candidates
0108
0109     if (!trk_selection_(v0daughter1) || !trk_selection_(v0daughter2))
0110       continue;
0111
0112     reco::TransientTrack v0daughter1_ttrack;  // 1st daughter, leading daughter to be assigned.
0113                                               // Proton mass will be assigned for the
0114                                               // Lambda->Proton Pion mode, Pion mass will be
0115                                               // assigned for the Kshort->PionPion mode.
0116     reco::TransientTrack v0daughter2_ttrack;  // 2nd daughter, subleading daughter to be
0117                                               // assigned. It hass always the pion mass
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     float Track1_mass = (isLambda_) ? bph::PROT_MASS : bph::PI_MASS;
0126     float Track1_sigma = bph::PI_SIGMA;
0127     float Track2_mass = bph::PI_MASS;
0128     float Track2_sigma = bph::PI_SIGMA;
0129     // create V0 vertex
0130     KinVtxFitter fitter;
0131
0132     try {
0133       fitter = KinVtxFitter(
0134           {v0daughter1_ttrack, v0daughter2_ttrack}, {Track1_mass, Track2_mass}, {Track1_sigma, Track2_sigma});
0135     } catch (const VertexException &e) {
0136       edm::LogWarning("KinematicFit") << "Skipping candidate due to fit failure: " << e.what();
0137       continue;  // Just skip this candidate and continue with the loop
0138     }
0139
0140     if (!fitter.success())
0141       continue;
0142     pat::CompositeCandidate cand;
0143     cand.setVertex(reco::Candidate::Point(fitter.fitted_vtx().x(), fitter.fitted_vtx().y(), fitter.fitted_vtx().z()));
0144     auto fit_p4 = fitter.fitted_p4();
0145     cand.setP4(fit_p4);
0146     cand.setCharge(v0daughter1.charge() + v0daughter2.charge());
0147     cand.addUserFloat("sv_chi2", fitter.chi2());
0148     cand.addUserFloat("sv_prob", fitter.prob());
0149     cand.addUserFloat("fitted_mass", fitter.fitted_candidate().mass());
0150     cand.addUserFloat("fitted_pt", fitter.fitted_p4().pt());
0151     cand.addUserFloat("fitted_eta", fitter.fitted_p4().eta());
0152     cand.addUserFloat("fitted_phi", fitter.fitted_p4().phi());
0153     cand.addUserFloat("massErr", sqrt(fitter.fitted_candidate().kinematicParametersError().matrix()(6, 6)));
0154     cand.addUserFloat("cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, cand.p4()));
0155     cand.addUserFloat("fitted_cos_theta_2D", bph::cos_theta_2D(fitter, *beamspot, fit_p4));
0156     auto lxy = bph::l_xy(fitter, *beamspot);
0157     cand.addUserFloat("l_xy", lxy.value());
0158     cand.addUserFloat("l_xy_unc", lxy.error());
0159     if (!post_vtx_selection_(cand))
0160       continue;
0161     const reco::BeamSpot &beamSpot = *beamspot;
0162     TrajectoryStateClosestToPoint theDCAXBS = fitter.fitted_candidate_ttrk().trajectoryStateClosestToPoint(
0163         GlobalPoint(beamSpot.position().x(), beamSpot.position().y(), beamSpot.position().z()));
0164     double DCAB0BS = -99.;
0165     double DCAB0BSErr = -99.;
0166
0167     if (theDCAXBS.isValid() == true) {
0168       DCAB0BS = theDCAXBS.perigeeParameters().transverseImpactParameter();
0169       DCAB0BSErr = theDCAXBS.perigeeError().transverseImpactParameterError();
0170     }
0171     cand.addUserFloat("dca", DCAB0BS);
0172     cand.addUserFloat("dcaErr", DCAB0BSErr);
0173
0174     cand.addUserFloat("vtx_x", cand.vx());
0175     cand.addUserFloat("vtx_y", cand.vy());
0176     cand.addUserFloat("vtx_z", cand.vz());
0177
0178     const auto &covMatrix = fitter.fitted_vtx_uncertainty();
0179     cand.addUserFloat("vtx_cxx", covMatrix.cxx());
0180     cand.addUserFloat("vtx_cyy", covMatrix.cyy());
0181     cand.addUserFloat("vtx_czz", covMatrix.czz());
0182     cand.addUserFloat("vtx_cyx", covMatrix.cyx());
0183     cand.addUserFloat("vtx_czx", covMatrix.czx());
0184     cand.addUserFloat("vtx_czy", covMatrix.czy());
0185     cand.addUserFloat("prefit_mass", v0->mass());
0186     int trk1 = 0;
0187     int trk2 = 1;
0188     if (fitter.daughter_p4(0).pt() < fitter.daughter_p4(1).pt()) {
0189       trk1 = 1;
0190       trk2 = 0;
0191     }
0192     cand.addUserFloat("trk1_pt", fitter.daughter_p4(trk1).pt());
0193     cand.addUserFloat("trk1_eta", fitter.daughter_p4(trk1).eta());
0194     cand.addUserFloat("trk1_phi", fitter.daughter_p4(trk1).phi());
0195     cand.addUserFloat("trk2_pt", fitter.daughter_p4(trk2).pt());
0196     cand.addUserFloat("trk2_eta", fitter.daughter_p4(trk2).eta());
0197     cand.addUserFloat("trk2_phi", fitter.daughter_p4(trk2).phi());
0198     // track match
0199     auto trk1_ptr = v0->daughterPtr(trk1);
0200     auto trk1_matched_ref = track_match.get(trk1_ptr.id(), trk1_ptr.key());
0201     auto trk2_ptr = v0->daughterPtr(trk2);
0202     auto trk2_matched_ref = track_match.get(trk2_ptr.id(), trk2_ptr.key());
0203
0204     cand.addUserInt("trk1_idx", trk1_matched_ref.key());
0205     cand.addUserInt("trk2_idx", trk2_matched_ref.key());
0206     cand.addUserInt("fit_trk1_charge", (int)v0daughter1_ttrack.charge());
0207     cand.addUserInt("fit_trk2_charge", (int)v0daughter2_ttrack.charge());
0208     // save
0209     ret_val->push_back(cand);
0210     auto V0TT = fitter.fitted_candidate_ttrk();
0211     trans_out->emplace_back(V0TT);
0212   }
0213
0214   evt.put(std::move(ret_val), "SelectedV0Collection");
0215   evt.put(std::move(trans_out), "SelectedV0TransientCollection");
0216 }
0217
0218 #include "FWCore/Framework/interface/MakerMacros.h"
0219 DEFINE_FWK_MODULE(V0ReBuilder);