Phase2L1ParticleFlow/src/NNVtxAssoc.cc

0001 // Package:      L1Trigger/Phase2L1ParticleFlow
0002 // Class:        NNVtxAssoc
0003 // Description:  Designed to run the track to vertex associations created by the E2E NNVtx.
0004 //               TTTrackNetworkSelector either accepts or rejects that a PF object's (t) track is associated to a vertex (v).
0005 // Authors:      Kai Hong Law and Benjamin Radburn-Smith
0006 // Created:      February 2025
0007
0008 #include "L1Trigger/Phase2L1ParticleFlow/interface/NNVtxAssoc.h"
0009 #include "DataFormats/L1TParticleFlow/interface/PFTrack.h"
0010 #include "PhysicsTools/TensorFlow/interface/TensorFlow.h"
0011 #include <iomanip>
0012
0013 NNVtxAssoc::NNVtxAssoc(std::string AssociationGraphPath,
0014                        const double AssociationThreshold,
0015                        const std::vector<double>& AssociationNetworkZ0binning,
0016                        const std::vector<double>& AssociationNetworkEtaBounds,
0017                        const std::vector<double>& AssociationNetworkZ0ResBins)
0018     : associationThreshold_(AssociationThreshold),
0019       z0_binning_(AssociationNetworkZ0binning),
0020       eta_bins_(AssociationNetworkEtaBounds),
0021       res_bins_(AssociationNetworkZ0ResBins) {
0022   tensorflow::GraphDef* associationGraph_ = tensorflow::loadGraphDef(AssociationGraphPath);
0023   associationSesh_ = tensorflow::createSession(associationGraph_);
0024   log_.setf(std::ios::fixed, std::ios::floatfield);
0025   log_.precision(3);
0026 }
0027
0028 template <typename T>
0029 bool NNVtxAssoc::TTTrackNetworkSelector(const PFRegionEmu& region, const T& t, const l1ct::PVObjEmu& v) {
0030   tensorflow::Tensor inputAssoc(tensorflow::DT_FLOAT, {1, 4});
0031   std::vector<tensorflow::Tensor> outputAssoc;
0032
0033   auto lower = std::lower_bound(eta_bins_.begin(), eta_bins_.end(), region.floatGlbEta(t.hwVtxEta()));
0034
0035   int resbin = std::distance(eta_bins_.begin(), lower);
0036   float binWidth = z0_binning_[2];
0037   // Calculate integer dZ from track z0 and vertex z0 (use floating point version and convert internally allowing use of both emulator and simulator vertex and track)
0038   float dZ =
0039       abs(floor(((t.floatZ0() + z0_binning_[1]) / (binWidth))) - floor(((v.floatZ0() + z0_binning_[1]) / (binWidth))));
0040
0041   // The following constants <22, 9> are defined by the quantisation of the Neural Network
0042   ap_ufixed<22, 9> ptEmulation_rescale = t.hwPt;
0043   ap_ufixed<22, 9> resBinEmulation_rescale = res_bins_[resbin];
0044   ap_ufixed<22, 9> MVAEmulation_rescale = 0;
0045   ap_ufixed<22, 9> dZEmulation_rescale = dZ;
0046
0047   // Deal with this template class using 2 different objects (t) which have different calls to their PFTracks:
0048   const l1t::PFTrack* srcTrack = nullptr;
0049   if constexpr (std::is_same_v<T, const l1ct::TkObjEmu>)
0050     srcTrack = t.src;
0051   else if constexpr (std::is_same_v<T, const l1ct::PFChargedObjEmu>)
0052     srcTrack = t.srcTrack;
0053   if (srcTrack)
0054     MVAEmulation_rescale = srcTrack->trackWord().getMVAQualityBits();
0055
0056   inputAssoc.tensor<float, 2>()(0, 0) = ptEmulation_rescale.to_double();
0057   inputAssoc.tensor<float, 2>()(0, 1) = MVAEmulation_rescale.to_double();
0058   inputAssoc.tensor<float, 2>()(0, 2) = resBinEmulation_rescale.to_double() / 16.0;
0059   inputAssoc.tensor<float, 2>()(0, 3) = dZEmulation_rescale.to_double();
0060
0061   // Run Association Network:
0062   tensorflow::run(associationSesh_, {{"NNvtx_track_association:0", inputAssoc}}, {"Identity:0"}, &outputAssoc);
0063
0064   double NNOutput = (double)outputAssoc[0].tensor<float, 2>()(0, 0);
0065   double NNOutput_exp = 1.0 / (1.0 + exp(-1.0 * (NNOutput)));
0066
0067   return NNOutput_exp >= associationThreshold_;
0068 }
0069
0070 #ifdef CMSSW_GIT_HASH
0071 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0072 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0073
0074 edm::ParameterSetDescription NNVtxAssoc::getParameterSetDescription() {
0075   edm::ParameterSetDescription description;
0076   description.add<double>("associationThreshold");
0077   description.add<std::string>("associationGraph");
0078   description.add<std::vector<double>>("associationNetworkZ0binning");
0079   description.add<std::vector<double>>("associationNetworkEtaBounds");
0080   description.add<std::vector<double>>("associationNetworkZ0ResBins");
0081   return description;
0082 }
0083 #endif
0084
0085 void NNVtxAssoc::NNVtxAssocDebug() {
0086   log_ << "-- NNVtxAssocDebug --\n";
0087   log_ << "AssociationThreshold: " << this->associationThreshold_ << "\n";
0088   log_ << "z0_binning: ";
0089   for (auto i : this->z0_binning_)
0090     log_ << i << " ";
0091   log_ << "\n";
0092   log_ << "eta_bins: ";
0093   for (auto i : this->eta_bins_)
0094     log_ << i << " ";
0095   log_ << "\n";
0096   log_ << "res_bins: ";
0097   for (auto i : this->res_bins_)
0098     log_ << i << " ";
0099   log_ << "\n";
0100   edm::LogPrint("NNVtxAssoc") << log_.str();
0101 }
0102
0103 template bool NNVtxAssoc::TTTrackNetworkSelector<const l1ct::TkObjEmu>(const PFRegionEmu&,
0104                                                                        const l1ct::TkObjEmu&,
0105                                                                        const l1ct::PVObjEmu&);
0106 template bool NNVtxAssoc::TTTrackNetworkSelector<const l1ct::PFChargedObjEmu>(const PFRegionEmu&,
0107                                                                               const l1ct::PFChargedObjEmu&,
0108                                                                               const l1ct::PVObjEmu&);