diff --git a/Configuration/ProcessModifiers/python/MLPFSonicTriton_cff.py b/Configuration/ProcessModifiers/python/MLPFSonicTriton_cff.py new file mode 100644 index 0000000000000..56e76e0cb77e8 --- /dev/null +++ b/Configuration/ProcessModifiers/python/MLPFSonicTriton_cff.py @@ -0,0 +1,3 @@ +import FWCore.ParameterSet.Config as cms +# For MLPF SONIC +MLPFSonicTriton = cms.Modifier() diff --git a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py index 7c8e6c428b0c6..cfd1949662b43 100644 --- a/Configuration/ProcessModifiers/python/allSonicTriton_cff.py +++ b/Configuration/ProcessModifiers/python/allSonicTriton_cff.py @@ -7,8 +7,8 @@ from Configuration.ProcessModifiers.deepTauSonicTriton_cff import deepTauSonicTriton from Configuration.ProcessModifiers.particleTransformerAK4SonicTriton_cff import particleTransformerAK4SonicTriton from Configuration.ProcessModifiers.unifiedparticleTransformerAK4SonicTriton_cff import unifiedparticleTransformerAK4SonicTriton +from Configuration.ProcessModifiers.MLPFSonicTriton_cff import MLPFSonicTriton # collect all SonicTriton-related process modifiers here -allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton) - +allSonicTriton = cms.ModifierChain(enableSonicTriton,deepMETSonicTriton,particleNetSonicTriton,particleNetPTSonicTriton,deepTauSonicTriton,particleTransformerAK4SonicTriton,unifiedparticleTransformerAK4SonicTriton,MLPFSonicTriton) diff --git a/RecoParticleFlow/Configuration/python/RecoParticleFlow_cff.py b/RecoParticleFlow/Configuration/python/RecoParticleFlow_cff.py index 586a7fe1ef071..59d9197a569f7 100644 --- a/RecoParticleFlow/Configuration/python/RecoParticleFlow_cff.py +++ b/RecoParticleFlow/Configuration/python/RecoParticleFlow_cff.py @@ -97,6 +97,23 @@ from Configuration.ProcessModifiers.mlpf_cff import mlpf from RecoParticleFlow.PFProducer.mlpfProducer_cfi import mlpfProducer mlpf.toReplaceWith(particleFlowTmp, mlpfProducer) +#For MLPF_SONIC +from Configuration.ProcessModifiers.MLPFSonicTriton_cff import MLPFSonicTriton +from RecoParticleFlow.PFProducer.mlpfsonicProducer_cfi import mlpfsonicProducer +MLPFSonicTriton.toReplaceWith(particleFlowTmp, mlpfsonicProducer.clone( + Client = cms.PSet( + timeout = cms.untracked.uint32(300), + mode = cms.string("Async"), + modelName = cms.string("MLPF"), + modelConfigPath = cms.FileInPath("RecoParticleFlow/PFProducer/data/models/MLPF/config.pbtxt"), + modelVersion = cms.string(""), + verbose = cms.untracked.bool(False), + allowedTries = cms.untracked.uint32(0), + useSharedMemory = cms.untracked.bool(True), + compression = cms.untracked.string(""), + ), +)) + # # switch from pfTICL to simPF diff --git a/RecoParticleFlow/PFProducer/BuildFile.xml b/RecoParticleFlow/PFProducer/BuildFile.xml index 98351dbc8143f..164e306f85db1 100644 --- a/RecoParticleFlow/PFProducer/BuildFile.xml +++ b/RecoParticleFlow/PFProducer/BuildFile.xml @@ -10,6 +10,7 @@ + diff --git a/RecoParticleFlow/PFProducer/plugins/BuildFile.xml b/RecoParticleFlow/PFProducer/plugins/BuildFile.xml index 6c4c7c4a9bcb8..ac104096fabea 100644 --- a/RecoParticleFlow/PFProducer/plugins/BuildFile.xml +++ b/RecoParticleFlow/PFProducer/plugins/BuildFile.xml @@ -81,5 +81,6 @@ + diff --git a/RecoParticleFlow/PFProducer/plugins/MLPFSONICProducer.cc b/RecoParticleFlow/PFProducer/plugins/MLPFSONICProducer.cc new file mode 100644 index 0000000000000..3bb7b5a38d024 --- /dev/null +++ b/RecoParticleFlow/PFProducer/plugins/MLPFSONICProducer.cc @@ -0,0 +1,226 @@ +#include "FWCore/Framework/interface/Frameworkfwd.h" +#include "FWCore/Framework/interface/Event.h" +#include "FWCore/Framework/interface/MakerMacros.h" +#include "FWCore/ParameterSet/interface/ParameterSet.h" +#include "FWCore/Utilities/interface/StreamID.h" + +#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h" +//#include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h" +#include "RecoParticleFlow/PFProducer/interface/MLPFModel.h" + +#include "DataFormats/ParticleFlowReco/interface/PFBlockElementTrack.h" + +#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h" +#include "HeterogeneousCore/SonicTriton/interface/TritonData.h" + +// using namespace cms::Ort; +class MLPFSONICProducer : public TritonEDProducer<> { +public: + explicit MLPFSONICProducer(const edm::ParameterSet &); + ~MLPFSONICProducer() override; + + void acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) override; + + void produce(edm::Event &iEvent, edm::EventSetup const &iSetup, Output const &iOutput) override; + static void fillDescriptions(edm::ConfigurationDescriptions &); + +private: + const edm::EDPutTokenT pfCandidatesPutToken_; + const edm::EDGetTokenT> gsfElectrons_; + const edm::EDGetTokenT inputTagBlocks_; + std::vector input_names_; + std::vector output_names_; + // Declare inputs as a private member variable + std::vector> inputs; +}; + +class SelectedElementsManager { +public: + static SelectedElementsManager &getInstance() { + static SelectedElementsManager instance; // Single instance for the program + return instance; + } + + void fill(const std::vector &all_elements) { + selected_elements_.clear(); + + for (const auto *pelem : all_elements) { + if (pelem->type() == reco::PFBlockElement::PS1 || pelem->type() == reco::PFBlockElement::PS2 || + pelem->type() == reco::PFBlockElement::BREM) { + continue; + } + selected_elements_.push_back(pelem); + } + } + + const std::vector &get() const { return selected_elements_; } + +private: + SelectedElementsManager() = default; + ~SelectedElementsManager() = default; + + SelectedElementsManager(const SelectedElementsManager &) = delete; + SelectedElementsManager &operator=(const SelectedElementsManager &) = delete; + + std::vector selected_elements_; +}; + +MLPFSONICProducer::MLPFSONICProducer(const edm::ParameterSet &iConfig) + : TritonEDProducer<>(iConfig), + pfCandidatesPutToken_{produces()}, + gsfElectrons_{consumes>(edm::InputTag("gedGsfElectronsTmp"))}, + inputTagBlocks_{consumes(iConfig.getParameter("src"))}, + input_names_(iConfig.getParameter>("input_names")), + output_names_(iConfig.getParameter>("output_names")) {} + +MLPFSONICProducer::~MLPFSONICProducer() {} +void MLPFSONICProducer::acquire(edm::Event const &iEvent, edm::EventSetup const &iSetup, Input &iInput) { + using namespace reco::mlpf; + const auto &blocks = iEvent.get(inputTagBlocks_); + const auto &all_elements = getPFElements(blocks); + + const auto &gsfElectrons = iEvent.get(gsfElectrons_); + + SelectedElementsManager::getInstance().fill(all_elements); // Fill data once + std::cout << "filled selected_elements." << std::endl; + const auto &selected_elements = SelectedElementsManager::getInstance().get(); + // Total Number of selected_elements + unsigned int num_elements_total = selected_elements.size(); + + const auto tensor_size = num_elements_total; + + //Fill the input tensor (batch, elems, features) = (1, tensor_size, NUM_ELEMENT_FEATURES) + inputs.resize(2); + inputs[0].assign(tensor_size * NUM_ELEMENT_FEATURES, 0.0); + inputs[1].assign(tensor_size, 0.0); + + unsigned int ielem = 0; + const auto &mask_name = input_names_[0]; + const auto &X_name = input_names_[1]; + + auto &data1 = iInput.at(mask_name); + auto &data2 = iInput.at(X_name); + data1.setShape(0, tensor_size); + data2.setShape(0, tensor_size); + auto tdata1 = data1.allocate(true); + auto tdata2 = data2.allocate(true); + for (const auto *pelem : selected_elements) { + if (ielem > tensor_size) { + continue; + } + + const auto &elem = *pelem; + + //prepare the input array from the PFElement + const auto &props = getElementProperties(elem, gsfElectrons).as_array(); + + //copy features to the input array + for (unsigned int iprop = 0; iprop < NUM_ELEMENT_FEATURES; iprop++) { + inputs[0][ielem * NUM_ELEMENT_FEATURES + iprop] = normalize(props[iprop]); + } + inputs[1][ielem] = 1.0; + ielem += 1; + } + auto &vdata1 = (*tdata1)[0]; + auto &vdata2 = (*tdata2)[0]; + vdata1 = inputs[1]; + vdata2 = inputs[0]; + data1.toServer(tdata1); + data2.toServer(tdata2); + std::cout << "check-point Producer-143_tensorsize_" << tensor_size << std::endl; +} +void MLPFSONICProducer::produce(edm::Event &iEvent, const edm::EventSetup &iSetup, Output const &iOutput) { + using namespace reco::mlpf; + const auto &bid_name = output_names_[0]; + const auto &id_name = output_names_[1]; + const auto &momentum_name = output_names_[2]; + const auto &output1 = iOutput.at(bid_name); + const auto &output2 = iOutput.at(id_name); + const auto &output3 = iOutput.at(momentum_name); + const auto &output_binary = output1.fromServer(); + const auto &output_pid = output2.fromServer(); + const auto &output_p4 = output3.fromServer(); + const auto &selected_elements = SelectedElementsManager::getInstance().get(); + // Total Number of selected_elements + unsigned int num_elements_total = selected_elements.size(); + unsigned int tensor_size = num_elements_total; + std::cout << "check-point pid-161_" << output_pid[0][0] << "__" << output_pid[0][1] << "__" << output_pid[0][2] + << std::endl; + std::cout << "check-point p4-162_" << output_p4[0][0] << "__" << output_p4[0][1] << "__" << output_p4[0][2] + << std::endl; + std::vector pOutputCandidateCollection; + for (size_t ielem = 0; ielem < num_elements_total; ielem++) { + std::vector pred_id_probas(pdgid_encoding.size(), 0.0); + const reco::PFBlockElement *elem = selected_elements[ielem]; + const auto logit_no_ptcl = output_binary[0][ielem * 2 + 0]; + const auto logit_ptcl = output_binary[0][ielem * 2 + 1]; + + // Check if the binary classifier of the model predicted a particle + int pred_pid = 0; + if (logit_ptcl > logit_no_ptcl) { + for (unsigned int idx_id = 0; idx_id < pred_id_probas.size(); idx_id++) { + auto pred_proba = output_pid[0][ielem * NUM_OUTPUT_FEATURES_CLS + idx_id]; + pred_id_probas[idx_id] = pred_proba; + } + + auto imax = argMax(pred_id_probas); + //get the most probable class PDGID + pred_pid = pdgid_encoding.at(imax); + } + //a particle was predicted for this PFElement, otherwise it was a spectator + if (pred_pid != 0) { + //muons and charged hadrons should only come from tracks, otherwise we won't have track references to pass downstream + if (((pred_pid == 13) || (pred_pid == 211)) && elem->type() != reco::PFBlockElement::TRACK) { + pred_pid = 130; + } + + float pred_charge = 0.0; + if (elem->type() == reco::PFBlockElement::TRACK) { + const auto *eltTrack = dynamic_cast(elem); + //for now, just take the charge from the track + if (eltTrack->trackRef().isNonnull()) { + pred_charge = eltTrack->trackRef()->charge(); + } + + //a track with no muon ref should not produce a muon candidate, instead we interpret it as a charged hadron here + if (pred_pid == 13 && eltTrack->muonRef().isNull()) { + pred_pid = 211; + } + + //taus are reconstructed downstream based on other criteria, instead we interpret it as a charged hadron here + if (pred_pid == 15) { + pred_pid = 211; + } + + //tracks from displaced vertices need reference debugging downstream as well, so we just treat them as neutrals for the moment + if ((pred_pid == 211) && (eltTrack->isLinkedToDisplacedVertex())) { + pred_pid = 130; + } + } + //get the predicted momentum components from the model + float pred_pt = output_p4[0][ielem * NUM_OUTPUT_FEATURES_P4 + IDX_PT]; + pred_pt = exp(pred_pt) * inputs[0][ielem * NUM_ELEMENT_FEATURES + 1]; + float pred_eta = output_p4[0][ielem * NUM_OUTPUT_FEATURES_P4 + IDX_ETA]; + float pred_sin_phi = output_p4[0][ielem * NUM_OUTPUT_FEATURES_P4 + IDX_SIN_PHI]; + float pred_cos_phi = output_p4[0][ielem * NUM_OUTPUT_FEATURES_P4 + IDX_COS_PHI]; + float pred_e = output_p4[0][ielem * NUM_OUTPUT_FEATURES_P4 + IDX_ENERGY]; + pred_e = exp(pred_e) * inputs[0][ielem * NUM_ELEMENT_FEATURES + 5]; + + auto cand = makeCandidate(pred_pid, pred_charge, pred_pt, pred_eta, pred_sin_phi, pred_cos_phi, pred_e); + setCandidateRefs(cand, selected_elements, ielem); + pOutputCandidateCollection.push_back(cand); + } + } //end loop of elements + iEvent.emplace(pfCandidatesPutToken_, pOutputCandidateCollection); +} + +void MLPFSONICProducer::fillDescriptions(edm::ConfigurationDescriptions &descriptions) { + edm::ParameterSetDescription desc; + TritonClient::fillPSetDescription(desc); + desc.add("src", edm::InputTag("particleFlowBlock")); + desc.add>("input_names", {"mask", "Xfeat_normed"}); + desc.add>("output_names", {"bid", "id", "momentum"}); + descriptions.addWithDefaultLabel(desc); +} + +DEFINE_FWK_MODULE(MLPFSONICProducer);