Merge pull request #34800 from lecriste/SimTS_fromAnyCP

[HGCAL] From CaloParticle to SimTrackster in the Validation linking section

RecoHGCal/Configuration/python/RecoHGCal_EventContent_cff.py

@@ -28,6 +28,7 @@
 TICL_FEVT = cms.PSet(
     outputCommands = cms.untracked.vstring(
       'keep *_ticlSimTracksters_*_*',
+      'keep *_ticlSimTrackstersFromCP_*_*',
       )
     )
 TICL_FEVT.outputCommands.extend(TICL_RECO.outputCommands)

RecoHGCal/TICL/interface/commons.h

@@ -0,0 +1,73 @@
+#ifndef RecoHGCal_TICL_interface_commons_h
+#define RecoHGCal_TICL_interface_commons_h
+#include <vector>
+#include "DataFormats/Common/interface/Ref.h"
+#include "DataFormats/CaloRecHit/interface/CaloClusterFwd.h"
+#include "DataFormats/HGCalReco/interface/Trackster.h"
+
+namespace ticl {
+
+  inline Trackster::ParticleType tracksterParticleTypeFromPdgId(int pdgId, int charge) {
+    if (pdgId == 111) {
+      return Trackster::ParticleType::neutral_pion;
+    } else {
+      pdgId = std::abs(pdgId);
+      if (pdgId == 22) {
+        return Trackster::ParticleType::photon;
+      } else if (pdgId == 11) {
+        return Trackster::ParticleType::electron;
+      } else if (pdgId == 13) {
+        return Trackster::ParticleType::muon;
+      } else {
+        bool isHadron = (pdgId > 100 and pdgId < 900) or (pdgId > 1000 and pdgId < 9000);
+        if (isHadron) {
+          if (charge != 0) {
+            return Trackster::ParticleType::charged_hadron;
+          } else {
+            return Trackster::ParticleType::neutral_hadron;
+          }
+        } else {
+          return Trackster::ParticleType::unknown;
+        }
+      }
+    }
+  }
+
+  static void addTrackster(
+      const int& index,
+      const std::vector<std::pair<edm::Ref<reco::CaloClusterCollection>, std::pair<float, float>>>& lcVec,
+      const std::vector<float>& inputClusterMask,
+      const float& fractionCut_,
+      const float& energy,
+      const int& pdgId,
+      const int& charge,
+      const edm::ProductID& seed,
+      std::vector<float>& output_mask,
+      std::vector<Trackster>& result) {
+    if (lcVec.empty())
+      return;
+
+    Trackster tmpTrackster;
+    tmpTrackster.zeroProbabilities();
+    tmpTrackster.vertices().reserve(lcVec.size());
+    tmpTrackster.vertex_multiplicity().reserve(lcVec.size());
+    for (auto const& [lc, energyScorePair] : lcVec) {
+      if (inputClusterMask[lc.index()] > 0) {
+        double fraction = energyScorePair.first / lc->energy();
+        if (fraction < fractionCut_)
+          continue;
+        tmpTrackster.vertices().push_back(lc.index());
+        output_mask[lc.index()] -= fraction;
+        tmpTrackster.vertex_multiplicity().push_back(1. / fraction);
+      }
+    }
+
+    tmpTrackster.setIdProbability(tracksterParticleTypeFromPdgId(pdgId, charge), 1.f);
+    tmpTrackster.setRegressedEnergy(energy);
+    tmpTrackster.setSeed(seed, index);
+    result.emplace_back(tmpTrackster);
+  }
+
+}  // namespace ticl
+
+#endif

RecoHGCal/TICL/plugins/TrackstersFromCaloParticlesProducer.cc

