[RFC] Outline of HGCAL ML ntuples using NanoAOD

CommonTools/RecoAlgos/plugins/RecHitToPFCandAssociationProducer.cc

@@ -0,0 +1,133 @@
+// system include files
+#include <memory>
+#include <string>
+
+// user include files
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "FWCore/Framework/interface/ESHandle.h"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "FWCore/MessageLogger/interface/MessageLogger.h"
+#include "SimDataFormats/TrackingAnalysis/interface/TrackingParticle.h"
+#include "SimDataFormats/CaloAnalysis/interface/SimClusterFwd.h"
+#include "SimDataFormats/CaloAnalysis/interface/SimCluster.h"
+#include "DataFormats/CaloRecHit/interface/CaloRecHit.h"
+#include "SimDataFormats/CaloHit/interface/PCaloHit.h"
+#include "SimDataFormats/CaloHit/interface/PCaloHitContainer.h"
+#include "DataFormats/HGCRecHit/interface/HGCRecHitCollections.h"
+
+#include "DataFormats/ParticleFlowCandidate/interface/PFCandidateFwd.h"
+#include "DataFormats/ParticleFlowCandidate/interface/PFCandidate.h"
+#include "DataFormats/ParticleFlowReco/interface/PFRecHit.h"
+#include "DataFormats/ParticleFlowReco/interface/PFRecHitFraction.h"
+#include "DataFormats/ParticleFlowReco/interface/PFBlock.h"
+#include "DataFormats/ParticleFlowReco/interface/PFBlockFwd.h"
+#include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
+#include "DataFormats/ParticleFlowReco/interface/PFClusterFwd.h"
+
+#include "DataFormats/Common/interface/Association.h"
+
+#include "FWCore/Utilities/interface/transform.h"
+#include "FWCore/Utilities/interface/EDGetToken.h"
+#include <set>
+
+//
+// class decleration
+//
+typedef std::pair<size_t, float> IdxAndFraction;
+
+class RecHitToPFCandAssociationProducer : public edm::stream::EDProducer<> {
+public:
+  explicit RecHitToPFCandAssociationProducer(const edm::ParameterSet&);
+  ~RecHitToPFCandAssociationProducer() override;
+
+private:
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+  std::vector<edm::InputTag> caloRechitTags_;
+  std::vector<edm::EDGetTokenT<edm::PCaloHitContainer>> caloSimhitCollectionTokens_;
+  //std::vector<edm::EDGetTokenT<std::vector<PSimHit>> trackSimhitCollectionTokens_;
+  std::vector<edm::EDGetTokenT<edm::View<CaloRecHit>>> caloRechitCollectionTokens_;
+  edm::EDGetTokenT<reco::PFCandidateCollection> pfCollectionToken_;
+};
+
+RecHitToPFCandAssociationProducer::RecHitToPFCandAssociationProducer(const edm::ParameterSet& pset)
+    :  //caloSimhitCollectionTokens_(edm::vector_transform(pset.getParameter<std::vector<edm::InputTag>>("caloSimHits"),
+      //  [this](const edm::InputTag& tag) {return mayConsume<edm::PCaloHitContainer>(tag); })),
+      //trackSimhitCollectionTokens_(edm::vector_transform(pset.getParameter<edm::InputTag>("trackSimHits"),
+      //    [this](const edm::InputTag& tag) {return mayConsume<std::vector<PSimHit>(tag); }),
+      caloRechitTags_(pset.getParameter<std::vector<edm::InputTag>>("caloRecHits")),
+      caloRechitCollectionTokens_(edm::vector_transform(
+          caloRechitTags_, [this](const edm::InputTag& tag) { return mayConsume<edm::View<CaloRecHit>>(tag); })),
+      pfCollectionToken_(consumes<reco::PFCandidateCollection>(pset.getParameter<edm::InputTag>("pfCands"))) {
+  for (auto& tag : caloRechitTags_) {
+    const std::string& label = tag.instance();
+    //TODO: Can this be an edm::View?
+    produces<edm::Association<reco::PFCandidateCollection>>(label + "ToPFCand");
+  }
+}
+
+RecHitToPFCandAssociationProducer::~RecHitToPFCandAssociationProducer() {}
+
+//
+// member functions
+//
+
+// ------------ method called to produce the data  ------------
+void RecHitToPFCandAssociationProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  edm::Handle<reco::PFCandidateCollection> pfCollection;
+  iEvent.getByToken(pfCollectionToken_, pfCollection);
+  std::unordered_map<size_t, IdxAndFraction> hitDetIdToIndex;
+
+  for (size_t j = 0; j < pfCollection->size(); j++) {
+    const auto& pfCand = pfCollection->at(j);
+    const reco::PFCandidate::ElementsInBlocks& elements = pfCand.elementsInBlocks();
+    for (auto& element : elements) {
+      const reco::PFBlockRef blockRef = element.first;
+      for (const auto& block : blockRef->elements()) {
+        if (block.type() == reco::PFBlockElement::HGCAL) {
+          const reco::PFClusterRef& cluster = block.clusterRef();
+          const std::vector<reco::PFRecHitFraction>& rhf = cluster->recHitFractions();
+          for (const auto& hf : rhf) {
+            auto& hit = hf.recHitRef();
+            if (!hit)
+              throw cms::Exception("RecHitToPFCandAssociationProducer") << "Invalid RecHit ref";
+            size_t detId = hit->detId();
+            auto entry = hitDetIdToIndex.find(detId);
+            if (entry == hitDetIdToIndex.end() || entry->second.second < hf.fraction())
+              hitDetIdToIndex[detId] = {j, hf.fraction()};
+          }
+        }
+      }
+    }
+  }
+
+  for (size_t i = 0; i < caloRechitCollectionTokens_.size(); i++) {
+    std::string label = caloRechitTags_.at(i).instance();
+    std::vector<size_t> rechitIndices;
+
+    edm::Handle<edm::View<CaloRecHit>> caloRechitCollection;
+    iEvent.getByToken(caloRechitCollectionTokens_.at(i), caloRechitCollection);
+
+    for (size_t h = 0; h < caloRechitCollection->size(); h++) {
+      const CaloRecHit& caloRh = caloRechitCollection->at(h);
+      size_t id = caloRh.detid().rawId();
+      int match = hitDetIdToIndex.find(id) == hitDetIdToIndex.end() ? -1 : hitDetIdToIndex.at(id).first;
+      rechitIndices.push_back(match);
+    }
+
+    auto assoc = std::make_unique<edm::Association<reco::PFCandidateCollection>>(pfCollection);
+    edm::Association<reco::PFCandidateCollection>::Filler filler(*assoc);
+    filler.insert(caloRechitCollection, rechitIndices.begin(), rechitIndices.end());
+    filler.fill();
+    iEvent.put(std::move(assoc), label + "ToPFCand");
+  }
+}
+
+// define this as a plug-in
+DEFINE_FWK_MODULE(RecHitToPFCandAssociationProducer);

