TrackTrigger/src/L1TrackQuality.cc

0001 /*
0002 Track Quality Body file
0003
0004 C.Brown & C.Savard 07/2020
0005 */
0006
0007 #include "L1Trigger/TrackTrigger/interface/L1TrackQuality.h"
0008 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0009
0010 //Constructors
0011
0012 L1TrackQuality::L1TrackQuality() {}
0013
0014 L1TrackQuality::L1TrackQuality(const edm::ParameterSet& qualityParams) {
0015   std::string AlgorithmString = qualityParams.getParameter<std::string>("qualityAlgorithm");
0016   // Unpacks EDM parameter set itself to save unecessary processing within TrackProducers
0017   if (AlgorithmString == "Cut") {
0018     setCutParameters(AlgorithmString,
0019                      (float)qualityParams.getParameter<double>("maxZ0"),
0020                      (float)qualityParams.getParameter<double>("maxEta"),
0021                      (float)qualityParams.getParameter<double>("chi2dofMax"),
0022                      (float)qualityParams.getParameter<double>("bendchi2Max"),
0023                      (float)qualityParams.getParameter<double>("minPt"),
0024                      qualityParams.getParameter<int>("nStubsmin"));
0025   }
0026
0027   else {
0028     setONNXModel(AlgorithmString,
0029                  qualityParams.getParameter<edm::FileInPath>("ONNXmodel"),
0030                  qualityParams.getParameter<std::string>("ONNXInputName"),
0031                  qualityParams.getParameter<std::vector<std::string>>("featureNames"));
0032     ONNXInvRScaling_ = qualityParams.getParameter<double>("ONNXInvRScale");
0033   }
0034 }
0035
0036 std::vector<float> L1TrackQuality::featureTransform(TTTrack<Ref_Phase2TrackerDigi_>& aTrack,
0037                                                     std::vector<std::string> const& featureNames) {
0038   // List input features for MVA in proper order below, the features options are
0039   // {"log_chi2","log_chi2rphi","log_chi2rz","log_bendchi2","nstubs","lay1_hits","lay2_hits",
0040   // "lay3_hits","lay4_hits","lay5_hits","lay6_hits","disk1_hits","disk2_hits","disk3_hits",
0041   // "disk4_hits","disk5_hits","rinv","tanl","z0","dtot","ltot","chi2","chi2rz","chi2rphi",
0042   // "bendchi2","pt","eta","nlaymiss_interior","phi","bendchi2_bin","chi2rz_bin","chi2rphi_bin"}
0043
0044   std::vector<float> transformedFeatures;
0045
0046   // Define feature map, filled as features are generated
0047   std::map<std::string, float> feature_map;
0048
0049   // The following converts the 7 bit hitmask in the TTTrackword to an expected
0050   // 11 bit hitmask based on the eta of the track
0051   std::vector<int> hitpattern_binary = {0, 0, 0, 0, 0, 0, 0};
0052   std::vector<int> hitpattern_expanded_binary = {0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0};
0053   std::vector<float> eta_bins = {0.0, 0.2, 0.41, 0.62, 0.9, 1.26, 1.68, 2.08, 2.4};
0054
0055   // Expected hitmap table, each row corresponds to an eta bin, each value corresponds to
0056   // the expected layer in the expanded hit pattern. The expanded hit pattern should be
0057   // 11 bits but contains a 12th element so this hitmap table is symmetric
0058   int hitmap[8][7] = {{0, 1, 2, 3, 4, 5, 11},
0059                       {0, 1, 2, 3, 4, 5, 11},
0060                       {0, 1, 2, 3, 4, 5, 11},
0061                       {0, 1, 2, 3, 4, 5, 11},
0062                       {0, 1, 2, 3, 4, 5, 11},
0063                       {0, 1, 2, 6, 7, 8, 9},
0064                       {0, 1, 7, 8, 9, 10, 11},
0065                       {0, 6, 7, 8, 9, 10, 11}};
0066
0067   // iterate through bits of the hitpattern and compare to 1 filling the hitpattern binary vector
0068   int tmp_trk_hitpattern = aTrack.hitPattern();
0069   for (int i = 6; i >= 0; i--) {
0070     int k = tmp_trk_hitpattern >> i;
0071     if (k & 1)
0072       hitpattern_binary[i] = 1;
0073   }
0074
0075   // calculate number of missed interior layers from hitpattern
0076   int nbits = floor(log2(tmp_trk_hitpattern)) + 1;
0077   int lay_i = 0;
0078   int tmp_trk_nlaymiss_interior = 0;
0079   bool seq = false;
0080   for (int i = 0; i < nbits; i++) {
0081     lay_i = ((1 << i) & tmp_trk_hitpattern) >> i;  //0 or 1 in ith bit (right to left)
0082
0083     if (lay_i && !seq)
0084       seq = true;  //sequence starts when first 1 found
0085     if (!lay_i && seq)
0086       tmp_trk_nlaymiss_interior++;
0087   }
0088
0089   float eta = abs(aTrack.eta());
0090   int eta_size = static_cast<int>(eta_bins.size());
0091   // First iterate through eta bins
0092
0093   for (int j = 1; j < eta_size; j++) {
0094     if (eta < eta_bins[j] && eta >= eta_bins[j - 1])  // if track in eta bin
0095     {
0096       // Iterate through hitpattern binary
0097       for (int k = 0; k <= 6; k++)
0098         // Fill expanded binary entries using the expected hitmap table positions
0099         hitpattern_expanded_binary[hitmap[j - 1][k]] = hitpattern_binary[k];
0100       break;
0101     }
0102   }
0103
0104   int tmp_trk_ltot = 0;
0105   //calculate number of layer hits
0106   for (int i = 0; i < 6; ++i) {
0107     tmp_trk_ltot += hitpattern_expanded_binary[i];
0108   }
0109
0110   int tmp_trk_dtot = 0;
0111   //calculate number of disk hits
0112   for (int i = 6; i < 11; ++i) {
0113     tmp_trk_dtot += hitpattern_expanded_binary[i];
0114   }
0115
0116   // bin bendchi2 variable (bins from https://twiki.cern.ch/twiki/bin/viewauth/CMS/HybridDataFormat#Fitted_Tracks_written_by_KalmanF)
0117   float tmp_trk_bendchi2 = aTrack.stubPtConsistency();
0118   std::array<float, 8> bendchi2_bins{{0, 0.5, 1.25, 2, 3, 5, 10, 50}};
0119   int n_bendchi2 = static_cast<int>(bendchi2_bins.size());
0120   float tmp_trk_bendchi2_bin = -1;
0121   for (int i = 0; i < (n_bendchi2 - 1); i++) {
0122     if (tmp_trk_bendchi2 >= bendchi2_bins[i] && tmp_trk_bendchi2 < bendchi2_bins[i + 1]) {
0123       tmp_trk_bendchi2_bin = i;
0124       break;
0125     }
0126   }
0127   if (tmp_trk_bendchi2_bin < 0)
0128     tmp_trk_bendchi2_bin = n_bendchi2;
0129
0130   // bin chi2rphi variable (bins from https://twiki.cern.ch/twiki/bin/viewauth/CMS/HybridDataFormat#Fitted_Tracks_written_by_KalmanF)
0131   float tmp_trk_chi2rphi = aTrack.chi2XY();
0132   std::array<float, 16> chi2rphi_bins{{0, 0.25, 0.5, 1, 2, 3, 5, 7, 10, 20, 40, 100, 200, 500, 1000, 3000}};
0133   int n_chi2rphi = static_cast<int>(chi2rphi_bins.size());
0134   float tmp_trk_chi2rphi_bin = -1;
0135   for (int i = 0; i < (n_chi2rphi - 1); i++) {
0136     if (tmp_trk_chi2rphi >= chi2rphi_bins[i] && tmp_trk_chi2rphi < chi2rphi_bins[i + 1]) {
0137       tmp_trk_chi2rphi_bin = i;
0138       break;
0139     }
0140   }
0141   if (tmp_trk_chi2rphi_bin < 0)
0142     tmp_trk_chi2rphi_bin = n_chi2rphi;
0143
0144   // bin chi2rz variable (bins from https://twiki.cern.ch/twiki/bin/viewauth/CMS/HybridDataFormat#Fitted_Tracks_written_by_KalmanF)
0145   float tmp_trk_chi2rz = aTrack.chi2Z();
0146   std::array<float, 16> chi2rz_bins{{0, 0.25, 0.5, 1, 2, 3, 5, 7, 10, 20, 40, 100, 200, 500, 1000, 3000}};
0147   int n_chi2rz = static_cast<int>(chi2rz_bins.size());
0148   float tmp_trk_chi2rz_bin = -1;
0149   for (int i = 0; i < (n_chi2rz - 1); i++) {
0150     if (tmp_trk_chi2rz >= chi2rz_bins[i] && tmp_trk_chi2rz < chi2rz_bins[i + 1]) {
0151       tmp_trk_chi2rz_bin = i;
0152       break;
0153     }
0154   }
0155   if (tmp_trk_chi2rz_bin < 0)
0156     tmp_trk_chi2rz_bin = n_chi2rz;
0157
0158   // get the nstub
0159   std::vector<edm::Ref<edmNew::DetSetVector<TTStub<Ref_Phase2TrackerDigi_>>, TTStub<Ref_Phase2TrackerDigi_>>> stubRefs =
0160       aTrack.getStubRefs();
0161
0162   // fill the feature map
0163   feature_map["nstub"] = stubRefs.size();
0164   feature_map["rinv"] = ONNXInvRScaling_ * abs(aTrack.rInv());
0165   feature_map["tanl"] = abs(aTrack.tanL());
0166   feature_map["z0"] = aTrack.z0();
0167   feature_map["phi"] = aTrack.phi();
0168   feature_map["pt"] = aTrack.momentum().perp();
0169   feature_map["eta"] = aTrack.eta();
0170
0171   float tmp_trk_chi2 = aTrack.chi2();
0172   feature_map["chi2"] = tmp_trk_chi2;
0173   feature_map["log_chi2"] = log(tmp_trk_chi2);
0174
0175   feature_map["chi2rphi"] = tmp_trk_chi2rphi;
0176   feature_map["log_chi2rphi"] = log(tmp_trk_chi2rphi);
0177
0178   feature_map["chi2rz"] = tmp_trk_chi2rz;
0179   feature_map["log_chi2rz"] = log(tmp_trk_chi2rz);
0180
0181   feature_map["chi2rz"] = tmp_trk_chi2rz;
0182   feature_map["log_chi2rz"] = log(tmp_trk_chi2rz);
0183
0184   feature_map["bendchi2"] = tmp_trk_bendchi2;
0185   feature_map["log_bendchi2"] = log(tmp_trk_bendchi2);
0186
0187   feature_map["lay1_hits"] = float(hitpattern_expanded_binary[0]);
0188   feature_map["lay2_hits"] = float(hitpattern_expanded_binary[1]);
0189   feature_map["lay3_hits"] = float(hitpattern_expanded_binary[2]);
0190   feature_map["lay4_hits"] = float(hitpattern_expanded_binary[3]);
0191   feature_map["lay5_hits"] = float(hitpattern_expanded_binary[4]);
0192   feature_map["lay6_hits"] = float(hitpattern_expanded_binary[5]);
0193   feature_map["disk1_hits"] = float(hitpattern_expanded_binary[6]);
0194   feature_map["disk2_hits"] = float(hitpattern_expanded_binary[7]);
0195   feature_map["disk3_hits"] = float(hitpattern_expanded_binary[8]);
0196   feature_map["disk4_hits"] = float(hitpattern_expanded_binary[9]);
0197   feature_map["disk5_hits"] = float(hitpattern_expanded_binary[10]);
0198
0199   feature_map["dtot"] = float(tmp_trk_dtot);
0200   feature_map["ltot"] = float(tmp_trk_ltot);
0201
0202   feature_map["nlaymiss_interior"] = float(tmp_trk_nlaymiss_interior);
0203   feature_map["bendchi2_bin"] = tmp_trk_bendchi2_bin;
0204   feature_map["chi2rphi_bin"] = tmp_trk_chi2rphi_bin;
0205   feature_map["chi2rz_bin"] = tmp_trk_chi2rz_bin;
0206
0207   // fill tensor with track params
0208   transformedFeatures.reserve(featureNames.size());
0209   for (const std::string& feature : featureNames)
0210     transformedFeatures.push_back(feature_map[feature]);
0211
0212   return transformedFeatures;
0213 }
0214
0215 void L1TrackQuality::setL1TrackQuality(TTTrack<Ref_Phase2TrackerDigi_>& aTrack) {
0216   if (this->qualityAlgorithm_ == QualityAlgorithm::Cut) {
0217     // Get Track parameters
0218     float trk_pt = aTrack.momentum().perp();
0219     float trk_bend_chi2 = aTrack.stubPtConsistency();
0220     float trk_z0 = aTrack.z0();
0221     float trk_eta = aTrack.momentum().eta();
0222     float trk_chi2 = aTrack.chi2();
0223     const auto& stubRefs = aTrack.getStubRefs();
0224     int nStubs = stubRefs.size();
0225
0226     float classification = 0.0;  // Default classification is 0
0227
0228     if (trk_pt >= this->minPt_ && abs(trk_z0) < this->maxZ0_ && abs(trk_eta) < this->maxEta_ &&
0229         trk_chi2 < this->chi2dofMax_ && trk_bend_chi2 < this->bendchi2Max_ && nStubs >= this->nStubsmin_)
0230       classification = 1.0;
0231     // Classification updated to 1 if conditions are met
0232
0233     aTrack.settrkMVA1(classification);
0234   }
0235
0236   if ((this->qualityAlgorithm_ == QualityAlgorithm::NN) || (this->qualityAlgorithm_ == QualityAlgorithm::GBDT)) {
0237     // Setup ONNX input and output names and arrays
0238     std::vector<std::string> ortinput_names;
0239     std::vector<std::string> ortoutput_names;
0240
0241     cms::Ort::FloatArrays ortinput;
0242     cms::Ort::FloatArrays ortoutputs;
0243
0244     std::vector<float> Transformed_features = featureTransform(aTrack, this->featureNames_);
0245     cms::Ort::ONNXRuntime Runtime(this->ONNXmodel_.fullPath());  //Setup ONNX runtime
0246
0247     ortinput_names.push_back(this->ONNXInputName_);
0248     ortoutput_names = Runtime.getOutputNames();
0249
0250     //ONNX runtime recieves a vector of vectors of floats so push back the input
0251     // vector of float to create a 1,1,21 ortinput
0252     ortinput.push_back(Transformed_features);
0253
0254     // batch_size 1 as only one set of transformed features is being processed
0255     int batch_size = 1;
0256     // Run classification
0257     ortoutputs = Runtime.run(ortinput_names, ortinput, {}, ortoutput_names, batch_size);
0258
0259     if (this->qualityAlgorithm_ == QualityAlgorithm::NN) {
0260       aTrack.settrkMVA1(ortoutputs[0][0]);
0261     }
0262
0263     else if (this->qualityAlgorithm_ == QualityAlgorithm::GBDT) {
0264       aTrack.settrkMVA1(ortoutputs[1][1]);
0265     }
0266     // Slight differences in the ONNX models of the GBDTs and NNs mean different
0267     // indices of the ortoutput need to be accessed
0268   }
0269
0270   else {
0271     aTrack.settrkMVA1(-999);
0272   }
0273 }
0274
0275 void L1TrackQuality::setCutParameters(std::string const& AlgorithmString,
0276                                       float maxZ0,
0277                                       float maxEta,
0278                                       float chi2dofMax,
0279                                       float bendchi2Max,
0280                                       float minPt,
0281                                       int nStubmin) {
0282   qualityAlgorithm_ = QualityAlgorithm::Cut;
0283   maxZ0_ = maxZ0;
0284   maxEta_ = maxEta;
0285   chi2dofMax_ = chi2dofMax;
0286   bendchi2Max_ = bendchi2Max;
0287   minPt_ = minPt;
0288   nStubsmin_ = nStubmin;
0289 }
0290
0291 void L1TrackQuality::setONNXModel(std::string const& AlgorithmString,
0292                                   edm::FileInPath const& ONNXmodel,
0293                                   std::string const& ONNXInputName,
0294                                   std::vector<std::string> const& featureNames) {
0295   //Convert algorithm string to Enum class for track by track comparison
0296   if (AlgorithmString == "NN") {
0297     qualityAlgorithm_ = QualityAlgorithm::NN;
0298   } else if (AlgorithmString == "GBDT") {
0299     qualityAlgorithm_ = QualityAlgorithm::GBDT;
0300   } else {
0301     qualityAlgorithm_ = QualityAlgorithm::None;
0302   }
0303   ONNXmodel_ = ONNXmodel;
0304   ONNXInputName_ = ONNXInputName;
0305   featureNames_ = featureNames;
0306 }