Skip to content

Commit

Permalink
factorize
Browse files Browse the repository at this point in the history
  • Loading branch information
@jpata
jpata committed Nov 5, 2020
1 parent 5b8b123 commit 85e6622
Show file tree
Hide file tree
Showing 6 changed files with 285 additions and 257 deletions.
2 changes: 0 additions & 2 deletions RecoParticleFlow/PFProducer/BuildFile.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -12,8 +12,6 @@
<use name="RecoParticleFlow/PFClusterTools"/>
<use name="RecoParticleFlow/PFTracking"/>
<use name="RecoEcal/EgammaCoreTools"/>
<use name="HeterogeneousCore/SonicCore"/>
<use name="HeterogeneousCore/SonicTriton"/>
<use name="boost"/>
<use name="clhep"/>
<use name="rootmath"/>
Expand Down
65 changes: 65 additions & 0 deletions RecoParticleFlow/PFProducer/interface/MLPFModel.h
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,65 @@
#ifndef RecoParticleFlow_PFProducer_interface_MLPFModel
#define RecoParticleFlow_PFProducer_interface_MLPFModel

#include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlock.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementSuperCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementGsfTrack.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementTrack.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementBrem.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElement.h"
#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;

//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;

//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};

//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);

int arg_max(std::vector<float> const& vec);
};

#endif
2 changes: 2 additions & 0 deletions RecoParticleFlow/PFProducer/plugins/BuildFile.xml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -80,5 +80,7 @@
<use name="RecoParticleFlow/PFTracking"/>
<use name="RecoEcal/EgammaCoreTools"/>
<use name="PhysicsTools/TensorFlow"/>
<use name="HeterogeneousCore/SonicCore"/>
<use name="HeterogeneousCore/SonicTriton"/>
<flags EDM_PLUGIN="1"/>
</library>
207 changes: 5 additions & 202 deletions RecoParticleFlow/PFProducer/plugins/MLPFProducer.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,200 +3,16 @@
#include "FWCore/Framework/interface/Event.h"
#include "FWCore/Framework/interface/MakerMacros.h"

#include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlock.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementSuperCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementGsfTrack.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementTrack.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementBrem.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElementCluster.h"
#include "DataFormats/ParticleFlowReco/interface/PFBlockElement.h"
#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"
#include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
#include "RecoParticleFlow/PFProducer/interface/MLPFModel.h"

#include "TLorentzVector.h"

//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;

//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;

//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},
};

struct MLPFCache {
std::atomic<tensorflow::GraphDef*> graph_def;
};

