PFProducer/interface/MLPFModel.h

0001 #ifndef RecoParticleFlow_PFProducer_interface_MLPFModel
0002 #define RecoParticleFlow_PFProducer_interface_MLPFModel
0003
0004 #include "FWCore/Framework/interface/Event.h"
0005 #include "DataFormats/ParticleFlowReco/interface/PFBlockElement.h"
0006 #include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"
0007
0008 namespace reco::mlpf {
0009   //The model takes the following number of features for each input PFElement
0010   static constexpr unsigned int NUM_ELEMENT_FEATURES = 25;
0011   static constexpr unsigned int NUM_OUTPUT_FEATURES = 14;
0012
0013   //these are defined at model creation time and set the random LSH codebook size
0014   static constexpr int LSH_BIN_SIZE = 64;
0015   static constexpr int NUM_MAX_ELEMENTS_BATCH = 200 * LSH_BIN_SIZE;
0016
0017   //In CPU mode, we want to evaluate each event separately
0018   static constexpr int BATCH_SIZE = 1;
0019
0020   //The model has 14 outputs for each particle:
0021   // out[0-7]: particle classification logits for each pdgId
0022   // out[8]: regressed charge
0023   // out[9]: regressed pt
0024   // out[10]: regressed eta
0025   // out[11]: regressed sin phi
0026   // out[12]: regressed cos phi
0027   // out[13]: regressed energy
0028   static constexpr unsigned int IDX_CLASS = 7;
0029
0030   static constexpr unsigned int IDX_CHARGE = 8;
0031
0032   static constexpr unsigned int IDX_PT = 9;
0033   static constexpr unsigned int IDX_ETA = 10;
0034   static constexpr unsigned int IDX_SIN_PHI = 11;
0035   static constexpr unsigned int IDX_COS_PHI = 12;
0036   static constexpr unsigned int IDX_ENERGY = 13;
0037
0038   //for consistency with the baseline PFAlgo
0039   static constexpr float PI_MASS = 0.13957;
0040
0041   //index [0, N_pdgids) -> PDGID
0042   //this maps the absolute values of the predicted PDGIDs to an array of ascending indices
0043   static const std::vector<int> pdgid_encoding = {0, 211, 130, 1, 2, 22, 11, 13};
0044
0045   //PFElement::type -> index [0, N_types)
0046   //this maps the type of the PFElement to an ascending index that is used by the model to distinguish between different elements
0047   static const std::map<int, int> elem_type_encoding = {
0048       {0, 0},
0049       {1, 1},
0050       {2, 2},
0051       {3, 3},
0052       {4, 4},
0053       {5, 5},
0054       {6, 6},
0055       {7, 7},
0056       {8, 8},
0057       {9, 9},
0058       {10, 10},
0059       {11, 11},
0060   };
0061
0062   std::array<float, NUM_ELEMENT_FEATURES> getElementProperties(const reco::PFBlockElement& orig);
0063   float normalize(float in);
0064
0065   int argMax(std::vector<float> const& vec);
0066
0067   reco::PFCandidate makeCandidate(int pred_pid,
0068                                   int pred_charge,
0069                                   float pred_pt,
0070                                   float pred_eta,
0071                                   float pred_sin_phi,
0072                                   float pred_cos_phi,
0073                                   float pred_e);
0074
0075   const std::vector<const reco::PFBlockElement*> getPFElements(const reco::PFBlockCollection& blocks);
0076
0077   void setCandidateRefs(reco::PFCandidate& cand,
0078                         const std::vector<const reco::PFBlockElement*> elems,
0079                         size_t ielem_originator);
0080 };  // namespace reco::mlpf
0081
0082 #endif