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