Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2024-04-06 12:22:10

0001 /*
0002 Track Quality Body file
0003 C.Brown & C.Savard 07/2020
0004 */
0005 
0006 #include "L1Trigger/TrackTrigger/interface/L1TrackQuality.h"
0007 
0008 //Constructors
0009 
0010 L1TrackQuality::L1TrackQuality() {}
0011 
0012 L1TrackQuality::L1TrackQuality(const edm::ParameterSet& qualityParams) : useHPH_(false), bonusFeatures_() {
0013   std::string AlgorithmString = qualityParams.getParameter<std::string>("qualityAlgorithm");
0014   // Unpacks EDM parameter set itself to save unecessary processing within TrackProducers
0015   if (AlgorithmString == "Cut") {
0016     setCutParameters(AlgorithmString,
0017                      (float)qualityParams.getParameter<double>("maxZ0"),
0018                      (float)qualityParams.getParameter<double>("maxEta"),
0019                      (float)qualityParams.getParameter<double>("chi2dofMax"),
0020                      (float)qualityParams.getParameter<double>("bendchi2Max"),
0021                      (float)qualityParams.getParameter<double>("minPt"),
0022                      qualityParams.getParameter<int>("nStubsmin"));
0023   }
0024 
0025   else {
0026     setONNXModel(AlgorithmString,
0027                  qualityParams.getParameter<edm::FileInPath>("ONNXmodel"),
0028                  qualityParams.getParameter<std::string>("ONNXInputName"),
0029                  qualityParams.getParameter<std::vector<std::string>>("featureNames"));
0030     if ((AlgorithmString == "GBDT") || (AlgorithmString == "NN"))
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   // get the nstub
0071   std::vector<edm::Ref<edmNew::DetSetVector<TTStub<Ref_Phase2TrackerDigi_>>, TTStub<Ref_Phase2TrackerDigi_>>> stubRefs =
0072       aTrack.getStubRefs();
0073   int tmp_trk_nstub = stubRefs.size();
0074 
0075   // get other variables directly from TTTrack
0076   float tmp_trk_z0 = aTrack.z0();
0077   float tmp_trk_z0_scaled = tmp_trk_z0 / abs(aTrack.minZ0);
0078   float tmp_trk_phi = aTrack.phi();
0079   float tmp_trk_eta = aTrack.eta();
0080   float tmp_trk_tanl = aTrack.tanL();
0081 
0082   // -------- fill the feature map ---------
0083 
0084   feature_map["nstub"] = float(tmp_trk_nstub);
0085   feature_map["z0"] = tmp_trk_z0;
0086   feature_map["z0_scaled"] = tmp_trk_z0_scaled;
0087   feature_map["phi"] = tmp_trk_phi;
0088   feature_map["eta"] = tmp_trk_eta;
0089   feature_map["nlaymiss_interior"] = float(tmp_trk_nlaymiss_interior);
0090   feature_map["bendchi2_bin"] = tmp_trk_bendchi2_bin;
0091   feature_map["chi2rphi_bin"] = tmp_trk_chi2rphi_bin;
0092   feature_map["chi2rz_bin"] = tmp_trk_chi2rz_bin;
0093   feature_map["tanl"] = tmp_trk_tanl;
0094 
0095   // fill tensor with track params
0096   transformedFeatures.reserve(featureNames.size());
0097   for (const std::string& feature : featureNames)
0098     transformedFeatures.push_back(feature_map[feature]);
0099 
0100   return transformedFeatures;
0101 }
0102 
0103 void L1TrackQuality::setL1TrackQuality(TTTrack<Ref_Phase2TrackerDigi_>& aTrack) {
0104   if (this->qualityAlgorithm_ == QualityAlgorithm::Cut) {
0105     // Get Track parameters
0106     float trk_pt = aTrack.momentum().perp();
0107     float trk_bend_chi2 = aTrack.stubPtConsistency();
0108     float trk_z0 = aTrack.z0();
0109     float trk_eta = aTrack.momentum().eta();
0110     float trk_chi2 = aTrack.chi2();
0111     const auto& stubRefs = aTrack.getStubRefs();
0112     int nStubs = stubRefs.size();
0113 
0114     float classification = 0.0;  // Default classification is 0
0115 
0116     if (trk_pt >= this->minPt_ && abs(trk_z0) < this->maxZ0_ && abs(trk_eta) < this->maxEta_ &&
0117         trk_chi2 < this->chi2dofMax_ && trk_bend_chi2 < this->bendchi2Max_ && nStubs >= this->nStubsmin_)
0118       classification = 1.0;
0119     // Classification updated to 1 if conditions are met
0120 
0121     aTrack.settrkMVA1(classification);
0122   }
0123 
0124   else if (this->qualityAlgorithm_ == QualityAlgorithm::GBDT_cpp) {
0125     // load in bdt
0126     conifer::BDT<float, float> bdt(this->ONNXmodel_.fullPath());
0127 
0128     // collect features and classify using bdt
0129     std::vector<float> inputs = featureTransform(aTrack, this->featureNames_);
0130     std::vector<float> output = bdt.decision_function(inputs);
0131     aTrack.settrkMVA1(1. / (1. + exp(-output.at(0))));  // need logistic sigmoid fcn applied to xgb output
0132   }
0133 
0134   else if ((this->qualityAlgorithm_ == QualityAlgorithm::NN) || (this->qualityAlgorithm_ == QualityAlgorithm::GBDT)) {
0135     // Setup ONNX input and output names and arrays
0136     std::vector<std::string> ortinput_names;
0137     std::vector<std::string> ortoutput_names;
0138 
0139     cms::Ort::FloatArrays ortinput;
0140     cms::Ort::FloatArrays ortoutputs;
0141 
0142     std::vector<float> Transformed_features = featureTransform(aTrack, this->featureNames_);
0143     //    cms::Ort::ONNXRuntime runTime(this->ONNXmodel_.fullPath());  //Setup ONNX runtime
0144 
0145     ortinput_names.push_back(this->ONNXInputName_);
0146     ortoutput_names = runTime_->getOutputNames();
0147 
0148     //ONNX runtime recieves a vector of vectors of floats so push back the input
0149     // vector of float to create a 1,1,21 ortinput
0150     ortinput.push_back(Transformed_features);
0151 
0152     // batch_size 1 as only one set of transformed features is being processed
0153     int batch_size = 1;
0154     // Run classification
0155     ortoutputs = runTime_->run(ortinput_names, ortinput, {}, ortoutput_names, batch_size);
0156 
0157     if (this->qualityAlgorithm_ == QualityAlgorithm::NN) {
0158       aTrack.settrkMVA1(ortoutputs[0][0]);
0159     }
0160 
0161     else if (this->qualityAlgorithm_ == QualityAlgorithm::GBDT) {
0162       aTrack.settrkMVA1(ortoutputs[1][1]);
0163     }
0164     // Slight differences in the ONNX models of the GBDTs and NNs mean different
0165     // indices of the ortoutput need to be accessed
0166   }
0167 
0168   else {
0169     aTrack.settrkMVA1(-999);
0170   }
0171 }
0172 
0173 float L1TrackQuality::runEmulatedTQ(std::vector<ap_fixed<10, 5>> inputFeatures) {
0174   // load in bdt
0175 
0176   conifer::BDT<ap_fixed<10, 5>, ap_fixed<10, 5>> bdt(this->ONNXmodel_.fullPath());
0177 
0178   // collect features and classify using bdt
0179   std::vector<ap_fixed<10, 5>> output = bdt.decision_function(inputFeatures);
0180   return output.at(0).to_float();  // need logistic sigmoid fcn applied to xgb output
0181 }
0182 
0183 void L1TrackQuality::setCutParameters(std::string const& AlgorithmString,
0184                                       float maxZ0,
0185                                       float maxEta,
0186                                       float chi2dofMax,
0187                                       float bendchi2Max,
0188                                       float minPt,
0189                                       int nStubmin) {
0190   qualityAlgorithm_ = QualityAlgorithm::Cut;
0191   maxZ0_ = maxZ0;
0192   maxEta_ = maxEta;
0193   chi2dofMax_ = chi2dofMax;
0194   bendchi2Max_ = bendchi2Max;
0195   minPt_ = minPt;
0196   nStubsmin_ = nStubmin;
0197 }
0198 
0199 void L1TrackQuality::setONNXModel(std::string const& AlgorithmString,
0200                                   edm::FileInPath const& ONNXmodel,
0201                                   std::string const& ONNXInputName,
0202                                   std::vector<std::string> const& featureNames) {
0203   //Convert algorithm string to Enum class for track by track comparison
0204   if (AlgorithmString == "NN") {
0205     qualityAlgorithm_ = QualityAlgorithm::NN;
0206   } else if (AlgorithmString == "GBDT") {
0207     qualityAlgorithm_ = QualityAlgorithm::GBDT;
0208   } else if (AlgorithmString == "GBDT_cpp") {
0209     qualityAlgorithm_ = QualityAlgorithm::GBDT_cpp;
0210   } else {
0211     qualityAlgorithm_ = QualityAlgorithm::None;
0212   }
0213   ONNXmodel_ = ONNXmodel;
0214   ONNXInputName_ = ONNXInputName;
0215   featureNames_ = featureNames;
0216 }
0217 
0218 void L1TrackQuality::setBonusFeatures(std::vector<float> bonusFeatures) {
0219   bonusFeatures_ = bonusFeatures;
0220   useHPH_ = true;
0221 }