@@ -0,0 +1,123 @@
+// Author: Leonardo Cristella - leonardo.cristella@cern.ch
+// Date: 08/2021
+
+// user include files
+
+#include "FWCore/Framework/interface/ESHandle.h"
+#include "FWCore/Framework/interface/Event.h"
+#include "FWCore/Framework/interface/Frameworkfwd.h"
+#include "FWCore/Framework/interface/MakerMacros.h"
+#include "FWCore/Framework/interface/stream/EDProducer.h"
+#include "FWCore/MessageLogger/interface/MessageLogger.h"
+#include "FWCore/ParameterSet/interface/ConfigurationDescriptions.h"
+#include "FWCore/ParameterSet/interface/ParameterSetDescription.h"
+
+#include "DataFormats/CaloRecHit/interface/CaloCluster.h"
+#include "DataFormats/ParticleFlowReco/interface/PFCluster.h"
+
+#include "DataFormats/HGCalReco/interface/Trackster.h"
+
+#include "DataFormats/Common/interface/ValueMap.h"
+#include "SimDataFormats/Associations/interface/LayerClusterToCaloParticleAssociator.h"
+
+#include "SimDataFormats/CaloAnalysis/interface/CaloParticle.h"
+#include "RecoLocalCalo/HGCalRecAlgos/interface/RecHitTools.h"
+
+#include "RecoHGCal/TICL/interface/commons.h"
+
+#include "TrackstersPCA.h"
+#include <vector>
+#include <iterator>
+#include <algorithm>
+
+using namespace ticl;
+
+class TrackstersFromCaloParticlesProducer : public edm::stream::EDProducer<> {
+public:
+  explicit TrackstersFromCaloParticlesProducer(const edm::ParameterSet&);
+  static void fillDescriptions(edm::ConfigurationDescriptions& descriptions);
+
+  void produce(edm::Event&, const edm::EventSetup&) override;
+
+private:
+  std::string detector_;
+  const bool doNose_ = false;
+  const edm::EDGetTokenT<std::vector<reco::CaloCluster>> clusters_token_;
+  const edm::EDGetTokenT<edm::ValueMap<std::pair<float, float>>> clustersTime_token_;
+  const edm::EDGetTokenT<std::vector<float>> filtered_layerclusters_mask_token_;
+
+  const edm::EDGetTokenT<std::vector<CaloParticle>> caloparticles_token_;
+
+  const edm::EDGetTokenT<hgcal::SimToRecoCollection> associatorMapCaloParticleToReco_token_;
+  const edm::ESGetToken<CaloGeometry, CaloGeometryRecord> geom_token_;
+  hgcal::RecHitTools rhtools_;
+  const double fractionCut_;
+};
+DEFINE_FWK_MODULE(TrackstersFromCaloParticlesProducer);
+
+TrackstersFromCaloParticlesProducer::TrackstersFromCaloParticlesProducer(const edm::ParameterSet& ps)
+    : detector_(ps.getParameter<std::string>("detector")),
+      doNose_(detector_ == "HFNose"),
+      clusters_token_(consumes(ps.getParameter<edm::InputTag>("layer_clusters"))),
+      clustersTime_token_(consumes(ps.getParameter<edm::InputTag>("time_layerclusters"))),
+      filtered_layerclusters_mask_token_(consumes(ps.getParameter<edm::InputTag>("filtered_mask"))),
+      caloparticles_token_(consumes(ps.getParameter<edm::InputTag>("caloparticles"))),
+      associatorMapCaloParticleToReco_token_(consumes(ps.getParameter<edm::InputTag>("associator"))),
+      geom_token_(esConsumes()),
+      fractionCut_(ps.getParameter<double>("fractionCut")) {
+  produces<std::vector<Trackster>>();
+  produces<std::vector<float>>();
+}
+
+void TrackstersFromCaloParticlesProducer::fillDescriptions(edm::ConfigurationDescriptions& descriptions) {
+  edm::ParameterSetDescription desc;
+  desc.add<std::string>("detector", "HGCAL");
+  desc.add<edm::InputTag>("layer_clusters", edm::InputTag("hgcalLayerClusters"));
+  desc.add<edm::InputTag>("time_layerclusters", edm::InputTag("hgcalLayerClusters", "timeLayerCluster"));
+  desc.add<edm::InputTag>("filtered_mask", edm::InputTag("filteredLayerClustersSimTracksters", "ticlSimTracksters"));
+  desc.add<edm::InputTag>("caloparticles", edm::InputTag("mix", "MergedCaloTruth"));
+  desc.add<edm::InputTag>("associator", edm::InputTag("layerClusterCaloParticleAssociationProducer"));
+  desc.add<double>("fractionCut", 0.);
+
+  descriptions.addWithDefaultLabel(desc);
+}
+
+void TrackstersFromCaloParticlesProducer::produce(edm::Event& evt, const edm::EventSetup& es) {
+  auto result = std::make_unique<std::vector<Trackster>>();
+  auto output_mask = std::make_unique<std::vector<float>>();
+  const auto& layerClusters = evt.get(clusters_token_);
+  const auto& layerClustersTimes = evt.get(clustersTime_token_);
+  const auto& inputClusterMask = evt.get(filtered_layerclusters_mask_token_);
+  output_mask->resize(layerClusters.size(), 1.f);
+
+  const auto& caloparticles = evt.get(caloparticles_token_);
+
+  const auto& caloParticlesToRecoColl = evt.get(associatorMapCaloParticleToReco_token_);
+
+  const auto& geom = es.getData(geom_token_);
+  rhtools_.setGeometry(geom);
+  result->reserve(caloparticles.size());
+
+  for (const auto& [key, values] : caloParticlesToRecoColl) {
+    auto const& cp = *(key);
+    auto cpIndex = &cp - &caloparticles[0];
+
+    addTrackster(cpIndex,
+                 values,
+                 inputClusterMask,
+                 fractionCut_,
+                 cp.g4Tracks()[0].getMomentumAtBoundary().energy(),
+                 cp.pdgId(),
+                 cp.charge(),
+                 key.id(),
+                 *output_mask,
+                 *result);
+  }
+
+  ticl::assignPCAtoTracksters(
+      *result, layerClusters, layerClustersTimes, rhtools_.getPositionLayer(rhtools_.lastLayerEE(doNose_)).z());
+  result->shrink_to_fit();
+
+  evt.put(std::move(result));
+  evt.put(std::move(output_mask));
+}