TICL/plugins/EGammaSuperclusterProducer.cc

0001 // Authors : Theo Cuisset <theo.cuisset@cern.ch>, Shamik Ghosh <shamik.ghosh@cern.ch>
0002 // Date : 01/2024
0003 /*
0004 Translates TICL superclusters to ECAL supercluster dataformats (reco::SuperCluster and reco::CaloCluster).
0005 Performs similar task as RecoEcal/EgammaCLusterProducers/PFECALSuperClusterProducer
0006 Note that all tracksters are translated to reco::CaloCluster, even those that are not in any SuperCluster
0007 */
0008
0009 #include "FWCore/Framework/interface/stream/EDProducer.h"
0010 #include "FWCore/Framework/interface/Event.h"
0011 #include "FWCore/Framework/interface/EventSetup.h"
0012 #include "FWCore/Framework/interface/ConsumesCollector.h"
0013 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0014 #include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
0015 #include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
0016 #include "FWCore/Utilities/interface/EDGetToken.h"
0017 #include "FWCore/Utilities/interface/FileInPath.h"
0018
0019 #include <vector>
0020 #include <array>
0021 #include <limits>
0022 #include <algorithm>
0023
0024 #include "PhysicsTools/ONNXRuntime/interface/ONNXRuntime.h"
0025
0026 #include "DataFormats/EgammaReco/interface/SuperClusterFwd.h"
0027 #include "DataFormats/EgammaReco/interface/SuperCluster.h"
0028 #include "DataFormats/EgammaReco/interface/ElectronSeed.h"
0029 #include "DataFormats/EgammaReco/interface/ElectronSeedFwd.h"
0030 #include "DataFormats/HGCalReco/interface/Trackster.h"
0031 #include "DataFormats/Math/interface/deltaPhi.h"
0032
0033 using cms::Ort::ONNXRuntime;
0034
0035 class EGammaSuperclusterProducer : public edm::stream::EDProducer<edm::GlobalCache<ONNXRuntime>> {
0036 public:
0037   EGammaSuperclusterProducer(const edm::ParameterSet&, const ONNXRuntime*);
0038
0039   void produce(edm::Event&, const edm::EventSetup&) override;
0040
0041   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0042
0043   static std::unique_ptr<ONNXRuntime> initializeGlobalCache(const edm::ParameterSet& iConfig);
0044   static void globalEndJob(const ONNXRuntime*);
0045
0046 private:
0047   edm::EDGetTokenT<ticl::TracksterCollection> ticlSuperClustersToken_;
0048   edm::EDGetTokenT<std::vector<std::vector<unsigned int>>> superClusterLinksToken_;
0049   edm::EDGetTokenT<ticl::TracksterCollection> ticlTrackstersEMToken_;
0050   edm::EDGetTokenT<reco::CaloClusterCollection> layerClustersToken_;
0051   float superclusterEtThreshold_;
0052   bool enableRegression_;
0053 };
0054
0055 EGammaSuperclusterProducer::EGammaSuperclusterProducer(const edm::ParameterSet& ps, const ONNXRuntime*)
0056     : ticlSuperClustersToken_(consumes<ticl::TracksterCollection>(ps.getParameter<edm::InputTag>("ticlSuperClusters"))),
0057       superClusterLinksToken_(consumes<std::vector<std::vector<unsigned int>>>(
0058           edm::InputTag(ps.getParameter<edm::InputTag>("ticlSuperClusters").label(),
0059                         "linkedTracksterIdToInputTracksterId",
0060                         ps.getParameter<edm::InputTag>("ticlSuperClusters").process()))),
0061       ticlTrackstersEMToken_(consumes<ticl::TracksterCollection>(ps.getParameter<edm::InputTag>("ticlTrackstersEM"))),
0062       layerClustersToken_(consumes<reco::CaloClusterCollection>(ps.getParameter<edm::InputTag>("layerClusters"))),
0063       superclusterEtThreshold_(ps.getParameter<double>("superclusterEtThreshold")),
0064       enableRegression_(ps.getParameter<bool>("enableRegression")) {
0065   produces<reco::SuperClusterCollection>();
0066   produces<reco::CaloClusterCollection>();  // The CaloCluster corresponding to each EM trackster
0067 }
0068
0069 void EGammaSuperclusterProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
0070   auto const& ticlSuperclusters = iEvent.get(ticlSuperClustersToken_);
0071   auto const& ticlSuperclusterLinks = iEvent.get(superClusterLinksToken_);
0072   auto emTracksters_h = iEvent.getHandle(ticlTrackstersEMToken_);
0073   auto const& emTracksters = *emTracksters_h;
0074   auto const& layerClusters = iEvent.get(layerClustersToken_);
0075   // Output collections :
0076   auto egammaSuperclusters = std::make_unique<reco::SuperClusterCollection>();
0077   auto caloClustersEM = std::make_unique<reco::CaloClusterCollection>();
0078
0079   // Fill reco::CaloCluster collection (1-1 mapping to TICL EM tracksters)
0080   for (ticl::Trackster const& emTrackster : emTracksters) {
0081     std::vector<std::pair<DetId, float>> hitsAndFractions;
0082     int iLC = 0;
0083     std::for_each(std::begin(emTrackster.vertices()), std::end(emTrackster.vertices()), [&](unsigned int lcId) {
0084       const auto fraction = 1.f / emTrackster.vertex_multiplicity(iLC++);
0085       for (const auto& cell : layerClusters[lcId].hitsAndFractions()) {
0086         hitsAndFractions.emplace_back(cell.first, cell.second * fraction);
0087       }
0088     });
0089
0090     reco::CaloCluster& caloCluster = caloClustersEM->emplace_back(
0091         emTrackster.raw_energy(),                  // energy
0092         math::XYZPoint(emTrackster.barycenter()),  // position
0093         reco::CaloID(reco::CaloID::DET_HGCAL_ENDCAP),
0094         hitsAndFractions,
0095         reco::CaloCluster::particleFlow,  // algoID (copying from output of PFECALCSuperClusterProducer)
0096         hitsAndFractions.at(0)
0097             .first  // seedId (this may need to be updated once a common definition of the seed of a cluster is adopted for Phase-2)
0098     );
0099     caloCluster.setCorrectedEnergy(emTrackster.raw_energy());  // Needs to be updated with new supercluster regression
0100   }
0101
0102   edm::OrphanHandle<reco::CaloClusterCollection> caloClustersEM_h = iEvent.put(std::move(caloClustersEM));
0103
0104   // Fill reco::SuperCluster collection and prepare regression inputs
0105   assert(ticlSuperclusters.size() == ticlSuperclusterLinks.size());
0106   const unsigned int regressionFeatureCount = 8;
0107   std::vector<float> regressionInputs;
0108   regressionInputs.reserve(ticlSuperclusters.size() * regressionFeatureCount);
0109   unsigned int superclustersPassingSelectionsCount = 0;
0110   for (std::size_t sc_i = 0; sc_i < ticlSuperclusters.size(); sc_i++) {
0111     ticl::Trackster const& ticlSupercluster = ticlSuperclusters[sc_i];
0112     if (ticlSupercluster.raw_pt() < superclusterEtThreshold_)
0113       continue;
0114     std::vector<unsigned int> const& superclusterLink = ticlSuperclusterLinks[sc_i];
0115     assert(!superclusterLink.empty());
0116
0117     reco::CaloClusterPtrVector trackstersEMInSupercluster;
0118     float max_eta = std::numeric_limits<float>::min();
0119     float max_phi = std::numeric_limits<float>::min();
0120     float min_eta = std::numeric_limits<float>::max();
0121     float min_phi = std::numeric_limits<float>::max();
0122     for (unsigned int tsInSc_id : superclusterLink) {
0123       trackstersEMInSupercluster.push_back(reco::CaloClusterPtr(caloClustersEM_h, tsInSc_id));
0124       ticl::Trackster const& constituentTrackster = emTracksters[tsInSc_id];
0125       max_eta = std::max(constituentTrackster.barycenter().eta(), max_eta);
0126       max_phi = std::max(constituentTrackster.barycenter().phi(), max_phi);
0127       min_eta = std::min(constituentTrackster.barycenter().eta(), min_eta);
0128       min_phi = std::min(constituentTrackster.barycenter().phi(), min_phi);
0129     }
0130     egammaSuperclusters->emplace_back(
0131         ticlSupercluster.raw_energy(),
0132         reco::SuperCluster::Point(ticlSupercluster.barycenter()),
0133         reco::CaloClusterPtr(caloClustersEM_h,
0134                              superclusterLink[0]),  // seed (first trackster in superclusterLink is the seed)
0135         trackstersEMInSupercluster,                 // clusters
0136         0.,                                         // Epreshower (not relevant for HGCAL)
0137         -1,                                         // phiwidth (not implemented yet)
0138         -1                                          // etawidth (not implemented yet)
0139     );
0140     superclustersPassingSelectionsCount++;
0141
0142     if (enableRegression_) {
0143       regressionInputs.insert(
0144           regressionInputs.end(),
0145           {ticlSupercluster.barycenter().eta(),
0146            ticlSupercluster.barycenter().phi(),
0147            ticlSupercluster.raw_energy(),
0148            std::abs(max_eta - min_eta),
0149            max_phi - min_phi > M_PI ? 2 * static_cast<float>(M_PI) - (max_phi - min_phi) : max_phi - min_phi,
0150            emTracksters[superclusterLink[0]].raw_energy() -
0151                (superclusterLink.size() >= 2 ? emTracksters[superclusterLink.back()].raw_energy() : 0.f),
0152            emTracksters[superclusterLink[0]].raw_energy() / ticlSupercluster.raw_energy(),
0153            static_cast<float>(superclusterLink.size())});
0154     }
0155   }
0156
0157   if (enableRegression_ && superclustersPassingSelectionsCount > 0) {
0158     // Run the regression
0159     // ONNXRuntime takes std::vector<std::vector<float>>& as input (non-const reference) so we have to make a new vector
0160     std::vector<std::vector<float>> inputs_for_onnx{{std::move(regressionInputs)}};
0161     std::vector<float> outputs =
0162         globalCache()->run({"input"}, inputs_for_onnx, {}, {}, superclustersPassingSelectionsCount)[0];
0163
0164     assert(egammaSuperclusters->size() == outputs.size());
0165     for (std::size_t sc_i = 0; sc_i < egammaSuperclusters->size(); sc_i++) {
0166       (*egammaSuperclusters)[sc_i].setCorrectedEnergy(outputs[sc_i]);
0167       // correctedEnergyUncertainty is left at its default value
0168     }
0169   }
0170
0171   iEvent.put(std::move(egammaSuperclusters));
0172 }
0173
0174 std::unique_ptr<ONNXRuntime> EGammaSuperclusterProducer::initializeGlobalCache(const edm::ParameterSet& iConfig) {
0175   if (iConfig.getParameter<bool>("enableRegression"))
0176     return std::make_unique<ONNXRuntime>(iConfig.getParameter<edm::FileInPath>("regressionModelPath").fullPath());
0177   else
0178     return std::unique_ptr<ONNXRuntime>(nullptr);
0179 }
0180
0181 void EGammaSuperclusterProducer::globalEndJob(const ONNXRuntime*) {}
0182
0183 void EGammaSuperclusterProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0184   edm::ParameterSetDescription desc;
0185   desc.add<edm::InputTag>("ticlSuperClusters", edm::InputTag("ticlTracksterLinksSuperclusteringDNN"));
0186   desc.add<edm::InputTag>("ticlTrackstersEM", edm::InputTag("ticlTrackstersCLUE3DHigh"))
0187       ->setComment("The trackster collection used before superclustering, ie CLUE3D EM tracksters");
0188   desc.add<edm::InputTag>("layerClusters", edm::InputTag("hgcalMergeLayerClusters"))
0189       ->setComment("The layer cluster collection that goes with ticlTrackstersEM");
0190   desc.add<double>("superclusterEtThreshold", 4.)->setComment("Minimum supercluster transverse energy.");
0191   desc.add<bool>("enableRegression", true)->setComment("Enable supercluster energy regression");
0192   desc.add<edm::FileInPath>("regressionModelPath",
0193                             edm::FileInPath("RecoHGCal/TICL/data/superclustering/regression_v1.onnx"))
0194       ->setComment("Path to regression network (as ONNX model)");
0195
0196   descriptions.add("ticlEGammaSuperClusterProducer", desc);
0197 }
0198
0199 #include "FWCore/Framework/interface/MakerMacros.h"
0200 DEFINE_FWK_MODULE(EGammaSuperclusterProducer);