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
0009 //Constructors
0010
0011 L1TrackQuality::L1TrackQuality() {}
0012
0013 L1TrackQuality::L1TrackQuality(const edm::ParameterSet& qualityParams) : setupHPH_(), useHPH_(false) {
0014   std::string AlgorithmString = qualityParams.getParameter<std::string>("qualityAlgorithm");
0015   // Unpacks EDM parameter set itself to save unecessary processing within TrackProducers
0016   if (AlgorithmString == "Cut") {
0017     setCutParameters(AlgorithmString,
0018                      (float)qualityParams.getParameter<double>("maxZ0"),
0019                      (float)qualityParams.getParameter<double>("maxEta"),
0020                      (float)qualityParams.getParameter<double>("chi2dofMax"),
0021                      (float)qualityParams.getParameter<double>("bendchi2Max"),
0022                      (float)qualityParams.getParameter<double>("minPt"),
0023                      qualityParams.getParameter<int>("nStubsmin"));
0024   }
0025
0026   else {
0027     setONNXModel(AlgorithmString,
0028                  qualityParams.getParameter<edm::FileInPath>("ONNXmodel"),
0029                  qualityParams.getParameter<std::string>("ONNXInputName"),
0030                  qualityParams.getParameter<std::vector<std::string>>("featureNames"));
0031     runTime_ = std::make_unique<cms::Ort::ONNXRuntime>(this->ONNXmodel_.fullPath());
0032   }
0033 }
0034
0035 std::vector<float> L1TrackQuality::featureTransform(TTTrack<Ref_Phase2TrackerDigi_>& aTrack,
0036                                                     std::vector<std::string> const& featureNames) {
0037   // List input features for MVA in proper order below, the current features options are
0038   // {"phi", "eta", "z0", "bendchi2_bin", "nstub", "nlaymiss_interior", "chi2rphi_bin",
0039   // "chi2rz_bin"}
0040   //
0041   // To use more features, they must be created here and added to feature_map below
0042
0043   std::vector<float> transformedFeatures;
0044
0045   // Define feature map, filled as features are generated
0046   std::map<std::string, float> feature_map;
0047
0048   // -------- calculate feature variables --------
0049
0050   // calculate number of missed interior layers from hitpattern
0051   int tmp_trk_hitpattern = aTrack.hitPattern();
0052   int nbits = floor(log2(tmp_trk_hitpattern)) + 1;
0053   int lay_i = 0;
0054   int tmp_trk_nlaymiss_interior = 0;
0055   bool seq = false;
0056   for (int i = 0; i < nbits; i++) {
0057     lay_i = ((1 << i) & tmp_trk_hitpattern) >> i;  //0 or 1 in ith bit (right to left)
0058
0059     if (lay_i && !seq)
0060       seq = true;  //sequence starts when first 1 found
0061     if (!lay_i && seq)
0062       tmp_trk_nlaymiss_interior++;
0063   }
0064
0065   // binned chi2 variables
0066   int tmp_trk_bendchi2_bin = aTrack.getBendChi2Bits();
0067   int tmp_trk_chi2rphi_bin = aTrack.getChi2RPhiBits();
0068   int tmp_trk_chi2rz_bin = aTrack.getChi2RZBits();
0069
0070   // Extra variables from HitPatternHelper don't improve performance, so commented out.
0071   //if (useHPH_) {
0072   //  hph::HitPatternHelper hph(setupHPH_, tmp_trk_hitpattern, tmp_trk_tanL, tmp_trk_z0);
0073   //  hitpattern_expanded_binary = hph.binary();
0074   //  tmp_trk_nlaymiss_PS = hph.numMissingPS();
0075   //  tmp_trk_nlaymiss_2S = hph.numMissing2S();
0076   //}
0077
0078   // get the nstub
0079   std::vector<edm::Ref<edmNew::DetSetVector<TTStub<Ref_Phase2TrackerDigi_>>, TTStub<Ref_Phase2TrackerDigi_>>> stubRefs =
0080       aTrack.getStubRefs();
0081   int tmp_trk_nstub = stubRefs.size();
0082
0083   // get other variables directly from TTTrack
0084   float tmp_trk_z0 = aTrack.z0();
0085   float tmp_trk_phi = aTrack.phi();
0086   float tmp_trk_eta = aTrack.eta();
0087
0088   // -------- fill the feature map ---------
0089
0090   feature_map["nstub"] = float(tmp_trk_nstub);
0091   feature_map["z0"] = tmp_trk_z0;
0092   feature_map["phi"] = tmp_trk_phi;
0093   feature_map["eta"] = tmp_trk_eta;
0094   feature_map["nlaymiss_interior"] = float(tmp_trk_nlaymiss_interior);
0095   feature_map["bendchi2_bin"] = tmp_trk_bendchi2_bin;
0096   feature_map["chi2rphi_bin"] = tmp_trk_chi2rphi_bin;
0097   feature_map["chi2rz_bin"] = tmp_trk_chi2rz_bin;
0098
0099   // fill tensor with track params
0100   transformedFeatures.reserve(featureNames.size());
0101   for (const std::string& feature : featureNames)
0102     transformedFeatures.push_back(feature_map[feature]);
0103
0104   return transformedFeatures;
0105 }
0106
0107 void L1TrackQuality::setL1TrackQuality(TTTrack<Ref_Phase2TrackerDigi_>& aTrack) {
0108   if (this->qualityAlgorithm_ == QualityAlgorithm::Cut) {
0109     // Get Track parameters
0110     float trk_pt = aTrack.momentum().perp();
0111     float trk_bend_chi2 = aTrack.stubPtConsistency();
0112     float trk_z0 = aTrack.z0();
0113     float trk_eta = aTrack.momentum().eta();
0114     float trk_chi2 = aTrack.chi2();
0115     const auto& stubRefs = aTrack.getStubRefs();
0116     int nStubs = stubRefs.size();
0117
0118     float classification = 0.0;  // Default classification is 0
0119
0120     if (trk_pt >= this->minPt_ && abs(trk_z0) < this->maxZ0_ && abs(trk_eta) < this->maxEta_ &&
0121         trk_chi2 < this->chi2dofMax_ && trk_bend_chi2 < this->bendchi2Max_ && nStubs >= this->nStubsmin_)
0122       classification = 1.0;
0123     // Classification updated to 1 if conditions are met
0124
0125     aTrack.settrkMVA1(classification);
0126   }
0127
0128   if ((this->qualityAlgorithm_ == QualityAlgorithm::NN) || (this->qualityAlgorithm_ == QualityAlgorithm::GBDT)) {
0129     // Setup ONNX input and output names and arrays
0130     std::vector<std::string> ortinput_names;
0131     std::vector<std::string> ortoutput_names;
0132
0133     cms::Ort::FloatArrays ortinput;
0134     cms::Ort::FloatArrays ortoutputs;
0135
0136     std::vector<float> Transformed_features = featureTransform(aTrack, this->featureNames_);
0137     //    cms::Ort::ONNXRuntime runTime(this->ONNXmodel_.fullPath());  //Setup ONNX runtime
0138
0139     ortinput_names.push_back(this->ONNXInputName_);
0140     ortoutput_names = runTime_->getOutputNames();
0141
0142     //ONNX runtime recieves a vector of vectors of floats so push back the input
0143     // vector of float to create a 1,1,21 ortinput
0144     ortinput.push_back(Transformed_features);
0145
0146     // batch_size 1 as only one set of transformed features is being processed
0147     int batch_size = 1;
0148     // Run classification
0149     ortoutputs = runTime_->run(ortinput_names, ortinput, {}, ortoutput_names, batch_size);
0150
0151     if (this->qualityAlgorithm_ == QualityAlgorithm::NN) {
0152       aTrack.settrkMVA1(ortoutputs[0][0]);
0153     }
0154
0155     else if (this->qualityAlgorithm_ == QualityAlgorithm::GBDT) {
0156       aTrack.settrkMVA1(ortoutputs[1][1]);
0157     }
0158     // Slight differences in the ONNX models of the GBDTs and NNs mean different
0159     // indices of the ortoutput need to be accessed
0160   }
0161
0162   else {
0163     aTrack.settrkMVA1(-999);
0164   }
0165 }
0166
0167 void L1TrackQuality::setCutParameters(std::string const& AlgorithmString,
0168                                       float maxZ0,
0169                                       float maxEta,
0170                                       float chi2dofMax,
0171                                       float bendchi2Max,
0172                                       float minPt,
0173                                       int nStubmin) {
0174   qualityAlgorithm_ = QualityAlgorithm::Cut;
0175   maxZ0_ = maxZ0;
0176   maxEta_ = maxEta;
0177   chi2dofMax_ = chi2dofMax;
0178   bendchi2Max_ = bendchi2Max;
0179   minPt_ = minPt;
0180   nStubsmin_ = nStubmin;
0181 }
0182
0183 void L1TrackQuality::setONNXModel(std::string const& AlgorithmString,
0184                                   edm::FileInPath const& ONNXmodel,
0185                                   std::string const& ONNXInputName,
0186                                   std::vector<std::string> const& featureNames) {
0187   //Convert algorithm string to Enum class for track by track comparison
0188   if (AlgorithmString == "NN") {
0189     qualityAlgorithm_ = QualityAlgorithm::NN;
0190   } else if (AlgorithmString == "GBDT") {
0191     qualityAlgorithm_ = QualityAlgorithm::GBDT;
0192   } else {
0193     qualityAlgorithm_ = QualityAlgorithm::None;
0194   }
0195   ONNXmodel_ = ONNXmodel;
0196   ONNXInputName_ = ONNXInputName;
0197   featureNames_ = featureNames;
0198 }
0199
0200 void L1TrackQuality::beginRun(const hph::Setup* setupHPH) {
0201   // Initialize helper for decoding hit patterns.
0202   setupHPH_ = setupHPH;
0203   useHPH_ = true;
0204 }