Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:29:15

0001 #ifndef RecoVertex_PrimaryVertexProducer_WeightedMeanFitter_h
0002 #define RecoVertex_PrimaryVertexProducer_WeightedMeanFitter_h
0003 
0004 #include <vector>
0005 #include "TrackingTools/TransientTrack/interface/TransientTrack.h"
0006 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0007 #include "DataFormats/BeamSpot/interface/BeamSpot.h"
0008 #include "RecoVertex/VertexPrimitives/interface/TransientVertex.h"
0009 #include "RecoVertex/PrimaryVertexProducer/interface/PrimaryVertexFitterBase.h"
0010 
0011 namespace WeightedMeanFitter {
0012 
0013   constexpr float startError = 20.0;
0014   constexpr float precision = 1e-24;
0015   constexpr float corr_x = 1.2;
0016   constexpr float corr_x_bs = 1.0;  // corr_x for beam spot
0017   constexpr float corr_z = 1.4;
0018   constexpr int maxIterations = 50;
0019   constexpr float muSquare = 9.;
0020 
0021   inline std::pair<GlobalPoint, double> nearestPoint(const GlobalPoint& vertex, reco::Track iclus) {
0022     double ox = iclus.vx();
0023     double oy = iclus.vy();
0024     double oz = iclus.vz();
0025 
0026     double vx = iclus.px();
0027     double vy = iclus.py();
0028     double vz = iclus.pz();
0029 
0030     double opx = vertex.x() - ox;
0031     double opy = vertex.y() - oy;
0032     double opz = vertex.z() - oz;
0033 
0034     double vnorm2 = (vx * vx + vy * vy + vz * vz);
0035     double t = (vx * opx + vy * opy + vz * opz) / (vnorm2);
0036 
0037     GlobalPoint p(ox + t * vx, oy + t * vy, oz + t * vz);
0038     return std::pair<GlobalPoint, double>(
0039         p,
0040         std::sqrt(std::pow(p.x() - vertex.x(), 2) + std::pow(p.y() - vertex.y(), 2) + std::pow(p.z() - vertex.z(), 2)));
0041   }
0042 
0043   inline TransientVertex weightedMeanOutlierRejection(const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0044                                                       std::vector<reco::TransientTrack> iclus) {
0045     float x = 0., y = 0., z = 0.;
0046     float s_wx = 0., s_wz = 0.;
0047     float s2_wx = 0., s2_wz = 0.;
0048     float wx = 0., wz = 0., chi2 = 0.;
0049     float ndof_x = 0.;
0050 
0051     AlgebraicSymMatrix33 err;
0052     err(0, 0) = startError / 10 * startError / 10;
0053     err(1, 1) = startError / 10 * startError / 10;
0054     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0055     for (const auto& p : points) {
0056       wx = p.second.x() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.x(), 2);
0057 
0058       wz = p.second.z() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.z(), 2);
0059 
0060       x += p.first.x() * wx;
0061       y += p.first.y() * wx;
0062       z += p.first.z() * wz;
0063 
0064       s_wx += wx;
0065       s_wz += wz;
0066     }
0067 
0068     if (s_wx == 0. || s_wz == 0.) {
0069       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0070       return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0071     }
0072 
0073     x /= s_wx;
0074     y /= s_wx;
0075     z /= s_wz;
0076 
0077     float old_x, old_y, old_z;
0078 
0079     float xpull;
0080     int niter = 0;
0081 
0082     float err_x, err_z;
0083 
0084     err_x = 1. / s_wx;
0085     err_z = 1. / s_wz;
0086 
0087     while ((niter++) < 2) {
0088       old_x = x;
0089       old_y = y;
0090       old_z = z;
0091       s_wx = 0;
0092       s_wz = 0;
0093       s2_wx = 0;
0094       s2_wz = 0;
0095       x = 0;
0096       y = 0;
0097       z = 0;
0098       ndof_x = 0;
0099 
0100       for (unsigned int i = 0; i < (unsigned int)points.size(); i++) {
0101         std::pair<GlobalPoint, double> p = nearestPoint(GlobalPoint(old_x, old_y, old_z), (iclus)[i].track());
0102 
0103         wx = points[i].second.x() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.x(), 2);
0104         wz = points[i].second.z() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.z(), 2);
0105 
0106         xpull = 0.;
0107 
0108         if (std::pow(p.first.x() - old_x, 2) / (wx + err_x) < muSquare &&
0109             std::pow(p.first.y() - old_y, 2) / (wx + err_x) < muSquare &&
0110             std::pow(p.first.z() - old_z, 2) / (wz + err_z) < muSquare)
0111           xpull = 1.;
0112 
0113         ndof_x += xpull;
0114 
0115         wx = xpull / wx;
0116         wz = xpull / wz;
0117 
0118         x += wx * p.first.x();
0119         y += wx * p.first.y();
0120         z += wz * p.first.z();
0121 
0122         s_wx += wx;
0123         s_wz += wz;
0124 
0125         s2_wx += wx * xpull;
0126         s2_wz += wz * xpull;
0127       }
0128 
0129       if (s_wx == 0. || s_wz == 0.) {
0130         edm::LogWarning("WeightedMeanFitter")
0131             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0132         return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0133       }
0134       x /= s_wx;
0135       y /= s_wx;
0136       z /= s_wz;
0137 
0138       err_x = (s2_wx / std::pow(s_wx, 2));
0139       err_z = (s2_wz / std::pow(s_wz, 2));
0140 
0141       if (std::abs(x - old_x) < (precision / 1.) && std::abs(y - old_y) < (precision / 1.) &&
0142           std::abs(z - old_z) < (precision / 1.)) {
0143         break;
0144       }
0145     }
0146 
0147     err(0, 0) = err_x * corr_x * corr_x;
0148     err(1, 1) = err_x * corr_x * corr_x;
0149     err(2, 2) = err_z * corr_z * corr_z;
0150 
0151     float dist = 0;
0152     for (const auto& p : points) {
0153       wx = p.second.x();
0154       wx = wx <= precision ? precision : wx;
0155 
0156       wz = p.second.z();
0157       wz = wz <= precision ? precision : wz;
0158 
0159       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + err(0, 0));
0160       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + err(1, 1));
0161       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + err(2, 2));
0162       chi2 += dist;
0163     }
0164     float ndof =
0165         ndof_x > 1 ? (2 * ndof_x - 3) : 0.00001;  // ndof_x is actually the number of tracks with non-zero weight
0166     TransientVertex v(GlobalPoint(x, y, z), err, iclus, chi2, ndof);
0167     return v;
0168   }
0169 
0170   inline TransientVertex weightedMeanOutlierRejectionBeamSpot(
0171       const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0172       std::vector<reco::TransientTrack> iclus,
0173       const reco::BeamSpot& beamSpot) {
0174     float x = 0., y = 0., z = 0.;
0175     float s_wx = 0., s_wz = 0.;
0176     float s2_wx = 0., s2_wz = 0.;
0177     float wx = 0., wz = 0., chi2 = 0.;
0178     float wy = 0., s_wy = 0., s2_wy = 0.;
0179     float ndof_x = 0.;
0180 
0181     AlgebraicSymMatrix33 err;
0182     err(0, 0) = startError / 10 * startError / 10;
0183     err(1, 1) = startError / 10 * startError / 10;
0184     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0185 
0186     GlobalError bse(beamSpot.rotatedCovariance3D());
0187     GlobalPoint bsp(Basic3DVector<float>(beamSpot.position()));
0188 
0189     for (const auto& p : points) {
0190       wx = p.second.x() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.x(), 2);
0191       wy = p.second.y() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.y(), 2);
0192 
0193       wz = p.second.z() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.z(), 2);
0194 
0195       x += p.first.x() * wx;
0196       y += p.first.y() * wy;
0197       z += p.first.z() * wz;
0198 
0199       s_wx += wx;
0200       s_wy += wy;
0201       s_wz += wz;
0202     }
0203 
0204     if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0205       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0206       return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0207     }
0208     // use the square of covariance element to increase it's weight: it will be the most important
0209     wx = bse.cxx() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(bse.cxx(), 2);
0210     wy = bse.cyy() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(bse.cyy(), 2);
0211 
0212     x += bsp.x() * wx;
0213     y += bsp.y() * wy;
0214 
0215     x /= (s_wx + wx);
0216     y /= (s_wy + wy);
0217     z /= s_wz;
0218 
0219     float old_x, old_y, old_z;
0220 
0221     float xpull;
0222     int niter = 0;
0223 
0224     float err_x, err_y, err_z;
0225 
0226     err_x = 1. / s_wx;
0227     err_y = 1. / s_wy;
0228     err_z = 1. / s_wz;
0229 
0230     while ((niter++) < 2) {
0231       old_x = x;
0232       old_y = y;
0233       old_z = z;
0234       s_wx = 0;
0235       s_wz = 0;
0236       s2_wx = 0;
0237       s2_wz = 0;
0238 
0239       s_wy = 0;
0240       s2_wy = 0;
0241 
0242       x = 0;
0243       y = 0;
0244       z = 0;
0245       ndof_x = 0;
0246 
0247       for (unsigned int i = 0; i < (unsigned int)points.size(); i++) {
0248         std::pair<GlobalPoint, double> p = nearestPoint(GlobalPoint(old_x, old_y, old_z), (iclus)[i].track());
0249 
0250         wx = points[i].second.x() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.x(), 2);
0251         wy = points[i].second.y() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.y(), 2);
0252 
0253         wz = points[i].second.z() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.z(), 2);
0254 
0255         xpull = 0.;
0256         if (std::pow(p.first.x() - old_x, 2) / (wx + err_x) < muSquare &&
0257             std::pow(p.first.y() - old_y, 2) / (wy + err_y) < muSquare &&
0258             std::pow(p.first.z() - old_z, 2) / (wz + err_z) < muSquare)
0259           xpull = 1.;
0260 
0261         ndof_x += xpull;
0262 
0263         wx = xpull / wx;
0264         wy = xpull / wy;
0265         wz = xpull / wz;
0266 
0267         x += wx * p.first.x();
0268         y += wy * p.first.y();
0269         z += wz * p.first.z();
0270 
0271         s_wx += wx;
0272         s_wy += wy;
0273         s_wz += wz;
0274 
0275         s2_wx += wx * xpull;
0276         s2_wy += wy * xpull;
0277         s2_wz += wz * xpull;
0278       }
0279 
0280       if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0281         edm::LogWarning("WeightedMeanFitter")
0282             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0283         return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0284       }
0285       wx = bse.cxx() <= std::pow(precision, 2) ? 1. / std::pow(precision, 2) : 1. / bse.cxx();
0286       wy = bse.cyy() <= std::pow(precision, 2) ? 1. / std::pow(precision, 2) : 1. / bse.cyy();
0287 
0288       x += bsp.x() * wx;
0289       y += bsp.y() * wy;
0290       s_wx += wx;
0291       s2_wx += wx;
0292       s_wy += wy;
0293       s2_wy += wy;
0294 
0295       x /= s_wx;
0296       y /= s_wy;
0297       z /= s_wz;
0298 
0299       err_x = (s2_wx / std::pow(s_wx, 2));
0300       err_y = (s2_wy / std::pow(s_wy, 2));
0301       err_z = (s2_wz / std::pow(s_wz, 2));
0302 
0303       if (std::abs(x - old_x) < (precision) && std::abs(y - old_y) < (precision) && std::abs(z - old_z) < (precision)) {
0304         break;
0305       }
0306     }
0307     err(0, 0) = err_x * corr_x_bs * corr_x_bs;
0308     err(1, 1) = err_y * corr_x_bs * corr_x_bs;
0309     err(2, 2) = err_z * corr_z * corr_z;
0310 
0311     float dist = 0;
0312     for (const auto& p : points) {
0313       wx = p.second.x();
0314       wx = wx <= precision ? precision : wx;
0315 
0316       wz = p.second.z();
0317       wz = wz <= precision ? precision : wz;
0318 
0319       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + err(0, 0));
0320       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + err(1, 1));
0321       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + err(2, 2));
0322       chi2 += dist;
0323     }
0324     TransientVertex v(GlobalPoint(x, y, z), err, iclus, chi2, (int)ndof_x);
0325     return v;
0326   }
0327 
0328   inline TransientVertex weightedMeanOutlierRejectionVarianceAsError(
0329       const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0330       std::vector<std::vector<reco::TransientTrack>>::const_iterator iclus) {
0331     float x = 0, y = 0, z = 0, s_wx = 0, s_wy = 0, s_wz = 0, wx = 0, wy = 0, wz = 0, chi2 = 0;
0332     float ndof_x = 0;
0333     AlgebraicSymMatrix33 err;
0334     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0335     err(0, 0) = err(1, 1) = err(2, 2) / 100.;
0336 
0337     for (const auto& p : points) {
0338       wx = p.second.x();
0339       wx = wx <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(wx, 2);
0340 
0341       wz = p.second.z();
0342       wz = wz <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(wz, 2);
0343 
0344       x += p.first.x() * wx;
0345       y += p.first.y() * wx;
0346       z += p.first.z() * wz;
0347 
0348       s_wx += wx;
0349       s_wz += wz;
0350     }
0351 
0352     if (s_wx == 0. || s_wz == 0.) {
0353       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0354       return TransientVertex(GlobalPoint(0, 0, 0), err, *iclus, 0, 0);
0355     }
0356 
0357     x /= s_wx;
0358     y /= s_wx;
0359     z /= s_wz;
0360 
0361     float old_x, old_y, old_z;
0362     float xpull;
0363     int niter = 0;
0364     float err_x, err_y, err_z;
0365     err_x = 1. / s_wx;
0366     err_y = 1. / s_wx;
0367     err_z = 1. / s_wz;
0368     float s_err_x = 0., s_err_y = 0., s_err_z = 0.;
0369     while ((niter++) < maxIterations) {
0370       old_x = x;
0371       old_y = y;
0372       old_z = z;
0373       s_wx = 0;
0374       s_wy = 0;
0375       s_wz = 0;
0376       x = 0;
0377       y = 0;
0378       z = 0;
0379       s_err_x = 0.;
0380       s_err_y = 0.;
0381       s_err_z = 0.;
0382 
0383       for (const auto& p : points) {
0384         wx = p.second.x();
0385         wx = wx <= precision ? precision : wx;
0386 
0387         wy = wx * wx + err_y;
0388         wx = wx * wx + err_x;
0389 
0390         wz = p.second.z();
0391         wz = wz <= precision ? precision : wz;
0392         wz = wz * wz + err_z;
0393 
0394         xpull = std::pow((p.first.x() - old_x), 2) / wx;
0395         xpull += std::pow((p.first.y() - old_y), 2) / wy;
0396         xpull += std::pow((p.first.z() - old_z), 2) / wz;
0397         xpull = 1. / (1. + std::exp(-0.5 * ((muSquare)-xpull)));
0398         ndof_x += xpull;
0399 
0400         wx = 1. / wx;
0401         wy = 1. / wy;
0402         wz = 1. / wz;
0403 
0404         wx *= xpull;
0405         wy *= xpull;
0406         wz *= xpull;
0407 
0408         x += wx * p.first.x();
0409         y += wy * p.first.y();
0410         z += wz * p.first.z();
0411 
0412         s_wx += wx;
0413         s_wy += wy;
0414         s_wz += wz;
0415 
0416         s_err_x += wx * pow(p.first.x() - old_x, 2);
0417         s_err_y += wy * pow(p.first.y() - old_y, 2);
0418         s_err_z += wz * pow(p.first.z() - old_z, 2);
0419       }
0420       if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0421         edm::LogWarning("WeightedMeanFitter")
0422             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0423         return TransientVertex(GlobalPoint(0, 0, 0), err, *iclus, 0, 0);
0424       }
0425       x /= s_wx;
0426       y /= s_wy;
0427       z /= s_wz;
0428 
0429       err_x = s_err_x / s_wx;
0430       err_y = s_err_y / s_wy;
0431       err_z = s_err_z / s_wz;
0432 
0433       if (std::abs(x - old_x) < (precision / 1.) && std::abs(y - old_y) < (precision / 1.) &&
0434           std::abs(z - old_z) < (precision / 1.))
0435         break;
0436     }
0437 
0438     err(0, 0) = err_x;
0439     err(1, 1) = err_y;
0440     err(2, 2) = err_z;
0441 
0442     float dist = 0.f;
0443     for (const auto& p : points) {
0444       wx = p.second.x();
0445       wx = wx <= precision ? precision : wx;
0446 
0447       wz = p.second.z();
0448       wz = wz <= precision ? precision : wz;
0449 
0450       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + std::pow(err(0, 0), 2));
0451       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + std::pow(err(1, 1), 2));
0452       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + std::pow(err(2, 2), 2));
0453       chi2 += dist;
0454     }
0455     TransientVertex v(GlobalPoint(x, y, z), err, *iclus, chi2, (int)ndof_x);
0456     return v;
0457   }
0458 
0459 };  // namespace WeightedMeanFitter
0460 
0461 // adapter for the multiprimaryvertexfitter scheme
0462 // this code was originally introduced as part of PrimaryVertexProducer.cc
0463 // by Adriano Di Florio <AdrianoDee>, Giorgio Pizzati <giorgiopizz> et.al. in #39995, then moved here with minor modifications
0464 class WeightedMeanPrimaryVertexEstimator : public PrimaryVertexFitterBase {
0465 public:
0466   WeightedMeanPrimaryVertexEstimator() = default;
0467   ~WeightedMeanPrimaryVertexEstimator() override = default;
0468 
0469   std::vector<TransientVertex> fit(const std::vector<reco::TransientTrack>& dummy,
0470                                    const std::vector<TransientVertex>& clusters,
0471                                    const reco::BeamSpot& beamSpot,
0472                                    const bool useBeamConstraint) override {
0473     std::vector<TransientVertex> pvs;
0474     std::vector<TransientVertex> seed(1);
0475 
0476     for (auto& cluster : clusters) {
0477       if (cluster.originalTracks().size() > 1) {
0478         std::vector<reco::TransientTrack> tracklist = cluster.originalTracks();
0479         TransientVertex::TransientTrackToFloatMap trkWeightMap;
0480         std::vector<std::pair<GlobalPoint, GlobalPoint>> points;
0481         if (useBeamConstraint && (tracklist.size() > 1)) {
0482           for (const auto& itrack : tracklist) {
0483             GlobalPoint p = itrack.stateAtBeamLine().trackStateAtPCA().position();
0484             GlobalPoint err(itrack.stateAtBeamLine().transverseImpactParameter().error(),
0485                             itrack.stateAtBeamLine().transverseImpactParameter().error(),
0486                             itrack.track().dzError());
0487             std::pair<GlobalPoint, GlobalPoint> p2(p, err);
0488             points.push_back(p2);
0489           }
0490 
0491           TransientVertex v = WeightedMeanFitter::weightedMeanOutlierRejectionBeamSpot(points, tracklist, beamSpot);
0492           if (!v.hasTrackWeight()) {
0493             // if the fitter doesn't provide weights, fill dummy values
0494             TransientVertex::TransientTrackToFloatMap trkWeightMap;
0495             for (const auto& trk : v.originalTracks()) {
0496               trkWeightMap[trk] = 1.;
0497             }
0498             v.weightMap(trkWeightMap);
0499           }
0500           if ((v.positionError().matrix())(2, 2) != (WeightedMeanFitter::startError * WeightedMeanFitter::startError))
0501             pvs.push_back(v);
0502         } else if (!(useBeamConstraint) && (tracklist.size() > 1)) {
0503           for (const auto& itrack : tracklist) {
0504             GlobalPoint p = itrack.impactPointState().globalPosition();
0505             GlobalPoint err(itrack.track().dxyError(), itrack.track().dxyError(), itrack.track().dzError());
0506             std::pair<GlobalPoint, GlobalPoint> p2(p, err);
0507             points.push_back(p2);
0508           }
0509 
0510           TransientVertex v = WeightedMeanFitter::weightedMeanOutlierRejection(points, tracklist);
0511           if (!v.hasTrackWeight()) {
0512             // if the fitter doesn't provide weights, fill dummy values
0513             TransientVertex::TransientTrackToFloatMap trkWeightMap;
0514             for (const auto& trk : v.originalTracks()) {
0515               trkWeightMap[trk] = 1.;
0516             }
0517             v.weightMap(trkWeightMap);
0518           }
0519           if ((v.positionError().matrix())(2, 2) != (WeightedMeanFitter::startError * WeightedMeanFitter::startError))
0520             pvs.push_back(v);  //FIX with constants
0521         }
0522       }
0523     }
0524     return pvs;
0525   }
0526 };
0527 
0528 #endif