Dynamic Reduction Network-based electron energy regression using the SONIC service

RecoEcal/EgammaClusterAlgos/BuildFile.xml

@@ -14,6 +14,7 @@
 <use name="Geometry/CaloTopology"/>
 <use name="RecoEgamma/EgammaTools"/>
 <use name="RecoParticleFlow/PFClusterTools"/>
+<use name="HeterogeneousCore/SonicTriton"/>
 <use name="clhep"/>
 <export>
   <lib name="1"/>

RecoEcal/EgammaClusterAlgos/interface/SCEnergyCorrectorDRN.h

@@ -0,0 +1,90 @@
+//--------------------------------------------------------------------------------------------------
+//
+// SCEnergyCorrectorDRN
+//
+// Helper Class for applying regression-based energy corrections with DRN implimentation
+//
+// Based on RecoEcal/EgammaClusterAlgos/SCEnergyCorrectorSemiParm
+//
+// Author: Simon Rothman (MIT, UMN)
+//
+//--------------------------------------------------------------------------------------------------
+
+#ifndef RecoEcal_EgammaClusterAlgos_SCEnergyCorrectorDRN_h
+#define RecoEcal_EgammaClusterAlgos_SCEnergyCorrectorDRN_h
+
+#include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
+
+#include "FWCore/Framework/interface/EventSetup.h"
+#include "FWCore/Framework/interface/ESHandle.h"
+#include "FWCore/Framework/interface/ConsumesCollector.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "Geometry/Records/interface/CaloTopologyRecord.h"
+#include "Geometry/CaloTopology/interface/CaloTopology.h"
+#include "Geometry/Records/interface/CaloGeometryRecord.h"
+#include "Geometry/CaloGeometry/interface/CaloGeometry.h"
+
+#include "DataFormats/EgammaReco/interface/SuperCluster.h"
+#include "DataFormats/EcalRecHit/interface/EcalRecHit.h"
+#include "DataFormats/EcalRecHit/interface/EcalRecHitCollections.h"
+#include "DataFormats/EcalDetId/interface/EcalSubdetector.h"
+#include "DataFormats/ParticleFlowReco/interface/PFRecHit.h"
+#include "DataFormats/EgammaReco/interface/SuperClusterFwd.h"
+#include "DataFormats/VertexReco/interface/VertexFwd.h"
+
+#include "CondFormats/GBRForest/interface/GBRForestD.h"
+#include "CondFormats/DataRecord/interface/GBRDWrapperRcd.h"
+
+#include "RecoEgamma/EgammaTools/interface/EgammaBDTOutputTransformer.h"
+#include "RecoEgamma/EgammaTools/interface/HGCalShowerShapeHelper.h"
+
+#include <sstream>
+#include <string>
+#include <vector>
+#include <random>
+
+class SCEnergyCorrectorDRN {
+public:
+  SCEnergyCorrectorDRN();
+  //if you want override the default on where conditions are consumed, you need to use
+  //the other constructor and then call setTokens approprately
+  SCEnergyCorrectorDRN(const edm::ParameterSet& iConfig, edm::ConsumesCollector cc);
+
+  static void fillPSetDescription(edm::ParameterSetDescription& desc);
+  static edm::ParameterSetDescription makePSetDescription();
+
+  template <edm::Transition tr = edm::Transition::BeginLuminosityBlock>
+  void setTokens(const edm::ParameterSet& iConfig, edm::ConsumesCollector cc);
+
+  void setEventSetup(const edm::EventSetup& es);
+  void setEvent(const edm::Event& e);
+
+  void makeInput(const edm::Event& iEvent, TritonInputMap& iInput, const reco::SuperClusterCollection& inputSCs) const;
+  TritonOutput<float> getOutput(const TritonOutputMap& iOutput);
+
+private:
+  const CaloTopology* caloTopo_;
+  const CaloGeometry* caloGeom_;
+  edm::ESGetToken<CaloTopology, CaloTopologyRecord> caloTopoToken_;
+  edm::ESGetToken<CaloGeometry, CaloGeometryRecord> caloGeomToken_;
+
+  edm::EDGetTokenT<EcalRecHitCollection> tokenEBRecHits_;
+  edm::EDGetTokenT<EcalRecHitCollection> tokenEERecHits_;
+  edm::EDGetTokenT<double> rhoToken_;
+
+  edm::Handle<EcalRecHitCollection> recHitsEB_;
+  edm::Handle<EcalRecHitCollection> recHitsEE_;
+
+  edm::Handle<double> rhoHandle_;
+};
+
+template <edm::Transition esTransition>
+void SCEnergyCorrectorDRN::setTokens(const edm::ParameterSet& iConfig, edm::ConsumesCollector cc) {
+  tokenEBRecHits_ = cc.consumes<EcalRecHitCollection>(iConfig.getParameter<edm::InputTag>("ecalRecHitsEB"));
+  tokenEERecHits_ = cc.consumes<EcalRecHitCollection>(iConfig.getParameter<edm::InputTag>("ecalRecHitsEE"));
+  caloGeomToken_ = cc.esConsumes<CaloGeometry, CaloGeometryRecord, esTransition>();
+  caloTopoToken_ = cc.esConsumes<CaloTopology, CaloTopologyRecord, esTransition>();
+  rhoToken_ = cc.consumes<double>(iConfig.getParameter<edm::InputTag>("rhoFastJet"));
+}
+#endif

