ONNXRuntime/interface/tensor_fillers.h

0001 #ifndef RecoBTag_ONNXRuntime_tensor_fillers_h
0002 #define RecoBTag_ONNXRuntime_tensor_fillers_h
0003
0004 #include "DataFormats/BTauReco/interface/DeepFlavourTagInfo.h"
0005 #include "DataFormats/BTauReco/interface/ParticleTransformerAK4Features.h"
0006 #include "DataFormats/BTauReco/interface/UnifiedParticleTransformerAK4Features.h"
0007 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0008 #include "RecoBTag/ONNXRuntime/interface/tensor_configs.h"
0009
0010 namespace btagbtvdeep {
0011
0012   void jet_tensor_filler(float*& ptr, const btagbtvdeep::DeepFlavourFeatures& features);
0013
0014   void cpf_tensor_filler(float*& ptr, const btagbtvdeep::ChargedCandidateFeatures& c_pf_features);
0015
0016   void npf_tensor_filler(float*& ptr, const btagbtvdeep::NeutralCandidateFeatures& n_pf_features);
0017
0018   void sv_tensor_filler(float*& ptr, const btagbtvdeep::SecondaryVertexFeatures& sv_features);
0019
0020   void jet4vec_tensor_filler(float*& ptr, const btagbtvdeep::JetFeatures& jet_features);
0021
0022   void seedTrack_tensor_filler(float*& ptr, const btagbtvdeep::SeedingTrackFeatures& seed_features);
0023
0024   void neighbourTrack_tensor_filler(float*& ptr, const btagbtvdeep::TrackPairFeatures& neighbourTrack_features);
0025
0026   std::vector<float> inputs_parT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features,
0027                                  parT::InputFeatures ifeature);
0028
0029   std::vector<float> inputs_parT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features,
0030                                  parT::InputFeatures ifeature);
0031
0032   std::vector<float> inputs_parT(const btagbtvdeep::SecondaryVertexFeatures& sv_features, parT::InputFeatures ifeature);
0033
0034   std::vector<float> inputs_UparT(const btagbtvdeep::ChargedCandidateFeatures& c_pf_features,
0035                                   UparT::InputFeatures ifeature);
0036
0037   std::vector<float> inputs_UparT(const btagbtvdeep::LostTracksFeatures& lt_features, UparT::InputFeatures ifeature);
0038
0039   std::vector<float> inputs_UparT(const btagbtvdeep::NeutralCandidateFeatures& n_pf_features,
0040                                   UparT::InputFeatures ifeature);
0041
0042   std::vector<float> inputs_UparT(const btagbtvdeep::SecondaryVertexFeatures& sv_features,
0043                                   UparT::InputFeatures ifeature);
0044
0045   template <class parT_features>
0046   void parT_tensor_filler(cms::Ort::FloatArrays& data,
0047                           const parT::InputFeatures ifeature,
0048                           const std::vector<parT_features>& features,
0049                           const unsigned int max_n,
0050                           const float*& start,
0051                           unsigned offset) {
0052     float* ptr = nullptr;
0053     for (std::size_t n = 0; n < max_n; n++) {
0054       const auto& f = features.at(n);
0055       ptr = &data[ifeature][offset + n * parT::N_InputFeatures.at(ifeature)];
0056       start = ptr;
0057       const std::vector<float>& inputs = inputs_parT(f, ifeature);
0058       for (unsigned int i = 0; i < inputs.size(); i++) {
0059         *ptr = inputs[i];
0060         ++ptr;
0061       }
0062       if (!inputs.empty())
0063         --ptr;
0064       assert(start + parT::N_InputFeatures.at(ifeature) - 1 == ptr);
0065     }
0066   }
0067
0068   template <class parT_features>
0069   void parT_tensor_filler(std::vector<float>& vdata,
0070                           const parT::InputFeatures ifeature,
0071                           const std::vector<parT_features>& features,
0072                           const unsigned int target_n) {
0073     unsigned int n =
0074         std::clamp((unsigned int)features.size(), (unsigned int)0, (unsigned int)parT::N_AcceptedFeatures.at(ifeature));
0075     for (unsigned int count = 0; count < n; count++) {
0076       const std::vector<float>& inputs = inputs_parT(features.at(count), ifeature);
0077       vdata.insert(vdata.end(), inputs.begin(), inputs.end());
0078     }
0079     unsigned int n_features = parT::N_InputFeatures.at(ifeature);
0080     if (n < target_n)
0081       vdata.insert(vdata.end(), (target_n - n) * n_features, 0);  // Add 0 to unfilled part as padding value
0082   }
0083
0084   template <class UparT_features>
0085   void UparT_tensor_filler(cms::Ort::FloatArrays& data,
0086                            const UparT::InputFeatures ifeature,
0087                            const std::vector<UparT_features>& features,
0088                            const unsigned int max_n,
0089                            const float*& start,
0090                            unsigned offset) {
0091     float* ptr = nullptr;
0092     for (std::size_t n = 0; n < max_n; n++) {
0093       const auto& f = features.at(n);
0094       ptr = &data[ifeature][offset + n * UparT::N_InputFeatures.at(ifeature)];
0095       start = ptr;
0096       const std::vector<float>& inputs = inputs_UparT(f, ifeature);
0097       for (unsigned int i = 0; i < inputs.size(); i++) {
0098         *ptr = inputs[i];
0099         ++ptr;
0100       }
0101       if (!inputs.empty())
0102         --ptr;
0103       assert(start + UparT::N_InputFeatures.at(ifeature) - 1 == ptr);
0104     }
0105   }
0106
0107   template <class UparT_features>
0108   void UparT_tensor_filler(std::vector<float>& vdata,
0109                            const UparT::InputFeatures ifeature,
0110                            const std::vector<UparT_features>& features,
0111                            const unsigned int target_n) {
0112     unsigned int n = std::clamp(
0113         (unsigned int)features.size(), (unsigned int)0, (unsigned int)UparT::N_AcceptedFeatures.at(ifeature));
0114     for (unsigned int count = 0; count < n; count++) {
0115       const std::vector<float>& inputs = inputs_UparT(features.at(count), ifeature);
0116       vdata.insert(vdata.end(), inputs.begin(), inputs.end());
0117     }
0118     unsigned int n_features = UparT::N_InputFeatures.at(ifeature);
0119     if (n < target_n)
0120       vdata.insert(vdata.end(), (target_n - n) * n_features, 0);  // Add 0 to unfilled part as padding value
0121   }
0122 }  // namespace btagbtvdeep
0123
0124 #endif