Skip to content

Commit

Permalink
format
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpata
jpata committed Nov 6, 2020
1 parent 93abcd8 commit 34a7931
Show file tree
Hide file tree
Showing 3 changed files with 178 additions and 179 deletions.
84 changes: 42 additions & 42 deletions RecoParticleFlow/PFProducer/interface/MLPFModel.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,54 +12,54 @@
#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"

namespace reco::mlpf {
//The model takes the following number of features for each input PFElement
static constexpr unsigned int NUM_ELEMENT_FEATURES = 15;
//The model takes the following number of features for each input PFElement
static constexpr unsigned int NUM_ELEMENT_FEATURES = 15;

//these are defined at model creation time and set the random LSH codebook size
static constexpr int NUM_MAX_ELEMENTS_BATCH = 20000;
static constexpr int LSH_BIN_SIZE = 100;
//these are defined at model creation time and set the random LSH codebook size
static constexpr int NUM_MAX_ELEMENTS_BATCH = 20000;
static constexpr int LSH_BIN_SIZE = 100;

//In CPU mode, we only want to evaluate each event separately
static constexpr int BATCH_SIZE = 1;
//In CPU mode, we only want to evaluate each event separately
static constexpr int BATCH_SIZE = 1;

//The model has 12 outputs for each particle:
// out[0-7]: particle classification logits
// out[8]: regressed eta
// out[9]: regressed phi
// out[10]: regressed energy
// out[11]: regressed charge logit
static constexpr unsigned int NUM_OUTPUTS = 12;
static constexpr unsigned int NUM_CLASS = 7;
static constexpr unsigned int IDX_ETA = 8;
static constexpr unsigned int IDX_PHI = 9;
static constexpr unsigned int IDX_ENERGY = 10;
static constexpr unsigned int IDX_CHARGE = 11;
//The model has 12 outputs for each particle:
// out[0-7]: particle classification logits
// out[8]: regressed eta
// out[9]: regressed phi
// out[10]: regressed energy
// out[11]: regressed charge logit
static constexpr unsigned int NUM_OUTPUTS = 12;
static constexpr unsigned int NUM_CLASS = 7;
static constexpr unsigned int IDX_ETA = 8;
static constexpr unsigned int IDX_PHI = 9;
static constexpr unsigned int IDX_ENERGY = 10;
static constexpr unsigned int IDX_CHARGE = 11;

//index [0, N_pdgids) -> PDGID
//this maps the absolute values of the predicted PDGIDs to an array of ascending indices
static const std::vector<int> pdgid_encoding = {0, 1, 2, 11, 13, 22, 130, 211};
//index [0, N_pdgids) -> PDGID
//this maps the absolute values of the predicted PDGIDs to an array of ascending indices
static const std::vector<int> pdgid_encoding = {0, 1, 2, 11, 13, 22, 130, 211};

//PFElement::type -> index [0, N_types)
//this maps the type of the PFElement to an ascending index that is used by the model to distinguish between different elements
static const std::map<int, int> elem_type_encoding = {
{0, 0},
{1, 1},
{2, 2},
{3, 3},
{4, 4},
{5, 5},
{6, 6},
{7, 7},
{8, 8},
{9, 9},
{10, 10},
{11, 11},
};
//PFElement::type -> index [0, N_types)
//this maps the type of the PFElement to an ascending index that is used by the model to distinguish between different elements
static const std::map<int, int> elem_type_encoding = {
{0, 0},
{1, 1},
{2, 2},
{3, 3},
{4, 4},
{5, 5},
{6, 6},
{7, 7},
{8, 8},
{9, 9},
{10, 10},
{11, 11},
};

std::array<float, NUM_ELEMENT_FEATURES> get_element_properties(const reco::PFBlockElement& orig);
float normalize(float in);
std::array<float, NUM_ELEMENT_FEATURES> get_element_properties(const reco::PFBlockElement& orig);
float normalize(float in);

int arg_max(std::vector<float> const& vec);
};
int arg_max(std::vector<float> const& vec);
}; // namespace reco::mlpf

#endif
13 changes: 6 additions & 7 deletions RecoParticleFlow/PFProducer/plugins/MLPFProducerSonic.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -36,7 +36,7 @@
// git clone https://github.com/jpata/mlpf-models.git
// mv mlpf-models/triton/models ./
// singularity exec --nv docker://nvcr.io/nvidia/tritonserver:20.10-py3 tritonserver --model-repository models --log-verbose=1 --log-error=1 --log-info=1
//
//
// process.mlpfproducergpu = cms.EDProducer("MLPFProducerSonicTriton",
// Client = cms.PSet(
// mode = cms.string("Sync"),
Expand All @@ -51,7 +51,6 @@
// )
// )


class MLPFProducerSonicTriton : public SonicEDProducer<TritonClient> {
public:
//needed because base class has dependent scope
Expand Down Expand Up @@ -136,24 +135,24 @@ class MLPFProducerSonicTriton : public SonicEDProducer<TritonClient> {
//get the coefficients in the output corresponding to the class probabilities (raw logits)
std::vector<float> pred_id_logits;
for (unsigned int idx_id = 0; idx_id <= NUM_CLASS; idx_id++) {
pred_id_logits.push_back(out_data[0][ielem*NUM_OUTPUTS + idx_id]);
pred_id_logits.push_back(out_data[0][ielem * NUM_OUTPUTS + idx_id]);
}

//get the most probable class PDGID
int pred_pid = pdgid_encoding[arg_max(pred_id_logits)];

//get the predicted momentum components
float pred_eta = out_data[0][ielem*NUM_OUTPUTS + IDX_ETA];
float pred_phi = out_data[0][ielem*NUM_OUTPUTS + IDX_PHI];
float pred_eta = out_data[0][ielem * NUM_OUTPUTS + IDX_ETA];
float pred_phi = out_data[0][ielem * NUM_OUTPUTS + IDX_PHI];

//put phi into a modular range
pred_phi = TMath::ATan2(TMath::Sin(pred_phi), TMath::Cos(pred_phi));
float pred_e = out_data[0][ielem*NUM_OUTPUTS + IDX_ENERGY];
float pred_e = out_data[0][ielem * NUM_OUTPUTS + IDX_ENERGY];

//currently, set the pT from a massless approximation.
//later versions of the model may predict predict both the energy and pT of the particle
float pred_pt = pred_e / cosh(pred_eta);
float pred_charge = out_data[0][ielem*NUM_OUTPUTS + IDX_CHARGE];
float pred_charge = out_data[0][ielem * NUM_OUTPUTS + IDX_CHARGE];

//a particle was predicted for this PFElement
if (pred_pid != 0) {
Expand Down
Loading

0 comments on commit 34a7931

Please sign in to comment.