RecoEcal/EgammaClusterAlgos/src/SCEnergyCorrectorDRN.cc

@@ -0,0 +1,122 @@
+#include "RecoEcal/EgammaClusterAlgos/interface/SCEnergyCorrectorDRN.h"
+
+#include "FWCore/Utilities/interface/isFinite.h"
+#include "FWCore/Utilities/interface/Transition.h"
+#include "DataFormats/EcalDetId/interface/EcalSubdetector.h"
+#include "DataFormats/VertexReco/interface/Vertex.h"
+#include "DataFormats/Math/interface/deltaPhi.h"
+#include "RecoEcal/EgammaCoreTools/interface/EcalTools.h"
+#include "RecoEcal/EgammaCoreTools/interface/EcalClusterTools.h"
+#include "RecoEgamma/EgammaTools/interface/EgammaHGCALIDParamDefaults.h"
+
+#include <vdt/vdtMath.h>
+
+static const float RHO_MAX = 15.0f;
+static const float X_MAX = 150.0f;
+static const float X_RANGE = 300.0f;
+static const float Y_MAX = 150.0f;
+static const float Y_RANGE = 300.0f;
+static const float Z_MAX = 330.0f;
+static const float Z_RANGE = 660.0f;
+static const float E_RANGE = 250.0f;
+
+SCEnergyCorrectorDRN::SCEnergyCorrectorDRN() : caloTopo_(nullptr), caloGeom_(nullptr) {}
+
+SCEnergyCorrectorDRN::SCEnergyCorrectorDRN(const edm::ParameterSet& iConfig, edm::ConsumesCollector cc)
+    : SCEnergyCorrectorDRN() {
+  setTokens(iConfig, cc);
+}
+
+void SCEnergyCorrectorDRN::fillPSetDescription(edm::ParameterSetDescription& desc) {
+  desc.add<edm::InputTag>("ecalRecHitsEE", edm::InputTag("ecalRecHit", "reducedEcalRecHitsEE"));
+  desc.add<edm::InputTag>("ecalRecHitsEB", edm::InputTag("ecalRecHit", "reducedEcalRecHitsEB"));
+  desc.add<edm::InputTag>("rhoFastJet", edm::InputTag("fixedGridRhoAll"));
+}
+
+edm::ParameterSetDescription SCEnergyCorrectorDRN::makePSetDescription() {
+  edm::ParameterSetDescription desc;
+  fillPSetDescription(desc);
+  return desc;
+}
+
+void SCEnergyCorrectorDRN::setEventSetup(const edm::EventSetup& es) {
+  caloTopo_ = &es.getData(caloTopoToken_);
+  caloGeom_ = &es.getData(caloGeomToken_);
+}
+
+void SCEnergyCorrectorDRN::setEvent(const edm::Event& event) {
+  event.getByToken(tokenEBRecHits_, recHitsEB_);
+  event.getByToken(tokenEERecHits_, recHitsEE_);
+  event.getByToken(rhoToken_, rhoHandle_);
+}
+
+void SCEnergyCorrectorDRN::makeInput(const edm::Event& iEvent,
+                                     TritonInputMap& iInput,
+                                     const reco::SuperClusterCollection& inputSCs) const {
+  std::vector<unsigned> nHits;
+  nHits.reserve(inputSCs.size());
+  unsigned totalHits = 0;
+  unsigned n;
+  for (const auto& inputSC : inputSCs) {
+    n = inputSC.hitsAndFractions().size();
+    totalHits += n;
+    nHits.push_back(n);
+  }
+
+  //set shapes
+  auto& input1 = iInput.at("x__0");
+  input1.setShape(0, totalHits);
+  auto data1 = input1.allocate<float>();
+  auto& vdata1 = (*data1)[0];
+
+  auto& input2 = iInput.at("batch__1");
+  input2.setShape(0, totalHits);
+  auto data2 = input2.allocate<int64_t>();
+  auto& vdata2 = (*data2)[0];
+
+  auto& input3 = iInput.at("graphx__2");
+  input3.setShape(0, 2 * nHits.size());
+  auto data3 = input3.allocate<float>();
+  auto& vdata3 = (*data3)[0];
+
+  //fill
+  unsigned batchNum = 0;
+  float En, frac, x, y, z;
+  for (const auto& inputSC : inputSCs) {
+    const auto& hits = inputSC.hitsAndFractions();
+    const bool isEB = hits[0].first.subdetId() == EcalBarrel;
+    const auto& recHitsProduct = isEB ? recHitsEB_.product() : recHitsEE_.product();
+    for (const auto& hit : hits) {
+      En = EcalClusterTools::recHitEnergy(hit.first, recHitsProduct);
+      frac = hit.second;
+      GlobalPoint position = caloGeom_->getGeometry(hit.first)->getPosition();
+      x = (position.x() + X_MAX) / X_RANGE;
+      y = (position.y() + Y_MAX) / Y_RANGE;
+      z = (position.z() + Z_MAX) / Z_RANGE;
+      vdata1.push_back(x);
+      vdata1.push_back(y);
+      vdata1.push_back(z);
+      vdata1.push_back(En * frac / E_RANGE);
+      //Triton does not currently support batching for pytorch GNNs
+      //We pass batch indices explicitely
+      vdata2.push_back(batchNum);
+    }
+    vdata3.push_back(*rhoHandle_ / RHO_MAX);
+    vdata3.push_back(0.0);
+    ++batchNum;
+  }
+
+  // convert to server format
+  input1.toServer(data1);
+  input2.toServer(data2);
+  input3.toServer(data3);
+}
+
+TritonOutput<float> SCEnergyCorrectorDRN::getOutput(const TritonOutputMap& iOutput) {
+  //check the results
+  const auto& output1 = iOutput.begin()->second;
+  // convert from server format
+  const auto& serverout = output1.fromServer<float>();
+
+  return serverout;
+}

