Back to home page

Project CMSSW displayed by LXR
 
 

    

File indexing completed on 2024-04-06 12:27:47 

0001 /*
0002  * \class DeepTauIdSONIC
0003  *
0004  * Tau identification using Deep NN with SONIC
0005  *
0006  */
0007 
0008 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0009 #include "RecoTauTag/RecoTau/interface/DeepTauIdBase.h"
0010 
0011 class DeepTauIdSonicProducer : public DeepTauIdBase<TritonEDProducer<>> {
0012 public:
0013   explicit DeepTauIdSonicProducer(edm::ParameterSet const& cfg) : DeepTauIdBase<TritonEDProducer<>>(cfg){};
0014 
0015   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) override;
0016   void produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) override;
0017   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0018 
0019 private:
0020   std::vector<size_t> tau_indices_;
0021 
0022   std::vector<float>* p_tauBlockInputs;
0023   std::vector<float>* p_egammaInnerBlockInputs;
0024   std::vector<float>* p_muonInnerBlockInputs;
0025   std::vector<float>* p_hadronInnerBlockInputs;
0026   std::vector<float>* p_egammaOuterBlockInputs;
0027   std::vector<float>* p_muonOuterBlockInputs;
0028   std::vector<float>* p_hadronOuterBlockInputs;
0029   std::vector<int64_t>* p_innerGridposInputs;
0030   std::vector<int64_t>* p_outerGridposInputs;
0031 
0032   template <typename CandidateCastType, typename TauCastType>
0033   void prepareInputsV2(TauCollection::const_reference& tau,
0034                        const size_t tau_index,
0035                        const edm::RefToBase<reco::BaseTau> tau_ref,
0036                        const std::vector<pat::Electron>* electrons,
0037                        const std::vector<pat::Muon>* muons,
0038                        const edm::View<reco::Candidate>& pfCands,
0039                        const reco::Vertex& pv,
0040                        double rho,
0041                        const TauFunc& tau_funcs);
0042 
0043   template <typename CandidateCastType, typename TauCastType>
0044   void createConvFeatures(const TauCastType& tau,
0045                           const size_t tau_index,
0046                           const edm::RefToBase<reco::BaseTau> tau_ref,
0047                           const reco::Vertex& pv,
0048                           double rho,
0049                           const std::vector<pat::Electron>* electrons,
0050                           const std::vector<pat::Muon>* muons,
0051                           const edm::View<reco::Candidate>& pfCands,
0052                           const CellGrid& grid,
0053                           const TauFunc& tau_funcs,
0054                           bool is_inner,
0055                           std::vector<float>* p_egammaBlockInputs,
0056                           std::vector<float>* p_muonBlockInputs,
0057                           std::vector<float>* p_hadronBlockInputs,
0058                           std::vector<int64_t>* p_GridposInputs);
0059 };
0060 
0061 void DeepTauIdSonicProducer::acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) {
0062   edm::Handle<TauCollection> taus;
0063   iEvent.getByToken(tausToken_, taus);
0064 
0065   loadPrediscriminants(iEvent, taus);
0066 
0067   const std::vector<pat::Electron> electron_collection_default;
0068   const std::vector<pat::Muon> muon_collection_default;
0069   const reco::TauDiscriminatorContainer basicTauDiscriminators_default;
0070   const reco::TauDiscriminatorContainer basicTauDiscriminatorsdR03_default;
0071   const edm::AssociationVector<reco::PFTauRefProd, std::vector<reco::PFTauTransverseImpactParameterRef>>
0072       pfTauTransverseImpactParameters_default;
0073 
0074   const std::vector<pat::Electron>* electron_collection;
0075   const std::vector<pat::Muon>* muon_collection;
0076   const reco::TauDiscriminatorContainer* basicTauDiscriminators;
0077   const reco::TauDiscriminatorContainer* basicTauDiscriminatorsdR03;
0078   const edm::AssociationVector<reco::PFTauRefProd, std::vector<reco::PFTauTransverseImpactParameterRef>>*
0079       pfTauTransverseImpactParameters;
0080 
0081   if (!is_online_) {
0082     electron_collection = &iEvent.get(electrons_token_);
0083     muon_collection = &iEvent.get(muons_token_);
0084     pfTauTransverseImpactParameters = &pfTauTransverseImpactParameters_default;
0085     basicTauDiscriminators = &basicTauDiscriminators_default;
0086     basicTauDiscriminatorsdR03 = &basicTauDiscriminatorsdR03_default;
0087   } else {
0088     electron_collection = &electron_collection_default;
0089     muon_collection = &muon_collection_default;
0090     pfTauTransverseImpactParameters = &iEvent.get(pfTauTransverseImpactParameters_token_);
0091     basicTauDiscriminators = &iEvent.get(basicTauDiscriminators_inputToken_);
0092     basicTauDiscriminatorsdR03 = &iEvent.get(basicTauDiscriminatorsdR03_inputToken_);
0093 
0094     // Get indices for discriminators
0095     if (!discrIndicesMapped_) {
0096       basicDiscrIndexMap_ =
0097           matchDiscriminatorIndices(iEvent, basicTauDiscriminators_inputToken_, requiredBasicDiscriminators_);
0098       basicDiscrdR03IndexMap_ =
0099           matchDiscriminatorIndices(iEvent, basicTauDiscriminatorsdR03_inputToken_, requiredBasicDiscriminatorsdR03_);
0100       discrIndicesMapped_ = true;
0101     }
0102   }
0103 
0104   TauFunc tauIDs = {basicTauDiscriminators,
0105                     basicTauDiscriminatorsdR03,
0106                     pfTauTransverseImpactParameters,
0107                     basicDiscrIndexMap_,
0108                     basicDiscrdR03IndexMap_};
0109 
0110   edm::Handle<edm::View<reco::Candidate>> pfCands;
0111   iEvent.getByToken(pfcandToken_, pfCands);
0112 
0113   edm::Handle<reco::VertexCollection> vertices;
0114   iEvent.getByToken(vtxToken_, vertices);
0115 
0116   edm::Handle<double> rho;
0117   iEvent.getByToken(rho_token_, rho);
0118 
0119   // vector to store the indices for the taus passing the selections
0120   tau_indices_.clear();
0121 
0122   for (size_t tau_index = 0; tau_index < taus->size(); ++tau_index) {
0123     const edm::RefToBase<reco::BaseTau> tauRef = taus->refAt(tau_index);
0124     bool passesPrediscriminants;
0125     if (is_online_) {
0126       passesPrediscriminants = tauIDs.passPrediscriminants<std::vector<TauDiscInfo<reco::PFTauDiscriminator>>>(
0127           recoPrediscriminants_, andPrediscriminants_, tauRef);
0128     } else {
0129       passesPrediscriminants = tauIDs.passPrediscriminants<std::vector<TauDiscInfo<pat::PATTauDiscriminator>>>(
0130           patPrediscriminants_, andPrediscriminants_, tauRef);
0131     }
0132     if (!passesPrediscriminants)
0133       continue;
0134 
0135     // tau index that passes the selection
0136     tau_indices_.push_back(tau_index);
0137   }
0138 
0139   if (tau_indices_.empty()) {
0140     // no tau passing the requirement
0141     // no need to run acquire and inference
0142     client_->setBatchSize(0);
0143     return;
0144   }
0145 
0146   int n_taus = tau_indices_.size();
0147   // for the regular non-split deep-tau model, set the
0148   // batch size to the number of taus per event
0149   client_->setBatchSize(n_taus);
0150 
0151   auto& input_tauBlock = iInput.at("input_tau");
0152   auto& input_innerEgammaBlock = iInput.at("input_inner_egamma");
0153   auto& input_outerEgammaBlock = iInput.at("input_outer_egamma");
0154   auto& input_innerMuonBlock = iInput.at("input_inner_muon");
0155   auto& input_outerMuonBlock = iInput.at("input_outer_muon");
0156   auto& input_innerHadronBlock = iInput.at("input_inner_hadrons");
0157   auto& input_outerHadronBlock = iInput.at("input_outer_hadrons");
0158 
0159   p_innerGridposInputs = nullptr;
0160   p_outerGridposInputs = nullptr;
0161 
0162   auto data_tauBlock = input_tauBlock.allocate<float>();
0163   auto data_innerEgammaBlock = input_innerEgammaBlock.allocate<float>();
0164   auto data_outerEgammaBlock = input_outerEgammaBlock.allocate<float>();
0165   auto data_innerMuonBlock = input_innerMuonBlock.allocate<float>();
0166   auto data_outerMuonBlock = input_outerMuonBlock.allocate<float>();
0167   auto data_innerHadronBlock = input_innerHadronBlock.allocate<float>();
0168   auto data_outerHadronBlock = input_outerHadronBlock.allocate<float>();
0169 
0170   for (unsigned itau_passed = 0; itau_passed < tau_indices_.size(); ++itau_passed) {
0171     // batch size is not one, needs to go to corresponding itau
0172     p_tauBlockInputs = &((*data_tauBlock)[itau_passed]);
0173 
0174     p_egammaInnerBlockInputs = &((*data_innerEgammaBlock)[itau_passed]);
0175     p_muonInnerBlockInputs = &((*data_innerMuonBlock)[itau_passed]);
0176     p_hadronInnerBlockInputs = &((*data_innerHadronBlock)[itau_passed]);
0177 
0178     p_egammaOuterBlockInputs = &((*data_outerEgammaBlock)[itau_passed]);
0179     p_muonOuterBlockInputs = &((*data_outerMuonBlock)[itau_passed]);
0180     p_hadronOuterBlockInputs = &((*data_outerHadronBlock)[itau_passed]);
0181 
0182     int tau_index = tau_indices_[itau_passed];
0183     const edm::RefToBase<reco::BaseTau> tauRef = taus->refAt(tau_index);
0184     prepareInputsV2<pat::PackedCandidate, pat::Tau>(taus->at(tau_index),
0185                                                     tau_index,
0186                                                     tauRef,
0187                                                     electron_collection,
0188                                                     muon_collection,
0189                                                     *pfCands,
0190                                                     vertices->at(0),
0191                                                     *rho,
0192                                                     tauIDs);
0193   }
0194 
0195   // set all input data to the server
0196   input_tauBlock.toServer(data_tauBlock);
0197 
0198   input_innerEgammaBlock.toServer(data_innerEgammaBlock);
0199   input_innerMuonBlock.toServer(data_innerMuonBlock);
0200   input_innerHadronBlock.toServer(data_innerHadronBlock);
0201 
0202   input_outerEgammaBlock.toServer(data_outerEgammaBlock);
0203   input_outerMuonBlock.toServer(data_outerMuonBlock);
0204   input_outerHadronBlock.toServer(data_outerHadronBlock);
0205 }
0206 
0207 void DeepTauIdSonicProducer::produce(edm::Event& iEvent, edm::EventSetup const& iSetup, Output const& iOutput) {
0208   edm::Handle<TauCollection> taus;
0209   iEvent.getByToken(tausToken_, taus);
0210 
0211   if (taus->empty()) {
0212     std::vector<std::vector<float>> pred_all(0, std::vector<float>(deep_tau::NumberOfOutputs, 0.));
0213     createOutputs(iEvent, pred_all, taus);
0214     return;
0215   }
0216 
0217   const auto& output_tauval = iOutput.at("main_output/Softmax");
0218   const auto& outputs_tauval = output_tauval.fromServer<float>();
0219 
0220   // fill the taus passing the selections with the results from produce,
0221   //  and the taus failing the selections with 2 and -1.0
0222   std::vector<std::vector<float>> pred_all(taus->size(), std::vector<float>(deep_tau::NumberOfOutputs, 0.));
0223   for (unsigned itau = 0; itau < taus->size(); ++itau) {
0224     for (int k = 0; k < deep_tau::NumberOfOutputs; ++k) {
0225       pred_all[itau][k] = (k == 2) ? -1.f : 2.f;
0226     }
0227   }
0228   for (unsigned itau_passed = 0; itau_passed < tau_indices_.size(); ++itau_passed) {
0229     int tau_index = tau_indices_[itau_passed];
0230     std::copy(outputs_tauval[itau_passed].begin(), outputs_tauval[itau_passed].end(), pred_all[tau_index].begin());
0231 
0232     if (debug_level >= 2) {
0233       for (int i = 0; i < 4; ++i) {
0234         std::cout << "tau index " << itau_passed << " k " << i << " pred " << pred_all[tau_index][i] << std::endl;
0235       }
0236     }
0237   }
0238   createOutputs(iEvent, pred_all, taus);
0239 }
0240 
0241 template <typename CandidateCastType, typename TauCastType>
0242 void DeepTauIdSonicProducer::prepareInputsV2(TauCollection::const_reference& tau,
0243                                              const size_t tau_index,
0244                                              const edm::RefToBase<reco::BaseTau> tau_ref,
0245                                              const std::vector<pat::Electron>* electrons,
0246                                              const std::vector<pat::Muon>* muons,
0247                                              const edm::View<reco::Candidate>& pfCands,
0248                                              const reco::Vertex& pv,
0249                                              double rho,
0250                                              const TauFunc& tau_funcs) {
0251   using namespace dnn_inputs_v2;
0252   CellGrid inner_grid(number_of_inner_cell, number_of_inner_cell, 0.02, 0.02, disable_CellIndex_workaround_);
0253   CellGrid outer_grid(number_of_outer_cell, number_of_outer_cell, 0.05, 0.05, disable_CellIndex_workaround_);
0254   auto tau_casted = dynamic_cast<const TauCastType&>(tau);
0255   // fill in the inner and outer grids for electrons, muons, and pfCands
0256   fillGrids(tau_casted, *electrons, inner_grid, outer_grid);
0257   fillGrids(tau_casted, *muons, inner_grid, outer_grid);
0258   fillGrids(tau_casted, pfCands, inner_grid, outer_grid);
0259 
0260   p_tauBlockInputs->insert(p_tauBlockInputs->end(), TauBlockInputs::NumberOfInputs, 0.);
0261   std::vector<float>::iterator tauIter = p_tauBlockInputs->begin();
0262 
0263   createTauBlockInputs<CandidateCastType>(tau_casted, tau_index, tau_ref, pv, rho, tau_funcs, tauIter);
0264 
0265   // egamma, muon, and hadron inner and outer inputs for the grids
0266   createConvFeatures<CandidateCastType>(tau_casted,
0267                                         tau_index,
0268                                         tau_ref,
0269                                         pv,
0270                                         rho,
0271                                         electrons,
0272                                         muons,
0273                                         pfCands,
0274                                         inner_grid,
0275                                         tau_funcs,
0276                                         true,
0277                                         p_egammaInnerBlockInputs,
0278                                         p_muonInnerBlockInputs,
0279                                         p_hadronInnerBlockInputs,
0280                                         p_innerGridposInputs);
0281   createConvFeatures<CandidateCastType>(tau_casted,
0282                                         tau_index,
0283                                         tau_ref,
0284                                         pv,
0285                                         rho,
0286                                         electrons,
0287                                         muons,
0288                                         pfCands,
0289                                         outer_grid,
0290                                         tau_funcs,
0291                                         false,
0292                                         p_egammaOuterBlockInputs,
0293                                         p_muonOuterBlockInputs,
0294                                         p_hadronOuterBlockInputs,
0295                                         p_outerGridposInputs);
0296 }
0297 
0298 template <typename CandidateCastType, typename TauCastType>
0299 void DeepTauIdSonicProducer::createConvFeatures(const TauCastType& tau,
0300                                                 const size_t tau_index,
0301                                                 const edm::RefToBase<reco::BaseTau> tau_ref,
0302                                                 const reco::Vertex& pv,
0303                                                 double rho,
0304                                                 const std::vector<pat::Electron>* electrons,
0305                                                 const std::vector<pat::Muon>* muons,
0306                                                 const edm::View<reco::Candidate>& pfCands,
0307                                                 const CellGrid& grid,
0308                                                 const TauFunc& tau_funcs,
0309                                                 bool is_inner,
0310                                                 std::vector<float>* p_egammaBlockInputs,
0311                                                 std::vector<float>* p_muonBlockInputs,
0312                                                 std::vector<float>* p_hadronBlockInputs,
0313                                                 std::vector<int64_t>* p_GridposInputs) {
0314   // fill in the block inputs with zeros
0315   int n_cells = 0;
0316   int n_cell_oneside = is_inner ? dnn_inputs_v2::number_of_inner_cell : dnn_inputs_v2::number_of_outer_cell;
0317 
0318   n_cells = is_inner ? (dnn_inputs_v2::number_of_inner_cell * dnn_inputs_v2::number_of_inner_cell)
0319                      : (dnn_inputs_v2::number_of_outer_cell * dnn_inputs_v2::number_of_outer_cell);
0320 
0321   p_egammaBlockInputs->insert(
0322       p_egammaBlockInputs->end(), n_cells * dnn_inputs_v2::EgammaBlockInputs::NumberOfInputs, 0.);
0323   std::vector<float>::iterator egammaIter = p_egammaBlockInputs->begin();
0324 
0325   p_muonBlockInputs->insert(p_muonBlockInputs->end(), n_cells * dnn_inputs_v2::MuonBlockInputs::NumberOfInputs, 0.);
0326   std::vector<float>::iterator muonIter = p_muonBlockInputs->begin();
0327 
0328   p_hadronBlockInputs->insert(
0329       p_hadronBlockInputs->end(), n_cells * dnn_inputs_v2::HadronBlockInputs::NumberOfInputs, 0.);
0330   std::vector<float>::iterator hadronIter = p_hadronBlockInputs->begin();
0331 
0332   unsigned idx = 0;
0333   for (int eta = -grid.maxEtaIndex(); eta <= grid.maxEtaIndex(); ++eta) {
0334     for (int phi = -grid.maxPhiIndex(); phi <= grid.maxPhiIndex(); ++phi) {
0335       if (debug_level >= 2) {
0336         std::cout << "processing ( eta = " << eta << ", phi = " << phi << " )" << std::endl;
0337       }
0338       const CellIndex cell_index{eta, phi};
0339 
0340       const auto cell_iter = grid.find(cell_index);
0341       if (cell_iter != grid.end()) {
0342         if (debug_level >= 2) {
0343           std::cout << " creating inputs for ( eta = " << eta << ", phi = " << phi << " ): idx = " << idx << std::endl;
0344         }
0345         const Cell& cell = cell_iter->second;
0346         const int eta_index = grid.getEtaTensorIndex(cell_index);
0347         const int phi_index = grid.getPhiTensorIndex(cell_index);
0348         std::vector<float>::iterator egammaIterCell =
0349             egammaIter + (eta_index * n_cell_oneside + phi_index) * dnn_inputs_v2::EgammaBlockInputs::NumberOfInputs;
0350         std::vector<float>::iterator muonIterCell =
0351             muonIter + (eta_index * n_cell_oneside + phi_index) * dnn_inputs_v2::MuonBlockInputs::NumberOfInputs;
0352         std::vector<float>::iterator hadronIterCell =
0353             hadronIter + (eta_index * n_cell_oneside + phi_index) * dnn_inputs_v2::HadronBlockInputs::NumberOfInputs;
0354 
0355         createEgammaBlockInputs<CandidateCastType>(
0356             idx, tau, tau_index, tau_ref, pv, rho, electrons, pfCands, cell, tau_funcs, is_inner, egammaIterCell);
0357         createMuonBlockInputs<CandidateCastType>(
0358             idx, tau, tau_index, tau_ref, pv, rho, muons, pfCands, cell, tau_funcs, is_inner, muonIterCell);
0359         createHadronsBlockInputs<CandidateCastType>(
0360             idx, tau, tau_index, tau_ref, pv, rho, pfCands, cell, tau_funcs, is_inner, hadronIterCell);
0361 
0362         idx += 1;
0363       } else {
0364         if (debug_level >= 2) {
0365           std::cout << " skipping creation of inputs, because ( eta = " << eta << ", phi = " << phi
0366                     << " ) is not in the grid !!" << std::endl;
0367         }
0368         // we need to fill in the zeros for the input tensors
0369         idx += 1;
0370       }
0371     }
0372   }
0373 }
0374 
0375 void DeepTauIdSonicProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0376   edm::ParameterSetDescription desc;
0377   TritonClient::fillPSetDescription(desc);
0378   fillDescriptionsHelper(desc);
0379   descriptions.add("DeepTauIdSonicProducer", desc);
0380 }
0381 
0382 #include "FWCore/Framework/interface/MakerMacros.h"
0383 DEFINE_FWK_MODULE(DeepTauIdSonicProducer);
[ Source navigation ] [ Diff markup ] [ Identifier search ] [ general search ]

This page was automatically generated by the 2.2.1 LXR engine.

The LXR team
Valid CSS 2.1! Valid HTML 4.01!