CommonTools/RecoAlgos/plugins/SimClusterRecHitAssocitionProducer.cc

@@ -0,0 +1,136 @@
+// system include files
+#include <memory>
+#include <string>
+
+// user include files
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+
+#include "FWCore/Framework/interface/ESHandle.h"
+
+#include "FWCore/ParameterSet/interface/ParameterSet.h"
+
+#include "FWCore/MessageLogger/interface/MessageLogger.h"
+#include "SimDataFormats/TrackingAnalysis/interface/TrackingParticle.h"
+#include "SimDataFormats/CaloAnalysis/interface/SimClusterFwd.h"
+#include "SimDataFormats/CaloAnalysis/interface/SimCluster.h"
+#include "DataFormats/CaloRecHit/interface/CaloRecHit.h"
+#include "SimDataFormats/CaloHit/interface/PCaloHit.h"
+#include "SimDataFormats/CaloHit/interface/PCaloHitContainer.h"
+
+#include "DataFormats/Common/interface/Association.h"
+
+#include "FWCore/Utilities/interface/transform.h"
+#include "FWCore/Utilities/interface/EDGetToken.h"
+#include <set>
+
+//
+// class decleration
+//
+typedef std::pair<size_t, float> IdxAndFraction;
+
+class SimClusterRecHitAssociationProducer : public edm::stream::EDProducer<> {
+public:
+  explicit SimClusterRecHitAssociationProducer(const edm::ParameterSet&);
+  ~SimClusterRecHitAssociationProducer() override;
+
+private:
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+  std::vector<edm::InputTag> caloRechitTags_;
+  std::vector<edm::EDGetTokenT<edm::PCaloHitContainer>> caloSimhitCollectionTokens_;
+  //std::vector<edm::EDGetTokenT<std::vector<PSimHit>> trackSimhitCollectionTokens_;
+  std::vector<edm::EDGetTokenT<edm::View<CaloRecHit>>> caloRechitCollectionTokens_;
+  edm::EDGetTokenT<SimClusterCollection> scCollectionToken_;
+};
+
+SimClusterRecHitAssociationProducer::SimClusterRecHitAssociationProducer(const edm::ParameterSet& pset)
+    :  //caloSimhitCollectionTokens_(edm::vector_transform(pset.getParameter<std::vector<edm::InputTag>>("caloSimHits"),
+      //  [this](const edm::InputTag& tag) {return mayConsume<edm::PCaloHitContainer>(tag); })),
+      //trackSimhitCollectionTokens_(edm::vector_transform(pset.getParameter<edm::InputTag>("trackSimHits"),
+      //    [this](const edm::InputTag& tag) {return mayConsume<std::vector<PSimHit>(tag); }),
+      caloRechitTags_(pset.getParameter<std::vector<edm::InputTag>>("caloRecHits")),
+      caloRechitCollectionTokens_(edm::vector_transform(
+          caloRechitTags_, [this](const edm::InputTag& tag) { return mayConsume<edm::View<CaloRecHit>>(tag); })),
+      scCollectionToken_(consumes<SimClusterCollection>(pset.getParameter<edm::InputTag>("simClusters"))) {
+  for (auto& tag : caloRechitTags_) {
+    const std::string& label = tag.instance();
+    produces<edm::Association<SimClusterCollection>>(label + "ToSimClus");
+  }
+  produces<std::unordered_map<int, float>>();
+}
+
+SimClusterRecHitAssociationProducer::~SimClusterRecHitAssociationProducer() {}
+
+//
+// member functions
+//
+
+// ------------ method called to produce the data  ------------
+void SimClusterRecHitAssociationProducer::produce(edm::Event& iEvent, const edm::EventSetup& iSetup) {
+  auto simClusterToRecEnergy = std::make_unique<std::unordered_map<int, float>>();
+  edm::Handle<SimClusterCollection> scCollection;
+  iEvent.getByToken(scCollectionToken_, scCollection);
+  std::unordered_map<size_t, IdxAndFraction> hitDetIdToIndex;
+  std::unordered_map<size_t, float> hitDetIdToTotalSimFrac;
+
+  for (size_t s = 0; s < scCollection->size(); s++) {
+    const auto& sc = scCollection->at(s);
+    (*simClusterToRecEnergy)[s] = 0.;
+    for (auto& hf : sc.hits_and_fractions()) {
+      auto entry = hitDetIdToIndex.find(hf.first);
+      // Update SimCluster assigment if detId has been found in no other SCs or if
+      // SC has greater fraction of energy in DetId than the SC already found
+      if (entry == hitDetIdToIndex.end()) {
+        hitDetIdToTotalSimFrac[hf.first] = hf.second;
+        hitDetIdToIndex[hf.first] = {s, hf.second};
+      } else {
+        hitDetIdToTotalSimFrac[hf.first] += hf.second;
+        if (entry->second.second < hf.second)
+          hitDetIdToIndex[hf.first] = {s, hf.second};
+      }
+    }
+  }
+
+  std::unordered_map<int, float> hitDetIdToTotalRecEnergy;
+
+  for (size_t i = 0; i < caloRechitCollectionTokens_.size(); i++) {
+    std::string label = caloRechitTags_.at(i).instance();
+    std::vector<size_t> rechitIndices;
+
+    edm::Handle<edm::View<CaloRecHit>> caloRechitCollection;
+    iEvent.getByToken(caloRechitCollectionTokens_.at(i), caloRechitCollection);
+
+    for (size_t h = 0; h < caloRechitCollection->size(); h++) {
+      const CaloRecHit& caloRh = caloRechitCollection->at(h);
+      size_t id = caloRh.detid().rawId();
+      auto entry = hitDetIdToTotalRecEnergy.find(id);
+      float energy = caloRh.energy();
+      if (entry == hitDetIdToTotalRecEnergy.end())
+        hitDetIdToTotalRecEnergy[id] = energy;
+      else
+        hitDetIdToTotalRecEnergy.at(id) += energy;
+
+      int match = hitDetIdToIndex.find(id) == hitDetIdToIndex.end() ? -1 : hitDetIdToIndex.at(id).first;
+      float fraction = match != -1 ? hitDetIdToTotalSimFrac.at(id) : 1.;
+      if (simClusterToRecEnergy->find(match) == simClusterToRecEnergy->end())
+        (*simClusterToRecEnergy)[match] = energy * fraction;
+      else
+        simClusterToRecEnergy->at(match) += energy * fraction;
+
+      rechitIndices.push_back(match);
+    }
+
+    auto assoc = std::make_unique<edm::Association<SimClusterCollection>>(scCollection);
+    edm::Association<SimClusterCollection>::Filler filler(*assoc);
+    filler.insert(caloRechitCollection, rechitIndices.begin(), rechitIndices.end());
+    filler.fill();
+    iEvent.put(std::move(assoc), label + "ToSimClus");
+  }
+  iEvent.put(std::move(simClusterToRecEnergy));
+}
+
+// define this as a plug-in
+DEFINE_FWK_MODULE(SimClusterRecHitAssociationProducer);