RecoEcal/EgammaClusterProducers/BuildFile.xml

@@ -1,4 +1,5 @@
 <use name="FWCore/Framework"/>
+<use name="HeterogeneousCore/SonicTriton"/>
 <use name="FWCore/ParameterSet"/>
 <use name="DataFormats/EcalDigi"/>
 <use name="DataFormats/EcalRecHit"/>

RecoEcal/EgammaClusterProducers/python/SCEnergyCorrectorDRNProducer_cfi.py

@@ -0,0 +1,26 @@
+import FWCore.ParameterSet.Config as cms
+
+DRNProducerEB = cms.EDProducer('SCEnergyCorrectorDRNProducer',
+   inputSCs = cms.InputTag('particleFlowSuperClusterECAL','particleFlowSuperClusterECALBarrel'),
+   Client = cms.PSet(
+        mode = cms.string("Async"),
+        modelName = cms.string("MustacheEB"),
+        modelConfigPath = cms.FileInPath("RecoEcal/EgammaClusterProducers/data/models/MustacheEB/config.pbtxt"),
+        allowedTries = cms.untracked.uint32(1),
+        timeout = cms.untracked.uint32(10),
+    ),
+ )
+
+
+DRNProducerEE = cms.EDProducer('SCEnergyCorrectorDRNProducer',
+   inputSCs = cms.InputTag('particleFlowSuperClusterECAL','particleFlowSuperClusterECALEndcapWithPreshower'),
+   Client = cms.PSet(
+        mode = cms.string("Async"),
+        modelName = cms.string('MustacheEE'),
+        modelConfigPath = cms.FileInPath("RecoEcal/EgammaClusterProducers/data/models/MustacheEE/config.pbtxt"),
+        allowedTries = cms.untracked.uint32(1),
+        timeout = cms.untracked.uint32(10),
+    ),
+ )
+
+

RecoEcal/EgammaClusterProducers/src/SCEnergyCorrectorDRNProducer.cc

