calorimetry/todo/SimpleClustering.h

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Whitney Armstrong, Sylvester Joosten, Wouter Deconinck, Chao Peng
0003
0004 #include <algorithm>
0005
0006 #include "Gaudi/Property.h"
0007 #include "GaudiAlg/GaudiAlgorithm.h"
0008 #include "GaudiAlg/GaudiTool.h"
0009 #include "GaudiAlg/Transformer.h"
0010 #include "GaudiKernel/PhysicalConstants.h"
0011 #include "GaudiKernel/RndmGenerators.h"
0012 #include "GaudiKernel/ToolHandle.h"
0013
0014 #include "DDRec/CellIDPositionConverter.h"
0015 #include "DDRec/Surface.h"
0016 #include "DDRec/SurfaceManager.h"
0017
0018 #include "JugBase/DataHandle.h"
0019 #include "JugBase/IGeoSvc.h"
0020
0021 // Event Model related classes
0022 #include "edm4hep/SimCalorimeterHitCollection.h"
0023 #include "edm4eic/CalorimeterHitCollection.h"
0024 #include "edm4eic/ClusterCollection.h"
0025 #include "edm4eic/ProtoClusterCollection.h"
0026 #include "edm4eic/RawCalorimeterHitCollection.h"
0027 #include "edm4hep/utils/vector_utils.h"
0028
0029 using namespace Gaudi::Units;
0030
0031 namespace Jug::Reco {
0032
0033   /** Simple clustering algorithm.
0034    *
0035    * \ingroup reco
0036    */
0037   class SimpleClustering : public GaudiAlgorithm {
0038   private:
0039     using RecHits  = edm4eic::CalorimeterHitCollection;
0040     using ProtoClusters = edm4eic::ProtoClusterCollection;
0041     using Clusters = edm4eic::ClusterCollection;
0042
0043     DataHandle<RecHits>       m_inputHitCollection{"inputHitCollection", Gaudi::DataHandle::Reader, this};
0044     DataHandle<ProtoClusters> m_outputProtoClusters{"outputProtoCluster", Gaudi::DataHandle::Writer, this};
0045     DataHandle<Clusters>      m_outputClusters{"outputClusterCollection", Gaudi::DataHandle::Writer, this};
0046
0047     Gaudi::Property<std::string> m_mcHits{this, "mcHits", ""};
0048
0049     Gaudi::Property<double>   m_minModuleEdep{this, "minModuleEdep", 5.0 * MeV};
0050     Gaudi::Property<double>   m_maxDistance{this, "maxDistance", 20.0 * cm};
0051
0052     /// Pointer to the geometry service
0053     SmartIF<IGeoSvc> m_geoSvc;
0054
0055     // Optional handle to MC hits
0056     std::unique_ptr<DataHandle<edm4hep::SimCalorimeterHitCollection>> m_inputMC;
0057
0058   public:
0059     SimpleClustering(const std::string& name, ISvcLocator* svcLoc)
0060       : GaudiAlgorithm(name, svcLoc) {
0061       declareProperty("inputHitCollection", m_inputHitCollection, "");
0062       declareProperty("outputProtoClusterCollection", m_outputClusters, "Output proto clusters");
0063       declareProperty("outputClusterCollection", m_outputClusters, "Output clusters");
0064     }
0065
0066     StatusCode initialize() override
0067     {
0068       if (GaudiAlgorithm::initialize().isFailure()) {
0069         return StatusCode::FAILURE;
0070       }
0071       // Initialize the MC input hit collection if requested
0072       if (m_mcHits != "") {
0073         m_inputMC =
0074             std::make_unique<DataHandle<edm4hep::SimCalorimeterHitCollection>>(m_mcHits, Gaudi::DataHandle::Reader, this);
0075       }
0076       m_geoSvc = service("GeoSvc");
0077       if (!m_geoSvc) {
0078         error() << "Unable to locate Geometry Service. "
0079                 << "Make sure you have GeoSvc and SimSvc in the right order in the configuration." << endmsg;
0080         return StatusCode::FAILURE;
0081       }
0082       return StatusCode::SUCCESS;
0083     }
0084
0085     StatusCode execute() override
0086     {
0087       // input collections
0088       const auto& hits = *m_inputHitCollection.get();
0089       // Create output collections
0090       auto& proto = *m_outputProtoClusters.createAndPut();
0091       auto& clusters = *m_outputClusters.createAndPut();
0092
0093       // Optional MC data
0094       // FIXME no connection between cluster and truth in edm4hep
0095       //const edm4hep::SimCalorimeterHitCollection* mcHits = nullptr;
0096       //if (m_inputMC) {
0097       //  mcHits = m_inputMC->get();
0098       //}
0099
0100       std::vector<std::pair<uint32_t, edm4eic::CalorimeterHit>> the_hits;
0101       std::vector<std::pair<uint32_t, edm4eic::CalorimeterHit>> remaining_hits;
0102
0103       double max_dist   = m_maxDistance.value() / mm;
0104       double min_energy = m_minModuleEdep.value() / GeV;
0105
0106       edm4eic::CalorimeterHit ref_hit;
0107       bool have_ref = false;
0108       // Collect all our hits, and get the highest energy hit
0109       {
0110         uint32_t idx  = 0;
0111         for (const auto& h : hits) {
0112           if (!have_ref || h.getEnergy() > ref_hit.getEnergy()) {
0113             ref_hit  = h;
0114             have_ref = true;
0115           }
0116           the_hits.emplace_back(idx, h);
0117           idx += 1;
0118         }
0119       }
0120
0121       if (msgLevel(MSG::DEBUG)) {
0122         debug() << " max_dist = " << max_dist << endmsg;
0123       }
0124
0125       while (have_ref && ref_hit.getEnergy() > min_energy) {
0126
0127         std::vector<std::pair<uint32_t, edm4eic::CalorimeterHit>> cluster_hits;
0128
0129         for (const auto& [idx, h] : the_hits) {
0130           if (edm4hep::utils::magnitude(h.getPosition() - ref_hit.getPosition()) < max_dist) {
0131             cluster_hits.emplace_back(idx, h);
0132           } else {
0133             remaining_hits.emplace_back(idx, h);
0134           }
0135         }
0136
0137         double total_energy = std::accumulate(
0138             std::begin(cluster_hits), std::end(cluster_hits), 0.0,
0139             [](double t, const std::pair<uint32_t, edm4eic::CalorimeterHit>& h1) { return (t + h1.second.getEnergy()); });
0140
0141         if (msgLevel(MSG::DEBUG)) {
0142           debug() << " total_energy = " << total_energy << endmsg;
0143           debug() << " cluster size " << cluster_hits.size() << endmsg;
0144         }
0145         auto cl = clusters.create();
0146         cl.setNhits(cluster_hits.size());
0147         auto pcl = proto.create();
0148         for (const auto& [idx, h] : cluster_hits) {
0149           cl.setEnergy(cl.getEnergy() + h.getEnergy());
0150           cl.setPosition(cl.getPosition() + (h.getPosition() * h.getEnergy() / total_energy));
0151           pcl.addToHits(h);
0152           pcl.addToWeights(1);
0153         }
0154         // Optionally store the MC truth associated with the first hit in this cluster
0155         // FIXME no connection between cluster and truth in edm4hep
0156         //if (mcHits) {
0157         //  const auto& mc_hit = (*mcHits)[ref_hit.ID().value];
0158         //  cl.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0159         //}
0160
0161         have_ref = false;
0162         if ((remaining_hits.size() > 5) && (clusters.size() < 10)) {
0163           for (const auto& [idx, h] : remaining_hits) {
0164             if (!have_ref || h.getEnergy() > ref_hit.getEnergy()) {
0165               ref_hit  = h;
0166               have_ref = true;
0167             }
0168           }
0169
0170           std::swap(remaining_hits, the_hits);
0171           remaining_hits.clear();
0172         }
0173       }
0174       if (msgLevel(MSG::DEBUG)) {
0175         debug() << " clusters found: " << clusters.size() << endmsg;
0176       }
0177
0178       return StatusCode::SUCCESS;
0179     }
0180   };
0181
0182   // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0183   DECLARE_COMPONENT(SimpleClustering)
0184
0185 } // namespace Jug::Reco