PrimaryVertexProducer/interface/WeightedMeanFitter.h

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
0010 namespace WeightedMeanFitter {
0011
0012   constexpr float startError = 20.0;
0013   constexpr float precision = 1e-24;
0014   constexpr float corr_x = 1.2;
0015   constexpr float corr_x_bs = 1.0;  // corr_x for beam spot
0016   constexpr float corr_z = 1.4;
0017   constexpr int maxIterations = 50;
0018   constexpr float muSquare = 9.;
0019
0020   inline std::pair<GlobalPoint, double> nearestPoint(const GlobalPoint& vertex, reco::Track iclus) {
0021     double ox = iclus.vx();
0022     double oy = iclus.vy();
0023     double oz = iclus.vz();
0024
0025     double vx = iclus.px();
0026     double vy = iclus.py();
0027     double vz = iclus.pz();
0028
0029     double opx = vertex.x() - ox;
0030     double opy = vertex.y() - oy;
0031     double opz = vertex.z() - oz;
0032
0033     double vnorm2 = (vx * vx + vy * vy + vz * vz);
0034     double t = (vx * opx + vy * opy + vz * opz) / (vnorm2);
0035
0036     GlobalPoint p(ox + t * vx, oy + t * vy, oz + t * vz);
0037     return std::pair<GlobalPoint, double>(
0038         p,
0039         std::sqrt(std::pow(p.x() - vertex.x(), 2) + std::pow(p.y() - vertex.y(), 2) + std::pow(p.z() - vertex.z(), 2)));
0040   }
0041
0042   inline TransientVertex weightedMeanOutlierRejection(const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0043                                                       std::vector<reco::TransientTrack> iclus) {
0044     float x = 0., y = 0., z = 0.;
0045     float s_wx = 0., s_wz = 0.;
0046     float s2_wx = 0., s2_wz = 0.;
0047     float wx = 0., wz = 0., chi2 = 0.;
0048     float ndof_x = 0.;
0049
0050     AlgebraicSymMatrix33 err;
0051     err(0, 0) = startError / 10 * startError / 10;
0052     err(1, 1) = startError / 10 * startError / 10;
0053     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0054     for (const auto& p : points) {
0055       wx = p.second.x() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.x(), 2);
0056
0057       wz = p.second.z() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.z(), 2);
0058
0059       x += p.first.x() * wx;
0060       y += p.first.y() * wx;
0061       z += p.first.z() * wz;
0062
0063       s_wx += wx;
0064       s_wz += wz;
0065     }
0066
0067     if (s_wx == 0. || s_wz == 0.) {
0068       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0069       return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0070     }
0071
0072     x /= s_wx;
0073     y /= s_wx;
0074     z /= s_wz;
0075
0076     float old_x, old_y, old_z;
0077
0078     float xpull;
0079     int niter = 0;
0080
0081     float err_x, err_z;
0082
0083     err_x = 1. / s_wx;
0084     err_z = 1. / s_wz;
0085
0086     while ((niter++) < 2) {
0087       old_x = x;
0088       old_y = y;
0089       old_z = z;
0090       s_wx = 0;
0091       s_wz = 0;
0092       s2_wx = 0;
0093       s2_wz = 0;
0094       x = 0;
0095       y = 0;
0096       z = 0;
0097       ndof_x = 0;
0098
0099       for (unsigned int i = 0; i < (unsigned int)points.size(); i++) {
0100         std::pair<GlobalPoint, double> p = nearestPoint(GlobalPoint(old_x, old_y, old_z), (iclus)[i].track());
0101
0102         wx = points[i].second.x() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.x(), 2);
0103         wz = points[i].second.z() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.z(), 2);
0104
0105         xpull = 0.;
0106
0107         if (std::pow(p.first.x() - old_x, 2) / (wx + err_x) < muSquare &&
0108             std::pow(p.first.y() - old_y, 2) / (wx + err_x) < muSquare &&
0109             std::pow(p.first.z() - old_z, 2) / (wz + err_z) < muSquare)
0110           xpull = 1.;
0111
0112         ndof_x += xpull;
0113
0114         wx = xpull / wx;
0115         wz = xpull / wz;
0116
0117         x += wx * p.first.x();
0118         y += wx * p.first.y();
0119         z += wz * p.first.z();
0120
0121         s_wx += wx;
0122         s_wz += wz;
0123
0124         s2_wx += wx * xpull;
0125         s2_wz += wz * xpull;
0126       }
0127
0128       if (s_wx == 0. || s_wz == 0.) {
0129         edm::LogWarning("WeightedMeanFitter")
0130             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0131         return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0132       }
0133       x /= s_wx;
0134       y /= s_wx;
0135       z /= s_wz;
0136
0137       err_x = (s2_wx / std::pow(s_wx, 2));
0138       err_z = (s2_wz / std::pow(s_wz, 2));
0139
0140       if (std::abs(x - old_x) < (precision / 1.) && std::abs(y - old_y) < (precision / 1.) &&
0141           std::abs(z - old_z) < (precision / 1.)) {
0142         break;
0143       }
0144     }
0145
0146     err(0, 0) = err_x * corr_x * corr_x;
0147     err(1, 1) = err_x * corr_x * corr_x;
0148     err(2, 2) = err_z * corr_z * corr_z;
0149
0150     float dist = 0;
0151     for (const auto& p : points) {
0152       wx = p.second.x();
0153       wx = wx <= precision ? precision : wx;
0154
0155       wz = p.second.z();
0156       wz = wz <= precision ? precision : wz;
0157
0158       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + err(0, 0));
0159       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + err(1, 1));
0160       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + err(2, 2));
0161       chi2 += dist;
0162     }
0163     TransientVertex v(GlobalPoint(x, y, z), err, iclus, chi2, (int)ndof_x);
0164     return v;
0165   }
0166
0167   inline TransientVertex weightedMeanOutlierRejectionBeamSpot(
0168       const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0169       std::vector<reco::TransientTrack> iclus,
0170       const reco::BeamSpot& beamSpot) {
0171     float x = 0., y = 0., z = 0.;
0172     float s_wx = 0., s_wz = 0.;
0173     float s2_wx = 0., s2_wz = 0.;
0174     float wx = 0., wz = 0., chi2 = 0.;
0175     float wy = 0., s_wy = 0., s2_wy = 0.;
0176     float ndof_x = 0.;
0177
0178     AlgebraicSymMatrix33 err;
0179     err(0, 0) = startError / 10 * startError / 10;
0180     err(1, 1) = startError / 10 * startError / 10;
0181     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0182
0183     GlobalError bse(beamSpot.rotatedCovariance3D());
0184     GlobalPoint bsp(Basic3DVector<float>(beamSpot.position()));
0185
0186     for (const auto& p : points) {
0187       wx = p.second.x() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.x(), 2);
0188       wy = p.second.y() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.y(), 2);
0189
0190       wz = p.second.z() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(p.second.z(), 2);
0191
0192       x += p.first.x() * wx;
0193       y += p.first.y() * wy;
0194       z += p.first.z() * wz;
0195
0196       s_wx += wx;
0197       s_wy += wy;
0198       s_wz += wz;
0199     }
0200
0201     if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0202       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0203       return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0204     }
0205     // use the square of covariance element to increase it's weight: it will be the most important
0206     wx = bse.cxx() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(bse.cxx(), 2);
0207     wy = bse.cyy() <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(bse.cyy(), 2);
0208
0209     x += bsp.x() * wx;
0210     y += bsp.y() * wy;
0211
0212     x /= (s_wx + wx);
0213     y /= (s_wy + wy);
0214     z /= s_wz;
0215
0216     float old_x, old_y, old_z;
0217
0218     float xpull;
0219     int niter = 0;
0220
0221     float err_x, err_y, err_z;
0222
0223     err_x = 1. / s_wx;
0224     err_y = 1. / s_wy;
0225     err_z = 1. / s_wz;
0226
0227     while ((niter++) < 2) {
0228       old_x = x;
0229       old_y = y;
0230       old_z = z;
0231       s_wx = 0;
0232       s_wz = 0;
0233       s2_wx = 0;
0234       s2_wz = 0;
0235
0236       s_wy = 0;
0237       s2_wy = 0;
0238
0239       x = 0;
0240       y = 0;
0241       z = 0;
0242       ndof_x = 0;
0243
0244       for (unsigned int i = 0; i < (unsigned int)points.size(); i++) {
0245         std::pair<GlobalPoint, double> p = nearestPoint(GlobalPoint(old_x, old_y, old_z), (iclus)[i].track());
0246
0247         wx = points[i].second.x() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.x(), 2);
0248         wy = points[i].second.y() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.y(), 2);
0249
0250         wz = points[i].second.z() <= precision ? std::pow(precision, 2) : std::pow(points[i].second.z(), 2);
0251
0252         xpull = 0.;
0253         if (std::pow(p.first.x() - old_x, 2) / (wx + err_x) < muSquare &&
0254             std::pow(p.first.y() - old_y, 2) / (wy + err_y) < muSquare &&
0255             std::pow(p.first.z() - old_z, 2) / (wz + err_z) < muSquare)
0256           xpull = 1.;
0257
0258         ndof_x += xpull;
0259
0260         wx = xpull / wx;
0261         wy = xpull / wy;
0262         wz = xpull / wz;
0263
0264         x += wx * p.first.x();
0265         y += wy * p.first.y();
0266         z += wz * p.first.z();
0267
0268         s_wx += wx;
0269         s_wy += wy;
0270         s_wz += wz;
0271
0272         s2_wx += wx * xpull;
0273         s2_wy += wy * xpull;
0274         s2_wz += wz * xpull;
0275       }
0276
0277       if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0278         edm::LogWarning("WeightedMeanFitter")
0279             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0280         return TransientVertex(GlobalPoint(0, 0, 0), err, iclus, 0, 0);
0281       }
0282       wx = bse.cxx() <= std::pow(precision, 2) ? 1. / std::pow(precision, 2) : 1. / bse.cxx();
0283       wy = bse.cyy() <= std::pow(precision, 2) ? 1. / std::pow(precision, 2) : 1. / bse.cyy();
0284
0285       x += bsp.x() * wx;
0286       y += bsp.y() * wy;
0287       s_wx += wx;
0288       s2_wx += wx;
0289       s_wy += wy;
0290       s2_wy += wy;
0291
0292       x /= s_wx;
0293       y /= s_wy;
0294       z /= s_wz;
0295
0296       err_x = (s2_wx / std::pow(s_wx, 2));
0297       err_y = (s2_wy / std::pow(s_wy, 2));
0298       err_z = (s2_wz / std::pow(s_wz, 2));
0299
0300       if (std::abs(x - old_x) < (precision) && std::abs(y - old_y) < (precision) && std::abs(z - old_z) < (precision)) {
0301         break;
0302       }
0303     }
0304     err(0, 0) = err_x * corr_x_bs * corr_x_bs;
0305     err(1, 1) = err_y * corr_x_bs * corr_x_bs;
0306     err(2, 2) = err_z * corr_z * corr_z;
0307
0308     float dist = 0;
0309     for (const auto& p : points) {
0310       wx = p.second.x();
0311       wx = wx <= precision ? precision : wx;
0312
0313       wz = p.second.z();
0314       wz = wz <= precision ? precision : wz;
0315
0316       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + err(0, 0));
0317       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + err(1, 1));
0318       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + err(2, 2));
0319       chi2 += dist;
0320     }
0321     TransientVertex v(GlobalPoint(x, y, z), err, iclus, chi2, (int)ndof_x);
0322     return v;
0323   }
0324
0325   inline TransientVertex weightedMeanOutlierRejectionVarianceAsError(
0326       const std::vector<std::pair<GlobalPoint, GlobalPoint>>& points,
0327       std::vector<std::vector<reco::TransientTrack>>::const_iterator iclus) {
0328     float x = 0, y = 0, z = 0, s_wx = 0, s_wy = 0, s_wz = 0, wx = 0, wy = 0, wz = 0, chi2 = 0;
0329     float ndof_x = 0;
0330     AlgebraicSymMatrix33 err;
0331     err(2, 2) = startError * startError;  // error is 20 cm, so cov -> is 20 ^ 2
0332     err(0, 0) = err(1, 1) = err(2, 2) / 100.;
0333
0334     for (const auto& p : points) {
0335       wx = p.second.x();
0336       wx = wx <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(wx, 2);
0337
0338       wz = p.second.z();
0339       wz = wz <= precision ? 1. / std::pow(precision, 2) : 1. / std::pow(wz, 2);
0340
0341       x += p.first.x() * wx;
0342       y += p.first.y() * wx;
0343       z += p.first.z() * wz;
0344
0345       s_wx += wx;
0346       s_wz += wz;
0347     }
0348
0349     if (s_wx == 0. || s_wz == 0.) {
0350       edm::LogWarning("WeightedMeanFitter") << "Vertex fitting failed at beginning\n";
0351       return TransientVertex(GlobalPoint(0, 0, 0), err, *iclus, 0, 0);
0352     }
0353
0354     x /= s_wx;
0355     y /= s_wx;
0356     z /= s_wz;
0357
0358     float old_x, old_y, old_z;
0359     float xpull;
0360     int niter = 0;
0361     float err_x, err_y, err_z;
0362     err_x = 1. / s_wx;
0363     err_y = 1. / s_wx;
0364     err_z = 1. / s_wz;
0365     float s_err_x = 0., s_err_y = 0., s_err_z = 0.;
0366     while ((niter++) < maxIterations) {
0367       old_x = x;
0368       old_y = y;
0369       old_z = z;
0370       s_wx = 0;
0371       s_wy = 0;
0372       s_wz = 0;
0373       x = 0;
0374       y = 0;
0375       z = 0;
0376       s_err_x = 0.;
0377       s_err_y = 0.;
0378       s_err_z = 0.;
0379
0380       for (const auto& p : points) {
0381         wx = p.second.x();
0382         wx = wx <= precision ? precision : wx;
0383
0384         wy = wx * wx + err_y;
0385         wx = wx * wx + err_x;
0386
0387         wz = p.second.z();
0388         wz = wz <= precision ? precision : wz;
0389         wz = wz * wz + err_z;
0390
0391         xpull = std::pow((p.first.x() - old_x), 2) / wx;
0392         xpull += std::pow((p.first.y() - old_y), 2) / wy;
0393         xpull += std::pow((p.first.z() - old_z), 2) / wz;
0394         xpull = 1. / (1. + std::exp(-0.5 * ((muSquare)-xpull)));
0395         ndof_x += xpull;
0396
0397         wx = 1. / wx;
0398         wy = 1. / wy;
0399         wz = 1. / wz;
0400
0401         wx *= xpull;
0402         wy *= xpull;
0403         wz *= xpull;
0404
0405         x += wx * p.first.x();
0406         y += wy * p.first.y();
0407         z += wz * p.first.z();
0408
0409         s_wx += wx;
0410         s_wy += wy;
0411         s_wz += wz;
0412
0413         s_err_x += wx * pow(p.first.x() - old_x, 2);
0414         s_err_y += wy * pow(p.first.y() - old_y, 2);
0415         s_err_z += wz * pow(p.first.z() - old_z, 2);
0416       }
0417       if (s_wx == 0. || s_wy == 0. || s_wz == 0.) {
0418         edm::LogWarning("WeightedMeanFitter")
0419             << "Vertex fitting failed, either all tracks are outliers or they have a very large error\n";
0420         return TransientVertex(GlobalPoint(0, 0, 0), err, *iclus, 0, 0);
0421       }
0422       x /= s_wx;
0423       y /= s_wy;
0424       z /= s_wz;
0425
0426       err_x = s_err_x / s_wx;
0427       err_y = s_err_y / s_wy;
0428       err_z = s_err_z / s_wz;
0429
0430       if (std::abs(x - old_x) < (precision / 1.) && std::abs(y - old_y) < (precision / 1.) &&
0431           std::abs(z - old_z) < (precision / 1.))
0432         break;
0433     }
0434
0435     err(0, 0) = err_x;
0436     err(1, 1) = err_y;
0437     err(2, 2) = err_z;
0438
0439     float dist = 0.f;
0440     for (const auto& p : points) {
0441       wx = p.second.x();
0442       wx = wx <= precision ? precision : wx;
0443
0444       wz = p.second.z();
0445       wz = wz <= precision ? precision : wz;
0446
0447       dist = std::pow(p.first.x() - x, 2) / (std::pow(wx, 2) + std::pow(err(0, 0), 2));
0448       dist += std::pow(p.first.y() - y, 2) / (std::pow(wx, 2) + std::pow(err(1, 1), 2));
0449       dist += std::pow(p.first.z() - z, 2) / (std::pow(wz, 2) + std::pow(err(2, 2), 2));
0450       chi2 += dist;
0451     }
0452     TransientVertex v(GlobalPoint(x, y, z), err, *iclus, chi2, (int)ndof_x);
0453     return v;
0454   }
0455
0456 };  // namespace WeightedMeanFitter
0457
0458 #endif