src/todo/ImagingTopoCluster.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Chao Peng, Sylvester Joosten, Whitney Armstrong
0003
0004 /*
0005  *  Topological Cell Clustering Algorithm for Imaging Calorimetry
0006  *  1. group all the adjacent pixels
0007  *
0008  *  Author: Chao Peng (ANL), 06/02/2021
0009  *  References: https://arxiv.org/pdf/1603.02934.pdf
0010  *
0011  */
0012 #include "fmt/format.h"
0013 #include <algorithm>
0014
0015 #include "Gaudi/Property.h"
0016 #include "GaudiAlg/GaudiAlgorithm.h"
0017 #include "GaudiAlg/GaudiTool.h"
0018 #include "GaudiAlg/Transformer.h"
0019 #include "GaudiKernel/PhysicalConstants.h"
0020 #include "GaudiKernel/RndmGenerators.h"
0021 #include "GaudiKernel/ToolHandle.h"
0022
0023 #include "DD4hep/BitFieldCoder.h"
0024 #include "DDRec/CellIDPositionConverter.h"
0025 #include "DDRec/Surface.h"
0026 #include "DDRec/SurfaceManager.h"
0027
0028 #include "JugBase/DataHandle.h"
0029 #include "JugBase/IGeoSvc.h"
0030 #include "JugReco/ClusterTypes.h"
0031
0032 // Event Model related classes
0033 #include "edm4eic/CalorimeterHitCollection.h"
0034 #include "edm4eic/ProtoClusterCollection.h"
0035 #include "edm4hep/utils/vector_utils.h"
0036
0037 using namespace Gaudi::Units;
0038
0039 namespace Jug::Reco {
0040
0041 /** Topological Cell Clustering Algorithm.
0042  *
0043  * Topological Cell Clustering Algorithm for Imaging Calorimetry
0044  *  1. group all the adjacent pixels
0045  *
0046  *  Author: Chao Peng (ANL), 06/02/2021
0047  *  References: https://arxiv.org/pdf/1603.02934.pdf
0048  *
0049  * \ingroup reco
0050  */
0051 class ImagingTopoCluster : public GaudiAlgorithm {
0052 private:
0053   // maximum difference in layer numbers that can be considered as neighbours
0054   Gaudi::Property<int> m_neighbourLayersRange{this, "neighbourLayersRange", 1};
0055   // maximum distance of local (x, y) to be considered as neighbors at the same layer
0056   Gaudi::Property<std::vector<double>> u_localDistXY{this, "localDistXY", {1.0 * mm, 1.0 * mm}};
0057   // maximum distance of global (eta, phi) to be considered as neighbors at different layers
0058   Gaudi::Property<std::vector<double>> u_layerDistEtaPhi{this, "layerDistEtaPhi", {0.01, 0.01}};
0059   // maximum global distance to be considered as neighbors in different sectors
0060   Gaudi::Property<double> m_sectorDist{this, "sectorDist", 1.0 * cm};
0061
0062   // minimum hit energy to participate clustering
0063   Gaudi::Property<double> m_minClusterHitEdep{this, "minClusterHitEdep", 0.};
0064   // minimum cluster center energy (to be considered as a seed for cluster)
0065   Gaudi::Property<double> m_minClusterCenterEdep{this, "minClusterCenterEdep", 0.};
0066   // minimum cluster energy (to save this cluster)
0067   Gaudi::Property<double> m_minClusterEdep{this, "minClusterEdep", 0.5 * MeV};
0068   // minimum number of hits (to save this cluster)
0069   Gaudi::Property<int> m_minClusterNhits{this, "minClusterNhits", 10};
0070   // input hits collection
0071   DataHandle<edm4eic::CalorimeterHitCollection> m_inputHitCollection{"inputHitCollection", Gaudi::DataHandle::Reader,
0072                                                                   this};
0073   // output clustered hits
0074   DataHandle<edm4eic::ProtoClusterCollection> m_outputProtoClusterCollection{"outputProtoClusterCollection",
0075                                                                           Gaudi::DataHandle::Writer, this};
0076
0077   // unitless counterparts of the input parameters
0078   double localDistXY[2]{0,0}, layerDistEtaPhi[2]{0,0}, sectorDist{0};
0079   double minClusterHitEdep{0}, minClusterCenterEdep{0}, minClusterEdep{0}, minClusterNhits{0};
0080
0081 public:
0082   ImagingTopoCluster(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0083     declareProperty("inputHitCollection", m_inputHitCollection, "");
0084     declareProperty("outputProtoClusterCollection", m_outputProtoClusterCollection, "");
0085   }
0086
0087   StatusCode initialize() override {
0088     if (GaudiAlgorithm::initialize().isFailure()) {
0089       return StatusCode::FAILURE;
0090     }
0091
0092     // unitless conversion
0093     // sanity checks
0094     if (u_localDistXY.size() != 2) {
0095       error() << "Expected 2 values (x_dist, y_dist) for localDistXY" << endmsg;
0096       return StatusCode::FAILURE;
0097     }
0098     if (u_layerDistEtaPhi.size() != 2) {
0099       error() << "Expected 2 values (eta_dist, phi_dist) for layerDistEtaPhi" << endmsg;
0100       return StatusCode::FAILURE;
0101     }
0102
0103     // using juggler internal units (GeV, mm, ns, rad)
0104     localDistXY[0]       = u_localDistXY.value()[0] / mm;
0105     localDistXY[1]       = u_localDistXY.value()[1] / mm;
0106     layerDistEtaPhi[0]   = u_layerDistEtaPhi.value()[0];
0107     layerDistEtaPhi[1]   = u_layerDistEtaPhi.value()[1] / rad;
0108     sectorDist           = m_sectorDist.value() / mm;
0109     minClusterHitEdep    = m_minClusterHitEdep.value() / GeV;
0110     minClusterCenterEdep = m_minClusterCenterEdep.value() / GeV;
0111     minClusterEdep       = m_minClusterEdep.value() / GeV;
0112
0113     // summarize the clustering parameters
0114     info() << fmt::format("Local clustering (same sector and same layer): "
0115                           "Local [x, y] distance between hits <= [{:.4f} mm, {:.4f} mm].",
0116                           localDistXY[0], localDistXY[1])
0117            << endmsg;
0118     info() << fmt::format("Neighbour layers clustering (same sector and layer id within +- {:d}: "
0119                           "Global [eta, phi] distance between hits <= [{:.4f}, {:.4f} rad].",
0120                           m_neighbourLayersRange.value(), layerDistEtaPhi[0], layerDistEtaPhi[1])
0121            << endmsg;
0122     info() << fmt::format("Neighbour sectors clustering (different sector): "
0123                           "Global distance between hits <= {:.4f} mm.",
0124                           sectorDist)
0125            << endmsg;
0126
0127     return StatusCode::SUCCESS;
0128   }
0129
0130   StatusCode execute() override {
0131     // input collections
0132     const auto& hits = *m_inputHitCollection.get();
0133     // Create output collections
0134     auto& proto = *m_outputProtoClusterCollection.createAndPut();
0135
0136     // group neighboring hits
0137     std::vector<bool> visits(hits.size(), false);
0138     std::vector<std::vector<std::pair<uint32_t, edm4eic::CalorimeterHit>>> groups;
0139     for (size_t i = 0; i < hits.size(); ++i) {
0140       if (msgLevel(MSG::DEBUG)) {
0141         debug() << fmt::format("hit {:d}: local position = ({}, {}, {}), global position = ({}, {}, {})", i + 1,
0142                                hits[i].getLocal().x, hits[i].getLocal().y, hits[i].getPosition().z,
0143                                hits[i].getPosition().x, hits[i].getPosition().y, hits[i].getPosition().z)
0144                 << endmsg;
0145       }
0146       // already in a group, or not energetic enough to form a cluster
0147       if (visits[i] || hits[i].getEnergy() < minClusterCenterEdep) {
0148         continue;
0149       }
0150       groups.emplace_back();
0151       // create a new group, and group all the neighboring hits
0152       dfs_group(groups.back(), i, hits, visits);
0153     }
0154     if (msgLevel(MSG::DEBUG)) {
0155       debug() << "found " << groups.size() << " potential clusters (groups of hits)" << endmsg;
0156       for (size_t i = 0; i < groups.size(); ++i) {
0157         debug() << fmt::format("group {}: {} hits", i, groups[i].size()) << endmsg;
0158       }
0159     }
0160
0161     // form clusters
0162     for (const auto& group : groups) {
0163       if (static_cast<int>(group.size()) < m_minClusterNhits.value()) {
0164         continue;
0165       }
0166       double energy = 0.;
0167       for (const auto& [idx, hit] : group) {
0168         energy += hit.getEnergy();
0169       }
0170       if (energy < minClusterEdep) {
0171         continue;
0172       }
0173       auto pcl = proto.create();
0174       for (const auto& [idx, hit] : group) {
0175         pcl.addToHits(hit);
0176         pcl.addToWeights(1);
0177       }
0178     }
0179
0180     return StatusCode::SUCCESS;
0181   }
0182
0183 private:
0184   template <typename T> static inline T pow2(const T& x) { return x * x; }
0185
0186   // helper function to group hits
0187   bool is_neighbor(const edm4eic::CalorimeterHit& h1, const edm4eic::CalorimeterHit& h2) const {
0188     // different sectors, simple distance check
0189     if (h1.getSector() != h2.getSector()) {
0190       return std::sqrt(pow2(h1.getPosition().x - h2.getPosition().x) + pow2(h1.getPosition().y - h2.getPosition().y) +
0191                        pow2(h1.getPosition().z - h2.getPosition().z)) <= sectorDist;
0192     }
0193
0194     // layer check
0195     int ldiff = std::abs(h1.getLayer() - h2.getLayer());
0196     // same layer, check local positions
0197     if (ldiff == 0) {
0198       return (std::abs(h1.getLocal().x - h2.getLocal().x) <= localDistXY[0]) &&
0199              (std::abs(h1.getLocal().y - h2.getLocal().y) <= localDistXY[1]);
0200     } else if (ldiff <= m_neighbourLayersRange) {
0201       return (std::abs(edm4hep::utils::eta(h1.getPosition()) - edm4hep::utils::eta(h2.getPosition())) <= layerDistEtaPhi[0]) &&
0202              (std::abs(edm4hep::utils::angleAzimuthal(h1.getPosition()) - edm4hep::utils::angleAzimuthal(h2.getPosition())) <=
0203               layerDistEtaPhi[1]);
0204     }
0205
0206     // not in adjacent layers
0207     return false;
0208   }
0209
0210   // grouping function with Depth-First Search
0211   void dfs_group(std::vector<std::pair<uint32_t, edm4eic::CalorimeterHit>>& group, int idx,
0212                  const edm4eic::CalorimeterHitCollection& hits, std::vector<bool>& visits) const {
0213     // not a qualified hit to participate in clustering, stop here
0214     if (hits[idx].getEnergy() < minClusterHitEdep) {
0215       visits[idx] = true;
0216       return;
0217     }
0218
0219     group.emplace_back(idx, hits[idx]);
0220     visits[idx] = true;
0221     for (size_t i = 0; i < hits.size(); ++i) {
0222       // visited, or not a neighbor
0223       if (visits[i] || !is_neighbor(hits[idx], hits[i])) {
0224         continue;
0225       }
0226       dfs_group(group, i, hits, visits);
0227     }
0228   }
0229 }; // namespace Jug::Reco
0230
0231 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0232 DECLARE_COMPONENT(ImagingTopoCluster)
0233
0234 } // namespace Jug::Reco