//Prepares the input array of floats for a single PFElement
std::array<float, NUM_ELEMENT_FEATURES> get_element_properties(const reco::PFBlockElement& orig) {
const auto type = orig.type();
float pt = 0.0;
//these are placeholders for the the future
[[maybe_unused]] float deltap = 0.0;
[[maybe_unused]] float sigmadeltap = 0.0;
[[maybe_unused]] float px = 0.0;
[[maybe_unused]] float py = 0.0;
[[maybe_unused]] float pz = 0.0;
float eta = 0.0;
float phi = 0.0;
float energy = 0.0;
float trajpoint = 0.0;
float eta_ecal = 0.0;
float phi_ecal = 0.0;
float eta_hcal = 0.0;
float phi_hcal = 0.0;
float charge = 0;
float layer = 0;
float depth = 0;
float muon_dt_hits = 0.0;
float muon_csc_hits = 0.0;

if (type == reco::PFBlockElement::TRACK) {
const auto& matched_pftrack = orig.trackRefPF();
if (matched_pftrack.isNonnull()) {
const auto& atECAL = matched_pftrack->extrapolatedPoint(reco::PFTrajectoryPoint::ECALShowerMax);
const auto& atHCAL = matched_pftrack->extrapolatedPoint(reco::PFTrajectoryPoint::HCALEntrance);
if (atHCAL.isValid()) {
eta_hcal = atHCAL.positionREP().eta();
phi_hcal = atHCAL.positionREP().phi();
}
if (atECAL.isValid()) {
eta_ecal = atECAL.positionREP().eta();
phi_ecal = atECAL.positionREP().phi();
}
}
const auto& ref = ((const reco::PFBlockElementTrack*)&orig)->trackRef();
pt = ref->pt();
px = ref->px();
py = ref->py();
pz = ref->pz();
eta = ref->eta();
phi = ref->phi();
energy = ref->pt() * cosh(ref->eta());
charge = ref->charge();

reco::MuonRef muonRef = orig.muonRef();
if (muonRef.isNonnull()) {
reco::TrackRef standAloneMu = muonRef->standAloneMuon();
if (standAloneMu.isNonnull()) {
muon_dt_hits = standAloneMu->hitPattern().numberOfValidMuonDTHits();
muon_csc_hits = standAloneMu->hitPattern().numberOfValidMuonCSCHits();
}
}

} else if (type == reco::PFBlockElement::BREM) {
const auto* orig2 = (const reco::PFBlockElementBrem*)&orig;
const auto& ref = orig2->GsftrackRef();
if (ref.isNonnull()) {
deltap = orig2->DeltaP();
sigmadeltap = orig2->SigmaDeltaP();
pt = ref->pt();
px = ref->px();
py = ref->py();
pz = ref->pz();
eta = ref->eta();
phi = ref->phi();
energy = ref->pt() * cosh(ref->eta());
trajpoint = orig2->indTrajPoint();
charge = ref->charge();
}
} else if (type == reco::PFBlockElement::GSF) {
//requires to keep GsfPFRecTracks
const auto* orig2 = (const reco::PFBlockElementGsfTrack*)&orig;
pt = orig2->Pin().pt();
px = orig2->Pin().px();
py = orig2->Pin().py();
pz = orig2->Pin().pz();
eta = orig2->Pin().eta();
phi = orig2->Pin().phi();
energy = pt * cosh(eta);
if (!orig2->GsftrackRefPF().isNull()) {
charge = orig2->GsftrackRefPF()->charge();
}
} else if (type == reco::PFBlockElement::ECAL || type == reco::PFBlockElement::PS1 ||
type == reco::PFBlockElement::PS2 || type == reco::PFBlockElement::HCAL ||
type == reco::PFBlockElement::HO || type == reco::PFBlockElement::HFHAD ||
type == reco::PFBlockElement::HFEM) {
const auto& ref = ((const reco::PFBlockElementCluster*)&orig)->clusterRef();
if (ref.isNonnull()) {
eta = ref->eta();
phi = ref->phi();
px = ref->position().x();
py = ref->position().y();
pz = ref->position().z();
energy = ref->energy();
layer = ref->layer();
depth = ref->depth();
}
} else if (type == reco::PFBlockElement::SC) {
const auto& clref = ((const reco::PFBlockElementSuperCluster*)&orig)->superClusterRef();
if (clref.isNonnull()) {
eta = clref->eta();
phi = clref->phi();
px = clref->position().x();
py = clref->position().y();
pz = clref->position().z();
energy = clref->energy();
}
}

float typ_idx = static_cast<float>(elem_type_encoding.at(orig.type()));

//Must be the same order as in tf_model.py
return std::array<float, NUM_ELEMENT_FEATURES>({{typ_idx,
pt,
eta,
phi,
energy,
layer,
depth,
charge,
trajpoint,
eta_ecal,
phi_ecal,
eta_hcal,
phi_hcal,
muon_dt_hits,
muon_csc_hits}});
}

class MLPFProducer : public edm::stream::EDProducer<edm::GlobalCache<MLPFCache> > {
public:
explicit MLPFProducer(const edm::ParameterSet&, const MLPFCache*);
Expand All @@ -221,22 +37,9 @@ MLPFProducer::MLPFProducer(const edm::ParameterSet& cfg, const MLPFCache* cache)
session_ = tensorflow::createSession(cache->graph_def);
}

//returns the argmax of a given std::vector<T>
template <typename T, typename A>
int arg_max(std::vector<T, A> const& vec) {
return static_cast<int>(std::distance(vec.begin(), max_element(vec.begin(), vec.end())));
}

float normalize(float in) {
if (std::abs(in) > 1e4) {
return 0.0;
} else if (std::isnan(in)) {
return 0.0;
}
return in;
}

void MLPFProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
using namespace reco::mlpf;

auto blocks_handle = event.getHandle(inputTagBlocks_);
const auto& blocks = *blocks_handle;
std::vector<reco::PFCandidate> pOutputCandidateCollection;
Expand Down Expand Up @@ -346,7 +149,7 @@ void MLPFProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
cand.setCharge(charge);

//set the track ref in case the originating element was a track
if (elem->type() == reco::PFBlockElement::TRACK && charge != 0) {
if (elem->type() == reco::PFBlockElement::TRACK && charge != 0 && elem->trackRef().isNonnull()) {
cand.setTrackRef(elem->trackRef());

//set the muon ref in case the originator was a muon
Expand All @@ -366,7 +169,7 @@ void MLPFProducer::produce(edm::Event& event, const edm::EventSetup& setup) {
} //loop over PFElements

edm::LogWarning("MLPFProducer") << "num_elements_total=" << num_elements_total
<< " num_pred_particles=" << num_pred_particles << std::endl;
<< " num_pred_particles=" << num_pred_particles << std::endl;

event.emplace(pfCandidatesToken_, pOutputCandidateCollection);
}
Expand Down
Loading

0 comments on commit 85e6622

Please sign in to comment.