@@ -0,0 +1,113 @@
+#include "HeterogeneousCore/SonicTriton/interface/TritonEDProducer.h"
+#include "HeterogeneousCore/SonicTriton/interface/TritonData.h"
+
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "DataFormats/EgammaReco/interface/SuperCluster.h"
+#include "DataFormats/EgammaReco/interface/SuperClusterFwd.h"
+#include "DataFormats/Common/interface/ValueMap.h"
+
+#include "DataFormats/EgammaCandidates/interface/GsfElectron.h"
+
+#include "RecoEcal/EgammaClusterAlgos/interface/SCEnergyCorrectorDRN.h"
+
+#include <sstream>
+#include <string>
+#include <vector>
+#include <random>
+
+/*
+ * SCEnergyCorrectorDRNProducer
+ *
+ * Simple producer to generate a set of corrected superclusters with the DRN regression
+ * Based on RecoEcal/EgammaClusterProducers/SCEnergyCorrectorProducer by S. Harper (RAL/CERN)
+ *
+ * Author: Simon Rothman (UMN, MIT)
+ *
+ */
+
+namespace {
+  float sigmoid(float x) { return 1.0f / (1.0f + exp(-x)); }
+
+  float logcorrection(float x) {
+    static float ln2 = log(2);
+    return ln2 * 2 * (sigmoid(x) - 0.5);
+  }
+
+  float correction(float x) { return exp(-logcorrection(x)); }
+}  // namespace
+
+class SCEnergyCorrectorDRNProducer : public TritonEDProducer<> {
+public:
+  explicit SCEnergyCorrectorDRNProducer(const edm::ParameterSet& iConfig);
+
+  void beginLuminosityBlock(const edm::LuminosityBlock& iLumi, const edm::EventSetup& iSetup) override;
+
+  void acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& input) override;
+  void produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const& iOutput) override;
+
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+private:
+  SCEnergyCorrectorDRN energyCorrector_;
+  edm::EDGetTokenT<reco::SuperClusterCollection> inputSCToken_;
+};
+
+SCEnergyCorrectorDRNProducer::SCEnergyCorrectorDRNProducer(const edm::ParameterSet& iConfig)
+    : TritonEDProducer<>(iConfig, "SCEnergyCorrectorDRNProducer"),
+      energyCorrector_(iConfig.getParameterSet("correctorCfg"), consumesCollector()),
+      inputSCToken_(consumes<reco::SuperClusterCollection>(iConfig.getParameter<edm::InputTag>("inputSCs"))) {
+  produces<reco::SuperClusterCollection>();
+}
+
+void SCEnergyCorrectorDRNProducer::beginLuminosityBlock(const edm::LuminosityBlock& iLumi,
+                                                        const edm::EventSetup& iSetup) {
+  energyCorrector_.setEventSetup(iSetup);
+}
+
+void SCEnergyCorrectorDRNProducer::acquire(edm::Event const& iEvent, edm::EventSetup const& iSetup, Input& iInput) {
+  const auto& inputSCs = iEvent.get(inputSCToken_);
+
+  if (inputSCs.empty()) {
+    client_->setBatchSize(0);
+    return;
+  } else {
+    client_->setBatchSize(1);
+  }
+
+  energyCorrector_.setEvent(iEvent);
+  energyCorrector_.makeInput(iEvent, iInput, inputSCs);
+}
+
+void SCEnergyCorrectorDRNProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup, Output const& iOutput) {
+  const auto& inputSCs = iEvent.get(inputSCToken_);
+  if (inputSCs.empty())
+    return;
+
+  const auto& serverout = energyCorrector_.getOutput(iOutput);
+
+  auto corrSCs = std::make_unique<reco::SuperClusterCollection>();
+  unsigned i = 0;
+  for (const auto& inputSC : inputSCs) {
+    float corrEn = correction(serverout[0][0 + 6 * i]) * inputSC.rawEnergy();
+    corrSCs->push_back(inputSC);
+    corrSCs->back().setEnergy(corrEn);
+    corrSCs->back().setCorrectedEnergy(corrEn);
+    ++i;
+  }
+
+  auto scHandle = iEvent.put(std::move(corrSCs));
+}
+
+void SCEnergyCorrectorDRNProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<edm::ParameterSetDescription>("correctorCfg", SCEnergyCorrectorDRN::makePSetDescription());
+  TritonClient::fillPSetDescription(desc);
+  desc.add<edm::InputTag>("inputSCs", edm::InputTag("particleFlowSuperClusterECAL"));
+  descriptions.add("scEnergyCorrectorDRNProducer", desc);
+}
+
+DEFINE_FWK_MODULE(SCEnergyCorrectorDRNProducer);

RecoEcal/EgammaClusterProducers/test/BuildFile.xml

@@ -8,3 +8,13 @@
   <use name="FWCore/ParameterSet"/>
   <flags EDM_PLUGIN="1"/>
 </library>
+
+<export>
+</export>
+<environment>
+  <bin name="DRNTest" file="runtestRecoEcalEgammaClusterProducers.cpp">
+    <flags TEST_RUNNER_ARGS="/bin/bash RecoEcal/EgammaClusterProducers/test runtests.sh"/>
+    <use name="FWCore/Utilities"/>
+  </bin>
+
+</environment>