BPHNano/plugins/KinVtxFitter.cc

0001 // original author: RK18 team
0002 #include "KinVtxFitter.h"
0003
0004 #include "RecoVertex/KinematicFit/interface/KinematicConstrainedVertexFitter.h"
0005 #include "RecoVertex/KinematicFit/interface/KinematicParticleVertexFitter.h"
0006 #include "RecoVertex/KinematicFit/interface/TwoTrackMassKinematicConstraint.h"  // MIGHT be useful for Phi->KK?
0007 #include "RecoVertex/KinematicFitPrimitives/interface/KinematicParticleFactoryFromTransientTrack.h"
0008
0009 KinVtxFitter::KinVtxFitter(const std::vector<reco::TransientTrack> tracks,
0010                            const std::vector<double> masses,
0011                            std::vector<float> sigmas)
0012     : n_particles_{masses.size()} {
0013   KinematicParticleFactoryFromTransientTrack factory;
0014   std::vector<RefCountedKinematicParticle> particles;
0015   for (size_t i = 0; i < tracks.size(); ++i) {
0016     particles.emplace_back(factory.particle(tracks.at(i), masses.at(i), kin_chi2_, kin_ndof_, sigmas[i]));
0017   }
0018
0019   KinematicParticleVertexFitter kcv_fitter;
0020   RefCountedKinematicTree vtx_tree = kcv_fitter.fit(particles);
0021
0022   if (vtx_tree->isEmpty() || !vtx_tree->isValid() || !vtx_tree->isConsistent()) {
0023     success_ = false;
0024     return;
0025   }
0026
0027   vtx_tree->movePointerToTheTop();
0028   fitted_particle_ = vtx_tree->currentParticle();
0029   fitted_vtx_ = vtx_tree->currentDecayVertex();
0030   if (!fitted_particle_->currentState().isValid() || !fitted_vtx_->vertexIsValid()) {
0031     success_ = false;
0032     return;
0033   }
0034   fitted_state_ = fitted_particle_->currentState();
0035   fitted_children_ = vtx_tree->finalStateParticles();
0036   if (fitted_children_.size() != n_particles_) {
0037     success_ = false;
0038     return;
0039   }
0040   fitted_track_ = fitted_particle_->refittedTransientTrack();
0041   success_ = true;
0042 }
0043
0044 KinVtxFitter::KinVtxFitter(const std::vector<reco::TransientTrack> tracks,
0045                            const std::vector<double> masses,
0046                            std::vector<float> sigmas,
0047                            ParticleMass dilep_mass)
0048     : n_particles_{masses.size()} {
0049   KinematicParticleFactoryFromTransientTrack factory;
0050   std::vector<RefCountedKinematicParticle> particles;
0051   for (size_t i = 0; i < tracks.size(); ++i) {
0052     particles.emplace_back(factory.particle(tracks.at(i), masses.at(i), kin_chi2_, kin_ndof_, sigmas[i]));
0053   }
0054
0055   auto dilep_c = std::make_unique<TwoTrackMassKinematicConstraint>(dilep_mass);
0056   KinematicConstrainedVertexFitter kcv_fitter;
0057   RefCountedKinematicTree vtx_tree = kcv_fitter.fit(particles, dilep_c.get());
0058
0059   if (vtx_tree->isEmpty() || !vtx_tree->isValid() || !vtx_tree->isConsistent()) {
0060     success_ = false;
0061     return;
0062   }
0063
0064   vtx_tree->movePointerToTheTop();
0065   fitted_particle_ = vtx_tree->currentParticle();
0066   fitted_vtx_ = vtx_tree->currentDecayVertex();
0067   if (!fitted_particle_->currentState().isValid() || !fitted_vtx_->vertexIsValid()) {
0068     success_ = false;
0069     return;
0070   }
0071   fitted_state_ = fitted_particle_->currentState();
0072   fitted_children_ = vtx_tree->finalStateParticles();
0073   if (fitted_children_.size() != n_particles_) {
0074     success_ = false;
0075     return;
0076   }
0077   fitted_track_ = fitted_particle_->refittedTransientTrack();
0078   success_ = true;
0079 }