src/components/SimpleClustering.cpp

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