Back to home page

Project CMSSW displayed by LXR

 
 

    


File indexing completed on 2022-08-30 22:39:07

0001 // Author: Marco Rovere - marco.rovere@cern.ch
0002 // Date: 10/2021
0003 #include <algorithm>
0004 #include <set>
0005 #include <vector>
0006 
0007 #include "tbb/task_arena.h"
0008 #include "tbb/tbb.h"
0009 
0010 #include "DataFormats/Math/interface/deltaR.h"
0011 #include "DataFormats/Math/interface/LorentzVector.h"
0012 #include "DataFormats/Candidate/interface/Candidate.h"
0013 #include "FWCore/Framework/interface/Event.h"
0014 #include "FWCore/MessageLogger/interface/MessageLogger.h"
0015 #include "FWCore/Utilities/interface/Exception.h"
0016 #include "PatternRecognitionbyFastJet.h"
0017 
0018 #include "TrackstersPCA.h"
0019 #include "Geometry/CaloGeometry/interface/CaloGeometry.h"
0020 #include "Geometry/Records/interface/CaloGeometryRecord.h"
0021 #include "FWCore/Framework/interface/EventSetup.h"
0022 
0023 #include "fastjet/ClusterSequence.hh"
0024 
0025 using namespace ticl;
0026 using namespace fastjet;
0027 
0028 template <typename TILES>
0029 PatternRecognitionbyFastJet<TILES>::PatternRecognitionbyFastJet(const edm::ParameterSet &conf,
0030                                                                 edm::ConsumesCollector iC)
0031     : PatternRecognitionAlgoBaseT<TILES>(conf, iC),
0032       caloGeomToken_(iC.esConsumes<CaloGeometry, CaloGeometryRecord>()),
0033       antikt_radius_(conf.getParameter<double>("antikt_radius")),
0034       minNumLayerCluster_(conf.getParameter<int>("minNumLayerCluster")),
0035       eidInputName_(conf.getParameter<std::string>("eid_input_name")),
0036       eidOutputNameEnergy_(conf.getParameter<std::string>("eid_output_name_energy")),
0037       eidOutputNameId_(conf.getParameter<std::string>("eid_output_name_id")),
0038       eidMinClusterEnergy_(conf.getParameter<double>("eid_min_cluster_energy")),
0039       eidNLayers_(conf.getParameter<int>("eid_n_layers")),
0040       eidNClusters_(conf.getParameter<int>("eid_n_clusters")){};
0041 
0042 template <typename TILES>
0043 void PatternRecognitionbyFastJet<TILES>::buildJetAndTracksters(std::vector<PseudoJet> &fjInputs,
0044                                                                std::vector<ticl::Trackster> &result) {
0045   if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Basic) {
0046     edm::LogVerbatim("PatternRecogntionbyFastJet")
0047         << "Creating FastJet with " << fjInputs.size() << " LayerClusters in input";
0048   }
0049   fastjet::ClusterSequence sequence(fjInputs, JetDefinition(antikt_algorithm, antikt_radius_));
0050   auto jets = fastjet::sorted_by_pt(sequence.inclusive_jets(0));
0051   if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Basic) {
0052     edm::LogVerbatim("PatternRecogntionbyFastJet") << "FastJet produced " << jets.size() << " jets/trackster";
0053   }
0054 
0055   auto trackster_idx = result.size();
0056   auto jetsSize = std::count_if(jets.begin(), jets.end(), [=](fastjet::PseudoJet jet) {
0057     return jet.constituents().size() > static_cast<unsigned int>(minNumLayerCluster_);
0058   });
0059   result.resize(trackster_idx + jetsSize);
0060 
0061   for (const auto &pj : jets) {
0062     if (pj.constituents().size() > static_cast<unsigned int>(minNumLayerCluster_)) {
0063       for (const auto &component : pj.constituents()) {
0064         result[trackster_idx].vertices().push_back(component.user_index());
0065         result[trackster_idx].vertex_multiplicity().push_back(1);
0066         if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Basic) {
0067           edm::LogVerbatim("PatternRecogntionbyFastJet")
0068               << "Jet has " << pj.constituents().size() << " components that are stored in trackster " << trackster_idx;
0069         }
0070       }
0071       trackster_idx++;
0072     } else {
0073       if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Advanced) {
0074         edm::LogVerbatim("PatternRecogntionbyFastJet")
0075             << "Jet with " << pj.constituents().size() << " constituents discarded since too small wrt "
0076             << minNumLayerCluster_;
0077       }
0078     }
0079   }
0080   fjInputs.clear();
0081 }
0082 
0083 template <typename TILES>
0084 void PatternRecognitionbyFastJet<TILES>::makeTracksters(
0085     const typename PatternRecognitionAlgoBaseT<TILES>::Inputs &input,
0086     std::vector<Trackster> &result,
0087     std::unordered_map<int, std::vector<int>> &seedToTracksterAssociation) {
0088   // Protect from events with no seeding regions
0089   if (input.regions.empty())
0090     return;
0091 
0092   edm::EventSetup const &es = input.es;
0093   const CaloGeometry &geom = es.getData(caloGeomToken_);
0094   rhtools_.setGeometry(geom);
0095 
0096   constexpr auto isHFnose = std::is_same<TILES, TICLLayerTilesHFNose>::value;
0097   constexpr int nEtaBin = TILES::constants_type_t::nEtaBins;
0098   constexpr int nPhiBin = TILES::constants_type_t::nPhiBins;
0099 
0100   // We need to partition the two sides of the HGCAL detector
0101   auto lastLayerPerSide = static_cast<unsigned int>(rhtools_.lastLayer(isHFnose)) - 1;
0102   unsigned int maxLayer = 2 * lastLayerPerSide - 1;
0103   std::vector<fastjet::PseudoJet> fjInputs;
0104   fjInputs.clear();
0105   for (unsigned int currentLayer = 0; currentLayer <= maxLayer; ++currentLayer) {
0106     if (currentLayer == lastLayerPerSide) {
0107       buildJetAndTracksters(fjInputs, result);
0108     }
0109     const auto &tileOnLayer = input.tiles[currentLayer];
0110     for (int ieta = 0; ieta <= nEtaBin; ++ieta) {
0111       auto offset = ieta * nPhiBin;
0112       if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Advanced) {
0113         edm::LogVerbatim("PatternRecogntionbyFastJet") << "offset: " << offset;
0114       }
0115       for (int iphi = 0; iphi <= nPhiBin; ++iphi) {
0116         if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Advanced) {
0117           edm::LogVerbatim("PatternRecogntionbyFastJet") << "iphi: " << iphi;
0118           edm::LogVerbatim("PatternRecogntionbyFastJet") << "Entries in tileBin: " << tileOnLayer[offset + iphi].size();
0119         }
0120         for (auto clusterIdx : tileOnLayer[offset + iphi]) {
0121           // Skip masked layer clusters
0122           if (input.mask[clusterIdx] == 0.) {
0123             if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Advanced) {
0124               edm::LogVerbatim("PatternRecogntionbyFastJet") << "Skipping masked layerIdx " << clusterIdx;
0125             }
0126             continue;
0127           }
0128           // Should we correct for the position of the PV?
0129           auto const &cl = input.layerClusters[clusterIdx];
0130           math::XYZVector direction(cl.x(), cl.y(), cl.z());
0131           direction = direction.Unit();
0132           direction *= cl.energy();
0133           auto fpj = fastjet::PseudoJet(direction.X(), direction.Y(), direction.Z(), cl.energy());
0134           fpj.set_user_index(clusterIdx);
0135           fjInputs.push_back(fpj);
0136         }  // End of loop on the clusters on currentLayer
0137       }    // End of loop over phi-bin region
0138     }      // End of loop over eta-bin region
0139   }        // End of loop over layers
0140 
0141   // Collect the jet from the other side wrt to the one taken care of inside the main loop above.
0142   buildJetAndTracksters(fjInputs, result);
0143 
0144   ticl::assignPCAtoTracksters(result,
0145                               input.layerClusters,
0146                               input.layerClustersTime,
0147                               rhtools_.getPositionLayer(rhtools_.lastLayerEE(isHFnose), isHFnose).z());
0148 
0149   // run energy regression and ID
0150   energyRegressionAndID(input.layerClusters, input.tfSession, result);
0151   if (PatternRecognitionAlgoBaseT<TILES>::algo_verbosity_ > VerbosityLevel::Basic) {
0152     for (auto const &t : result) {
0153       edm::LogVerbatim("PatternRecogntionbyFastJet") << "Barycenter: " << t.barycenter();
0154       edm::LogVerbatim("PatternRecogntionbyFastJet") << "LCs: " << t.vertices().size();
0155       edm::LogVerbatim("PatternRecogntionbyFastJet") << "Energy: " << t.raw_energy();
0156       edm::LogVerbatim("PatternRecogntionbyFastJet") << "Regressed: " << t.regressed_energy();
0157     }
0158   }
0159 }
0160 
0161 template <typename TILES>
0162 void PatternRecognitionbyFastJet<TILES>::energyRegressionAndID(const std::vector<reco::CaloCluster> &layerClusters,
0163                                                                const tensorflow::Session *eidSession,
0164                                                                std::vector<Trackster> &tracksters) {
0165   // Energy regression and particle identification strategy:
0166   //
0167   // 1. Set default values for regressed energy and particle id for each trackster.
0168   // 2. Store indices of tracksters whose total sum of cluster energies is above the
0169   //    eidMinClusterEnergy_ (GeV) treshold. Inference is not applied for soft tracksters.
0170   // 3. When no trackster passes the selection, return.
0171   // 4. Create input and output tensors. The batch dimension is determined by the number of
0172   //    selected tracksters.
0173   // 5. Fill input tensors with layer cluster features. Per layer, clusters are ordered descending
0174   //    by energy. Given that tensor data is contiguous in memory, we can use pointer arithmetic to
0175   //    fill values, even with batching.
0176   // 6. Zero-fill features for empty clusters in each layer.
0177   // 7. Batched inference.
0178   // 8. Assign the regressed energy and id probabilities to each trackster.
0179   //
0180   // Indices used throughout this method:
0181   // i -> batch element / trackster
0182   // j -> layer
0183   // k -> cluster
0184   // l -> feature
0185 
0186   // set default values per trackster, determine if the cluster energy threshold is passed,
0187   // and store indices of hard tracksters
0188   std::vector<int> tracksterIndices;
0189   for (int i = 0; i < static_cast<int>(tracksters.size()); i++) {
0190     // calculate the cluster energy sum (2)
0191     // note: after the loop, sumClusterEnergy might be just above the threshold which is enough to
0192     // decide whether to run inference for the trackster or not
0193     float sumClusterEnergy = 0.;
0194     for (const unsigned int &vertex : tracksters[i].vertices()) {
0195       sumClusterEnergy += static_cast<float>(layerClusters[vertex].energy());
0196       // there might be many clusters, so try to stop early
0197       if (sumClusterEnergy >= eidMinClusterEnergy_) {
0198         // set default values (1)
0199         tracksters[i].setRegressedEnergy(0.f);
0200         tracksters[i].zeroProbabilities();
0201         tracksterIndices.push_back(i);
0202         break;
0203       }
0204     }
0205   }
0206 
0207   // do nothing when no trackster passes the selection (3)
0208   int batchSize = static_cast<int>(tracksterIndices.size());
0209   if (batchSize == 0) {
0210     return;
0211   }
0212 
0213   // create input and output tensors (4)
0214   tensorflow::TensorShape shape({batchSize, eidNLayers_, eidNClusters_, eidNFeatures_});
0215   tensorflow::Tensor input(tensorflow::DT_FLOAT, shape);
0216   tensorflow::NamedTensorList inputList = {{eidInputName_, input}};
0217 
0218   std::vector<tensorflow::Tensor> outputs;
0219   std::vector<std::string> outputNames;
0220   if (!eidOutputNameEnergy_.empty()) {
0221     outputNames.push_back(eidOutputNameEnergy_);
0222   }
0223   if (!eidOutputNameId_.empty()) {
0224     outputNames.push_back(eidOutputNameId_);
0225   }
0226 
0227   // fill input tensor (5)
0228   for (int i = 0; i < batchSize; i++) {
0229     const Trackster &trackster = tracksters[tracksterIndices[i]];
0230 
0231     // per layer, we only consider the first eidNClusters_ clusters in terms of energy, so in order
0232     // to avoid creating large / nested structures to do the sorting for an unknown number of total
0233     // clusters, create a sorted list of layer cluster indices to keep track of the filled clusters
0234     std::vector<int> clusterIndices(trackster.vertices().size());
0235     for (int k = 0; k < (int)trackster.vertices().size(); k++) {
0236       clusterIndices[k] = k;
0237     }
0238     sort(clusterIndices.begin(), clusterIndices.end(), [&layerClusters, &trackster](const int &a, const int &b) {
0239       return layerClusters[trackster.vertices(a)].energy() > layerClusters[trackster.vertices(b)].energy();
0240     });
0241 
0242     // keep track of the number of seen clusters per layer
0243     std::vector<int> seenClusters(eidNLayers_);
0244 
0245     // loop through clusters by descending energy
0246     for (const int &k : clusterIndices) {
0247       // get features per layer and cluster and store the values directly in the input tensor
0248       const reco::CaloCluster &cluster = layerClusters[trackster.vertices(k)];
0249       int j = rhtools_.getLayerWithOffset(cluster.hitsAndFractions()[0].first) - 1;
0250       if (j < eidNLayers_ && seenClusters[j] < eidNClusters_) {
0251         // get the pointer to the first feature value for the current batch, layer and cluster
0252         float *features = &input.tensor<float, 4>()(i, j, seenClusters[j], 0);
0253 
0254         // fill features
0255         *(features++) = float(cluster.energy() / float(trackster.vertex_multiplicity(k)));
0256         *(features++) = float(std::abs(cluster.eta()));
0257         *(features) = float(cluster.phi());
0258 
0259         // increment seen clusters
0260         seenClusters[j]++;
0261       }
0262     }
0263 
0264     // zero-fill features of empty clusters in each layer (6)
0265     for (int j = 0; j < eidNLayers_; j++) {
0266       for (int k = seenClusters[j]; k < eidNClusters_; k++) {
0267         float *features = &input.tensor<float, 4>()(i, j, k, 0);
0268         for (int l = 0; l < eidNFeatures_; l++) {
0269           *(features++) = 0.f;
0270         }
0271       }
0272     }
0273   }
0274 
0275   // run the inference (7)
0276   tensorflow::run(const_cast<tensorflow::Session *>(eidSession), inputList, outputNames, &outputs);
0277 
0278   // store regressed energy per trackster (8)
0279   if (!eidOutputNameEnergy_.empty()) {
0280     // get the pointer to the energy tensor, dimension is batch x 1
0281     float *energy = outputs[0].flat<float>().data();
0282 
0283     for (const int &i : tracksterIndices) {
0284       tracksters[i].setRegressedEnergy(*(energy++));
0285     }
0286   }
0287 
0288   // store id probabilities per trackster (8)
0289   if (!eidOutputNameId_.empty()) {
0290     // get the pointer to the id probability tensor, dimension is batch x id_probabilities.size()
0291     int probsIdx = eidOutputNameEnergy_.empty() ? 0 : 1;
0292     float *probs = outputs[probsIdx].flat<float>().data();
0293 
0294     for (const int &i : tracksterIndices) {
0295       tracksters[i].setProbabilities(probs);
0296       probs += tracksters[i].id_probabilities().size();
0297     }
0298   }
0299 }
0300 
0301 template <typename TILES>
0302 void PatternRecognitionbyFastJet<TILES>::fillPSetDescription(edm::ParameterSetDescription &iDesc) {
0303   iDesc.add<int>("algo_verbosity", 0);
0304   iDesc.add<double>("antikt_radius", 0.09)->setComment("Radius to be used while running the Anti-kt clustering");
0305   iDesc.add<int>("minNumLayerCluster", 5)->setComment("Not Inclusive");
0306   iDesc.add<std::string>("eid_input_name", "input");
0307   iDesc.add<std::string>("eid_output_name_energy", "output/regressed_energy");
0308   iDesc.add<std::string>("eid_output_name_id", "output/id_probabilities");
0309   iDesc.add<double>("eid_min_cluster_energy", 1.);
0310   iDesc.add<int>("eid_n_layers", 50);
0311   iDesc.add<int>("eid_n_clusters", 10);
0312 }
0313 
0314 template class ticl::PatternRecognitionbyFastJet<TICLLayerTiles>;