DPGAnalysis/CommonNanoAOD/BuildFile.xml

@@ -0,0 +1,6 @@
+<use name="DataFormats/Common"/>
+<use name="DataFormats/NanoAOD"/>
+<use name="boost"/>
+<export>
+  <lib name="1"/>
+</export>

DPGAnalysis/CommonNanoAOD/plugins/BuildFile.xml

@@ -0,0 +1,19 @@
+<use name="FWCore/Framework"/>
+<use name="FWCore/ParameterSet"/>
+<use name="FWCore/ServiceRegistry"/>
+<use name="FWCore/Utilities"/>
+<use name="Geometry/HGCalGeometry"/>
+<use name="Geometry/CaloGeometry"/>
+<use name="Geometry/CSCGeometry"/>
+<use name="Geometry/TrackerGeometryBuilder"/>
+<use name="PhysicsTools/PatAlgos"/>
+<use name="DataFormats/NanoAOD"/>
+<use name="SimDataFormats/CaloAnalysis"/>
+<use name="SimDataFormats/Track"/>
+<use name="IOPool/Provenance"/>
+<use name="CondFormats/RunInfo"/>
+<use name="CondFormats/DataRecord"/>
+<use name="RecoLocalCalo/HGCalRecAlgos"/>
+<library file="*.cc" name="DPGAnalysisCommonNanoAODPlugins">
+  <flags EDM_PLUGIN="1"/>
+</library>