EgammaTools/plugins/DRNCorrectionProducerT.cc

0001 #include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
0002 #include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
0003
0004 #include "FWCore/Framework/interface/Frameworkfwd.h"
0005 #include "FWCore/Framework/interface/stream/EDProducer.h"
0006 #include "FWCore/Framework/interface/Event.h"
0007 #include "FWCore/Framework/interface/EventSetup.h"
0008 #include "FWCore/Framework/interface/MakerMacros.h"
0009
0010 #include "FWCore/ParameterSet/interface/ParameterSet.h"
0011
0012 #include "DataFormats/PatCandidates/interface/Photon.h"
0013 #include "DataFormats/PatCandidates/interface/Electron.h"
0014 #include "DataFormats/EgammaCandidates/interface/Photon.h"
0015 #include "DataFormats/EgammaCandidates/interface/GsfElectron.h"
0016 #include "DataFormats/EgammaCandidates/interface/PhotonFwd.h"
0017 #include "DataFormats/EgammaCandidates/interface/PhotonCore.h"
0018 #include "DataFormats/Common/interface/ValueMap.h"
0019 #include "DataFormats/Common/interface/Handle.h"
0020 #include "DataFormats/EcalDetId/interface/EcalSubdetector.h"
0021 #include "DataFormats/EcalDetId/interface/EBDetId.h"
0022 #include "DataFormats/EcalDetId/interface/EEDetId.h"
0023 #include "DataFormats/EcalDetId/interface/ESDetId.h"
0024 #include "DataFormats/EcalRecHit/interface/EcalRecHit.h"
0025 #include "DataFormats/EcalRecHit/interface/EcalRecHitCollections.h"
0026 #include "DataFormats/EgammaReco/interface/SuperClusterFwd.h"
0027 #include "DataFormats/EgammaReco/interface/SuperCluster.h"
0028
0029 #include "Geometry/CaloGeometry/interface/CaloSubdetectorGeometry.h"
0030 #include "Geometry/CaloGeometry/interface/CaloGeometry.h"
0031 #include "Geometry/CaloGeometry/interface/CaloCellGeometry.h"
0032 #include "Geometry/CaloGeometry/interface/TruncatedPyramid.h"
0033 #include "Geometry/EcalAlgo/interface/EcalPreshowerGeometry.h"
0034 #include "Geometry/CaloTopology/interface/EcalPreshowerTopology.h"
0035 #include "Geometry/Records/interface/CaloGeometryRecord.h"
0036 #include "Geometry/CaloTopology/interface/CaloTopology.h"
0037
0038 #include "CondFormats/EcalObjects/interface/EcalPedestals.h"
0039 #include "CondFormats/DataRecord/interface/EcalPedestalsRcd.h"
0040
0041 #include "RecoEcal/EgammaCoreTools/interface/PositionCalc.h"
0042
0043 #include <cmath>
0044
0045 /*
0046  * DRNCorrectionProducerT
0047  *
0048  * Producer template generate a ValueMap of corrected energies and resolutions
0049  * ValueMap contains std::pair<float, float> of corrected energy, resolution
0050  *
0051  * Author: Simon Rothman (MIT)
0052  * Written 2022
0053  *
0054  */
0055
0056 namespace {
0057   float sigmoid(float x) { return 1.0f / (1.0f + std::exp(-x)); }
0058
0059   float logcorrection(float x) {
0060     static float ln2 = std::log(2);
0061     return ln2 * 2 * (sigmoid(x) - 0.5);
0062   }
0063
0064   //correction factor is transformed by sigmoid and "logratioflip target"
0065   float correction(float x) { return std::exp(-logcorrection(x)); }
0066
0067   inline float rescale(float x, float min, float range) { return (x - min) / range; }
0068
0069   //resolution is transformed by softplus function
0070   float resolution(float x) { return std::log(1 + std::exp(x)); }
0071
0072   const float RHO_MIN = 0.0f;
0073   const float RHO_RANGE = 13.0f;
0074
0075   const float HOE_MIN = 0.0f;
0076   const float HOE_RANGE = 0.05f;
0077
0078   const float XY_MIN = -150.0f;
0079   const float XY_RANGE = 300.0f;
0080
0081   const float Z_MIN = -330.0f;
0082   const float Z_RANGE = 660.0f;
0083
0084   const float NOISE_MIN = 0.9f;
0085   const float NOISE_RANGE = 3.0f;
0086
0087   const float ECAL_MIN = 0.0f;
0088   const float ECAL_RANGE = 250.0f;
0089
0090   const float ES_MIN = 0.0f;
0091   const float ES_RANGE = 0.1f;
0092
0093 }  // namespace
0094
0095 template <typename T>
0096 class DRNCorrectionProducerT : public TritonEDProducer<> {
0097 public:
0098   explicit DRNCorrectionProducerT(const edm::ParameterSet& iConfig);
0099
0100   static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
0101
0102   void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& input) override;
0103   void produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const& iOutput) override;
0104
0105 private:
0106   edm::EDGetTokenT<edm::View<T>> particleToken_;
0107
0108   edm::EDGetTokenT<double> rhoToken_;
0109
0110   edm::EDGetTokenT<EcalRecHitCollection> EBRecHitsToken_, EERecHitsToken_, ESRecHitsToken_;
0111
0112   edm::ESGetToken<EcalPedestals, EcalPedestalsRcd> pedToken_;
0113   edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geomToken_;
0114
0115   size_t nPart_, nValidPart_;
0116
0117   bool isEB(const T& part);
0118   bool isEE(const T& part);
0119   bool skip(const T& part);
0120 };
0121
0122 template <typename T>
0123 DRNCorrectionProducerT<T>::DRNCorrectionProducerT(const edm::ParameterSet& iConfig)
0124     : TritonEDProducer<>(iConfig),
0125       particleToken_(consumes(iConfig.getParameter<edm::InputTag>("particleSource"))),
0126       rhoToken_(consumes(iConfig.getParameter<edm::InputTag>("rhoName"))),
0127       EBRecHitsToken_(consumes<EcalRecHitCollection>(iConfig.getParameter<edm::InputTag>("reducedEcalRecHitsEB"))),
0128       EERecHitsToken_(consumes<EcalRecHitCollection>(iConfig.getParameter<edm::InputTag>("reducedEcalRecHitsEE"))),
0129       ESRecHitsToken_(consumes<EcalRecHitCollection>(iConfig.getParameter<edm::InputTag>("reducedEcalRecHitsES"))),
0130       pedToken_(esConsumes()),
0131       geomToken_(esConsumes()) {
0132   produces<edm::ValueMap<std::pair<float, float>>>();
0133 }
0134
0135 template <typename T>
0136 bool DRNCorrectionProducerT<T>::isEB(const T& part) {
0137   return part.superCluster()->seed()->hitsAndFractions().at(0).first.subdetId() == EcalBarrel;
0138 }
0139
0140 template <typename T>
0141 bool DRNCorrectionProducerT<T>::isEE(const T& part) {
0142   return part.superCluster()->seed()->hitsAndFractions().at(0).first.subdetId() == EcalEndcap;
0143 }
0144
0145 template <typename T>
0146 bool DRNCorrectionProducerT<T>::skip(const T& part) {
0147   /*
0148    * Separated out from acquire() and produce() to ensure that skipping check is identical in both
0149    * N.B. in MiniAOD there are sometimes particles with no RecHits
0150    * We can not apply our regression to these, so we skip them
0151    */
0152   return (!isEB(part) && !isEE(part)) || part.superCluster()->hitsAndFractions().empty();
0153 }
0154
0155 template <typename T>
0156 void DRNCorrectionProducerT<T>::acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) {
0157   /*
0158    * Get products from event and event setup
0159    */
0160   const auto& particles_ = iEvent.getHandle(particleToken_);
0161   float rho = iEvent.get(rhoToken_);
0162
0163   const auto& ped = &iSetup.getData(pedToken_);
0164   const auto& geo = &iSetup.getData(geomToken_);
0165
0166   const CaloSubdetectorGeometry* ecalEBGeom =
0167       static_cast<const CaloSubdetectorGeometry*>(geo->getSubdetectorGeometry(DetId::Ecal, EcalBarrel));
0168   const CaloSubdetectorGeometry* ecalEEGeom =
0169       static_cast<const CaloSubdetectorGeometry*>(geo->getSubdetectorGeometry(DetId::Ecal, EcalEndcap));
0170   const CaloSubdetectorGeometry* ecalESGeom =
0171       static_cast<const CaloSubdetectorGeometry*>(geo->getSubdetectorGeometry(DetId::Ecal, EcalPreshower));
0172
0173   const auto& recHitsEB = iEvent.get(EBRecHitsToken_);
0174   const auto& recHitsEE = iEvent.get(EERecHitsToken_);
0175   const auto& recHitsES = iEvent.get(ESRecHitsToken_);
0176
0177   nPart_ = particles_->size();
0178
0179   if (nPart_ == 0) {
0180     client_->setBatchSize(0);
0181     return;
0182   } else {
0183     client_->setBatchSize(1);
0184   }
0185
0186   /*
0187    * Determine how many particles, how many RecHits there are in each subdetector
0188    */
0189   unsigned nHitsECAL = 0, nHitsES = 0;
0190   nValidPart_ = 0;
0191   for (auto& part : *particles_) {
0192     const reco::SuperClusterRef& sc = part.superCluster();
0193
0194     if (skip(part))
0195       continue;
0196
0197     nHitsECAL += sc->hitsAndFractions().size();
0198
0199     for (auto iES = sc->preshowerClustersBegin(); iES != sc->preshowerClustersEnd(); ++iES) {
0200       nHitsES += (*iES)->hitsAndFractions().size();
0201     }
0202
0203     ++nValidPart_;
0204   }
0205
0206   /*
0207    * Allocate DRN inputs ({SB} is one of ECAL, ES):
0208    * x{SB}: (x, y, z, energy, [noise]) continuous-valued inputs per RecHit
0209    * f{SB}: (flagVal) integer denoting RecHit flag values
0210    * gainECAL: (gain) integer in (0, 1, 2) denoting gain value
0211    * gx: (rho, H/E) additional high-level features.
0212    * batch{SB}: graph models require explicitely passing the particle index for each vertex
0213    */
0214   auto& inputxECAL = iInput.at("xECAL__0");
0215   inputxECAL.setShape(0, nHitsECAL);
0216   auto dataxECAL = inputxECAL.allocate<float>();
0217   auto& vdataxECAL = (*dataxECAL)[0];
0218
0219   auto& inputfECAL = iInput.at("fECAL__1");
0220   inputfECAL.setShape(0, nHitsECAL);
0221   auto datafECAL = inputfECAL.allocate<int64_t>();
0222   auto& vdatafECAL = (*datafECAL)[0];
0223
0224   auto& inputGainECAL = iInput.at("gain__2");
0225   inputGainECAL.setShape(0, nHitsECAL);
0226   auto dataGainECAL = inputGainECAL.allocate<int64_t>();
0227   auto& vdataGainECAL = (*dataGainECAL)[0];
0228
0229   auto& inputGx = iInput.at("graph_x__5");
0230   inputGx.setShape(0, nValidPart_);
0231   auto dataGx = inputGx.allocate<float>();
0232   auto& vdataGx = (*dataGx)[0];
0233
0234   auto& inputBatchECAL = iInput.at("xECAL_batch__6");
0235   inputBatchECAL.setShape(0, nHitsECAL);
0236   auto dataBatchECAL = inputBatchECAL.allocate<int64_t>();
0237   auto& vdataBatchECAL = (*dataBatchECAL)[0];
0238
0239   auto& inputxES = iInput.at("xES__3");
0240   inputxES.setShape(0, nHitsES);
0241   auto dataxES = inputxES.allocate<float>();
0242   auto& vdataxES = (*dataxES)[0];
0243
0244   auto& inputfES = iInput.at("fES__4");
0245   inputfES.setShape(0, nHitsES);
0246   auto datafES = inputfES.allocate<int64_t>();
0247   auto& vdatafES = (*datafES)[0];
0248
0249   auto& inputBatchES = iInput.at("xES_batch__7");
0250   inputBatchES.setShape(0, nHitsES);
0251   auto dataBatchES = inputBatchES.allocate<int64_t>();
0252   auto& vdataBatchES = (*dataBatchES)[0];
0253
0254   /*
0255    * Fill input tensors by iterating over particles...
0256    */
0257   int64_t partNum = 0;
0258   std::shared_ptr<const CaloCellGeometry> geom;
0259   for (auto& part : *particles_) {
0260     const reco::SuperClusterRef& sc = part.superCluster();
0261
0262     if (skip(part))
0263       continue;
0264
0265     std::vector<std::pair<DetId, float>> hitsAndFractions = sc->hitsAndFractions();
0266     EcalRecHitCollection::const_iterator hit;
0267
0268     //iterate over ECAL hits...
0269     for (const auto& detitr : hitsAndFractions) {
0270       DetId id = detitr.first.rawId();
0271       if (isEB(part)) {
0272         geom = ecalEBGeom->getGeometry(id);
0273         hit = recHitsEB.find(detitr.first);
0274       } else {
0275         geom = ecalEEGeom->getGeometry(id);
0276         hit = recHitsEE.find(detitr.first);
0277       }
0278
0279       //fill xECAL
0280       auto pos = geom->getPosition();
0281       vdataxECAL.push_back(rescale(pos.x(), XY_MIN, XY_RANGE));
0282       vdataxECAL.push_back(rescale(pos.y(), XY_MIN, XY_RANGE));
0283       vdataxECAL.push_back(rescale(pos.z(), Z_MIN, Z_RANGE));
0284       vdataxECAL.push_back(rescale(hit->energy() * detitr.second, ECAL_MIN, ECAL_RANGE));
0285       vdataxECAL.push_back(rescale(ped->find(detitr.first)->rms(1), NOISE_MIN, NOISE_RANGE));
0286
0287       //fill fECAL
0288       int64_t flagVal = 0;
0289       if (hit->checkFlag(EcalRecHit::kGood))
0290         flagVal += 1;
0291       if (hit->checkFlag(EcalRecHit::kOutOfTime))
0292         flagVal += 2;
0293       if (hit->checkFlag(EcalRecHit::kPoorCalib))
0294         flagVal += 4;
0295
0296       vdatafECAL.push_back(flagVal);
0297
0298       //fill gain
0299       int64_t gainVal = 0;
0300       if (hit->checkFlag(EcalRecHit::kHasSwitchToGain6))
0301         gainVal = 1;
0302       else if (hit->checkFlag(EcalRecHit::kHasSwitchToGain1))
0303         gainVal = 0;
0304       else
0305         gainVal = 2;
0306
0307       vdataGainECAL.push_back(gainVal);
0308
0309       //fill batch number
0310       vdataBatchECAL.push_back(partNum);
0311     }  //end iterate over ECAL hits
0312
0313     //iterate over ES clusters...
0314     for (auto iES = sc->preshowerClustersBegin(); iES != sc->preshowerClustersEnd(); ++iES) {
0315       for (const auto& ESitr : (*iES)->hitsAndFractions()) {  //iterate over ES hits
0316         hit = recHitsES.find(ESitr.first);
0317         geom = ecalESGeom->getGeometry(ESitr.first);
0318         auto& pos = geom->getPosition();
0319
0320         //fill xES
0321         vdataxES.push_back(rescale(pos.x(), XY_MIN, XY_RANGE));
0322         vdataxES.push_back(rescale(pos.y(), XY_MIN, XY_RANGE));
0323         vdataxES.push_back(rescale(pos.z(), Z_MIN, Z_RANGE));
0324         vdataxES.push_back(rescale(hit->energy(), ES_MIN, ES_RANGE));
0325
0326         //fill fES
0327         int64_t flagVal = 0;
0328         if (hit->checkFlag(EcalRecHit::kESGood))
0329           flagVal += 1;
0330
0331         vdatafES.push_back(flagVal);
0332
0333         //fill batchES
0334         vdataBatchES.push_back(partNum);
0335       }  //end iterate over ES hits
0336     }    //end iterate over ES clusters
0337
0338     //fill gx
0339     vdataGx.push_back(rescale(rho, RHO_MIN, RHO_RANGE));
0340     vdataGx.push_back(rescale(part.hadronicOverEm(), HOE_MIN, HOE_RANGE));
0341
0342     //increment particle number
0343     ++partNum;
0344   }  // end iterate over particles
0345
0346   /*
0347    * Convert input tensors to server data format
0348    */
0349
0350   inputxECAL.toServer(dataxECAL);
0351   inputfECAL.toServer(datafECAL);
0352   inputGainECAL.toServer(dataGainECAL);
0353   inputBatchECAL.toServer(dataBatchECAL);
0354
0355   inputGx.toServer(dataGx);
0356
0357   inputxES.toServer(dataxES);
0358   inputfES.toServer(datafES);
0359   inputBatchES.toServer(dataBatchES);
0360 }
0361
0362 template <typename T>
0363 void DRNCorrectionProducerT<T>::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const& iOutput) {
0364   const auto& particles_ = iEvent.getHandle(particleToken_);
0365
0366   std::vector<std::pair<float, float>> corrections;
0367   corrections.reserve(nPart_);
0368
0369   //if there are no particles, the fromServer() call will fail
0370   //but we can just put() an empty valueMap
0371   if (nPart_) {
0372     const auto& out = iOutput.at("combined_output__0").fromServer<float>();
0373
0374     unsigned i = 0;
0375     float mu, sigma, Epred, sigmaPred, rawE;
0376     for (unsigned iPart = 0; iPart < nPart_; ++iPart) {
0377       const auto& part = particles_->at(iPart);
0378       if (!skip(part)) {
0379         mu = correction(out[0][0 + 11 * i]);
0380         sigma = resolution(out[0][6 + 11 * i]);
0381         ++i;
0382
0383         rawE = part.superCluster()->rawEnergy();
0384         Epred = mu * rawE;
0385         sigmaPred = sigma * rawE;
0386         corrections.emplace_back(Epred, sigmaPred);
0387       } else {
0388         corrections.emplace_back(-1.0f, -1.0f);
0389       }
0390     }
0391   }
0392
0393   //fill
0394   auto out = std::make_unique<edm::ValueMap<std::pair<float, float>>>();
0395   edm::ValueMap<std::pair<float, float>>::Filler filler(*out);
0396   filler.insert(particles_, corrections.begin(), corrections.end());
0397   filler.fill();
0398
0399   iEvent.put(std::move(out));
0400 }
0401
0402 template <typename T>
0403 void DRNCorrectionProducerT<T>::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
0404   edm::ParameterSetDescription desc;
0405   TritonClient::fillPSetDescription(desc);
0406   desc.add<edm::InputTag>("particleSource");
0407   desc.add<edm::InputTag>("rhoName");
0408   desc.add<edm::InputTag>("reducedEcalRecHitsEB", edm::InputTag("reducedEcalRecHitsEB"));
0409   desc.add<edm::InputTag>("reducedEcalRecHitsEE", edm::InputTag("reducedEcalRecHitsEE"));
0410   desc.add<edm::InputTag>("reducedEcalRecHitsES", edm::InputTag("reducedEcalRecHitsES"));
0411   descriptions.addWithDefaultLabel(desc);
0412 }
0413
0414 //reco:: template instances are supported
0415 //uncomment the lines below to enable them
0416
0417 //using GsfElectronDRNCorrectionProducer = DRNCorrectionProducerT<reco::GsfElectron>;
0418 //using GedPhotonDRNCorrectionProducer = DRNCorrectionProducerT<reco::Photon>;
0419 //DEFINE_FWK_MODULE(GedPhotonDRNCorrectionProducer);
0420 //DEFINE_FWK_MODULE(GsfElectronDRNCorrectionProducer);
0421
0422 using PatElectronDRNCorrectionProducer = DRNCorrectionProducerT<pat::Electron>;
0423 using PatPhotonDRNCorrectionProducer = DRNCorrectionProducerT<pat::Photon>;
0424
0425 DEFINE_FWK_MODULE(PatPhotonDRNCorrectionProducer);
0426 DEFINE_FWK_MODULE(PatElectronDRNCorrectionProducer);