algorithms/calorimetry/CalorimeterIslandCluster.cc

0001 // Copyright (C) 2022, 2023 Chao Peng, Wouter Deconinck, Sylvester Joosten, Dmitry Kalinkin, David Lawrence
0002 // SPDX-License-Identifier: LGPL-3.0-or-later
0003
0004 // References:
0005 //   https://cds.cern.ch/record/687345/files/note01_034.pdf
0006 //   https://www.jlab.org/primex/weekly_meetings/primexII/slides_2012_01_20/island_algorithm.pdf
0007
0008 #include <DD4hep/Handle.h>
0009 #include <DD4hep/Readout.h>
0010 #include <Evaluator/DD4hepUnits.h>
0011 #include <algorithms/service.h>
0012 #include <edm4hep/Vector2f.h>
0013 #include <edm4hep/Vector3f.h>
0014 #include <edm4hep/utils/vector_utils.h>
0015 #include <fmt/ranges.h>
0016 #include <cmath>
0017 #include <format>
0018 #include <iterator>
0019 #include <map>
0020 #include <ranges>
0021 #include <set>
0022 #include <stdexcept>
0023 #include <string>
0024 #include <tuple>
0025 #include <unordered_map>
0026 #include <utility>
0027 #include <variant>
0028 #include <vector>
0029
0030 #include "CalorimeterIslandCluster.h"
0031 #include "algorithms/calorimetry/CalorimeterIslandClusterConfig.h"
0032 #include "algorithms/interfaces/detail/multilambda.h"
0033 #include "services/evaluator/EvaluatorSvc.h"
0034
0035 using namespace edm4eic;
0036
0037 namespace eicrecon {
0038
0039 static double Phi_mpi_pi(double phi) { return std::remainder(phi, 2 * M_PI); }
0040
0041 static edm4hep::Vector2f localDistXY(const CaloHit& h1, const CaloHit& h2) {
0042   const auto delta = h1.getLocal() - h2.getLocal();
0043   return {delta.x, delta.y};
0044 }
0045 static edm4hep::Vector2f localDistXZ(const CaloHit& h1, const CaloHit& h2) {
0046   const auto delta = h1.getLocal() - h2.getLocal();
0047   return {delta.x, delta.z};
0048 }
0049 static edm4hep::Vector2f localDistYZ(const CaloHit& h1, const CaloHit& h2) {
0050   const auto delta = h1.getLocal() - h2.getLocal();
0051   return {delta.y, delta.z};
0052 }
0053 static edm4hep::Vector2f dimScaledLocalDistXY(const CaloHit& h1, const CaloHit& h2) {
0054   const auto delta = h1.getLocal() - h2.getLocal();
0055
0056   const auto dimsum = h1.getDimension() + h2.getDimension();
0057
0058   return {2 * delta.x / dimsum.x, 2 * delta.y / dimsum.y};
0059 }
0060 static edm4hep::Vector2f globalDistRPhi(const CaloHit& h1, const CaloHit& h2) {
0061   using vector_type = decltype(edm4hep::Vector2f::a);
0062   return {static_cast<vector_type>(edm4hep::utils::magnitude(h1.getPosition()) -
0063                                    edm4hep::utils::magnitude(h2.getPosition())),
0064           static_cast<vector_type>(Phi_mpi_pi(edm4hep::utils::angleAzimuthal(h1.getPosition()) -
0065                                               edm4hep::utils::angleAzimuthal(h2.getPosition())))};
0066 }
0067 static edm4hep::Vector2f globalDistEtaPhi(const CaloHit& h1, const CaloHit& h2) {
0068   using vector_type = decltype(edm4hep::Vector2f::a);
0069   return {static_cast<vector_type>(edm4hep::utils::eta(h1.getPosition()) -
0070                                    edm4hep::utils::eta(h2.getPosition())),
0071           static_cast<vector_type>(Phi_mpi_pi(edm4hep::utils::angleAzimuthal(h1.getPosition()) -
0072                                               edm4hep::utils::angleAzimuthal(h2.getPosition())))};
0073 }
0074
0075 //------------------------
0076 // AlgorithmInit
0077 //------------------------
0078 void CalorimeterIslandCluster::init() {
0079
0080   multilambda _toDouble = {
0081       [](const std::string& v) { return dd4hep::_toDouble(v); },
0082       [](const double& v) { return v; },
0083   };
0084
0085   if (m_cfg.localDistXY.size() == 2) {
0086     m_localDistXY.push_back(std::visit(_toDouble, m_cfg.localDistXY[0]));
0087     m_localDistXY.push_back(std::visit(_toDouble, m_cfg.localDistXY[1]));
0088   }
0089
0090   static std::map<std::string,
0091                   std::tuple<std::function<edm4hep::Vector2f(const CaloHit&, const CaloHit&)>,
0092                              std::vector<double>>>
0093       distMethods{{"localDistXY", {localDistXY, {dd4hep::mm, dd4hep::mm}}},
0094                   {"localDistXZ", {localDistXZ, {dd4hep::mm, dd4hep::mm}}},
0095                   {"localDistYZ", {localDistYZ, {dd4hep::mm, dd4hep::mm}}},
0096                   {"dimScaledLocalDistXY", {dimScaledLocalDistXY, {1., 1.}}},
0097                   {"globalDistRPhi", {globalDistRPhi, {dd4hep::mm, dd4hep::rad}}},
0098                   {"globalDistEtaPhi", {globalDistEtaPhi, {1., dd4hep::rad}}}};
0099
0100   // set coordinate system
0101   auto set_dist_method = [this](std::pair<std::string, std::vector<double>> uprop) {
0102     if (uprop.second.empty()) {
0103       return false;
0104     }
0105     auto& [method, units] = distMethods[uprop.first];
0106     if (uprop.second.size() != units.size()) {
0107       warning("Expect {} values from {}, received {}. ignored it.", units.size(), uprop.first,
0108               uprop.second.size());
0109       return false;
0110     }
0111     for (std::size_t i = 0; i < units.size(); ++i) {
0112       // NOLINTNEXTLINE(cppcoreguidelines-pro-bounds-constant-array-index)
0113       neighbourDist[i] = uprop.second[i] / units[i];
0114     }
0115     hitsDist = method;
0116     info("Clustering uses {} with distances <= [{}]", uprop.first, fmt::join(neighbourDist, ","));
0117
0118     return true;
0119   };
0120
0121   std::vector<std::pair<std::string, std::vector<double>>> uprops{
0122       {"localDistXY", m_localDistXY},
0123       {"localDistXZ", m_cfg.localDistXZ},
0124       {"localDistYZ", m_cfg.localDistYZ},
0125       {"globalDistRPhi", m_cfg.globalDistRPhi},
0126       {"globalDistEtaPhi", m_cfg.globalDistEtaPhi},
0127       // default one should be the last one
0128       {"dimScaledLocalDistXY", m_cfg.dimScaledLocalDistXY}};
0129
0130   auto& serviceSvc = algorithms::ServiceSvc::instance();
0131
0132   std::function hit_pair_to_map = [this](const edm4eic::CalorimeterHit& h1,
0133                                          const edm4eic::CalorimeterHit& h2) {
0134     std::unordered_map<std::string, double> params;
0135     for (const auto& p : m_idSpec.fields()) {
0136       const std::string& name                  = p.first;
0137       const dd4hep::IDDescriptor::Field* field = p.second;
0138       params.emplace(name + "_1", field->value(h1.getCellID()));
0139       params.emplace(name + "_2", field->value(h2.getCellID()));
0140       trace("{}_1 = {}", name, field->value(h1.getCellID()));
0141       trace("{}_2 = {}", name, field->value(h2.getCellID()));
0142     }
0143     return params;
0144   };
0145
0146   if (m_cfg.readout.empty()) {
0147     if ((!m_cfg.adjacencyMatrix.empty()) || (!m_cfg.peakNeighbourhoodMatrix.empty())) {
0148       throw std::runtime_error(
0149           "'readout' is not provided, it is needed to know the fields in readout ids");
0150     }
0151   } else {
0152     m_idSpec = m_detector->readout(m_cfg.readout).idSpec();
0153   }
0154
0155   bool method_found = false;
0156
0157   // Adjacency matrix methods
0158   if (!m_cfg.adjacencyMatrix.empty()) {
0159     is_neighbour = serviceSvc.service<EvaluatorSvc>("EvaluatorSvc")
0160                        ->compile(m_cfg.adjacencyMatrix, hit_pair_to_map);
0161     method_found = true;
0162   }
0163
0164   // Coordinate distance methods
0165   if (not method_found) {
0166     for (auto& uprop : uprops) {
0167       if (set_dist_method(uprop)) {
0168         method_found = true;
0169
0170         is_neighbour = [this](const CaloHit& h1, const CaloHit& h2) {
0171           // in the same sector
0172           if (h1.getSector() == h2.getSector()) {
0173             auto dist = hitsDist(h1, h2);
0174             return (std::abs(dist.a) <= neighbourDist[0]) && (std::abs(dist.b) <= neighbourDist[1]);
0175             // different sector, local coordinates do not work, using global coordinates
0176           } // sector may have rotation (barrel), so z is included
0177           // (EDM4hep units are mm, so convert sectorDist to mm)
0178           return (edm4hep::utils::magnitude(h1.getPosition() - h2.getPosition()) <=
0179                   m_cfg.sectorDist / dd4hep::mm);
0180         };
0181
0182         break;
0183       }
0184     }
0185   }
0186
0187   if (not method_found) {
0188     throw std::runtime_error("Cannot determine the clustering coordinates");
0189   }
0190
0191   if (m_cfg.splitCluster) {
0192     if (!m_cfg.peakNeighbourhoodMatrix.empty()) {
0193       is_maximum_neighbourhood = serviceSvc.service<EvaluatorSvc>("EvaluatorSvc")
0194                                      ->compile(m_cfg.peakNeighbourhoodMatrix, hit_pair_to_map);
0195     } else {
0196       is_maximum_neighbourhood = is_neighbour;
0197     }
0198
0199     auto transverseEnergyProfileMetric_it = std::ranges::find_if(
0200         distMethods, [&](auto& p) { return m_cfg.transverseEnergyProfileMetric == p.first; });
0201     if (transverseEnergyProfileMetric_it == distMethods.end()) {
0202       throw std::runtime_error(
0203           std::format(R"(Unsupported value "{}" for "transverseEnergyProfileMetric")",
0204                       m_cfg.transverseEnergyProfileMetric));
0205     }
0206     transverseEnergyProfileMetric = std::get<0>(transverseEnergyProfileMetric_it->second);
0207     std::vector<double>& units    = std::get<1>(transverseEnergyProfileMetric_it->second);
0208     for (auto unit : units) {
0209       if (unit != units[0]) {
0210         throw std::runtime_error(std::format("Metric {} has incompatible dimension units",
0211                                              m_cfg.transverseEnergyProfileMetric));
0212       }
0213     }
0214     transverseEnergyProfileScaleUnits = units[0];
0215   }
0216 }
0217
0218 void CalorimeterIslandCluster::process(const CalorimeterIslandCluster::Input& input,
0219                                        const CalorimeterIslandCluster::Output& output) const {
0220
0221   const auto [hits]     = input;
0222   auto [proto_clusters] = output;
0223
0224   // group neighboring hits
0225   std::vector<std::set<std::size_t>> groups;
0226
0227   std::vector<bool> visits(hits->size(), false);
0228   for (std::size_t i = 0; i < hits->size(); ++i) {
0229
0230     {
0231       const auto& hit = (*hits)[i];
0232       debug("hit {:d}: energy = {:.4f} MeV, local = ({:.4f}, {:.4f}) mm, global=({:.4f}, {:.4f}, "
0233             "{:.4f}) mm",
0234             i, hit.getEnergy() * 1000., hit.getLocal().x, hit.getLocal().y, hit.getPosition().x,
0235             hit.getPosition().y, hit.getPosition().z);
0236     }
0237     // already in a group
0238     if (visits[i]) {
0239       continue;
0240     }
0241     groups.emplace_back();
0242     // create a new group, and group all the neighboring hits
0243     bfs_group(*hits, groups.back(), i, visits);
0244   }
0245
0246   for (auto& group : groups) {
0247     if (group.empty()) {
0248       continue;
0249     }
0250     auto maxima = find_maxima(*hits, group, !m_cfg.splitCluster);
0251     split_group(*hits, group, maxima, proto_clusters);
0252
0253     debug("hits in a group: {}, local maxima: {}", group.size(), maxima.size());
0254   }
0255 }
0256
0257 // grouping function with Breadth-First Search
0258 void CalorimeterIslandCluster::bfs_group(const edm4eic::CalorimeterHitCollection& hits,
0259                                          std::set<std::size_t>& group, std::size_t idx,
0260                                          std::vector<bool>& visits) const {
0261   visits[idx] = true;
0262
0263   // not a qualified hit to participate clustering, stop here
0264   if (hits[idx].getEnergy() < m_cfg.minClusterHitEdep) {
0265     return;
0266   }
0267
0268   group.insert(idx);
0269   std::size_t prev_size = 0;
0270
0271   while (prev_size != group.size()) {
0272     prev_size = group.size();
0273     for (std::size_t idx1 : group) {
0274       // check neighbours
0275       for (std::size_t idx2 = 0; idx2 < hits.size(); ++idx2) {
0276         // not a qualified hit to participate clustering, skip
0277         if (hits[idx2].getEnergy() < m_cfg.minClusterHitEdep) {
0278           continue;
0279         }
0280         if ((!visits[idx2]) && is_neighbour(hits[idx1], hits[idx2])) {
0281           group.insert(idx2);
0282           visits[idx2] = true;
0283         }
0284       }
0285     }
0286   }
0287 }
0288
0289 // find local maxima that above a certain threshold
0290 std::vector<std::size_t>
0291 CalorimeterIslandCluster::find_maxima(const edm4eic::CalorimeterHitCollection& hits,
0292                                       const std::set<std::size_t>& group, bool global) const {
0293   std::vector<std::size_t> maxima;
0294   if (group.empty()) {
0295     return maxima;
0296   }
0297
0298   if (global) {
0299     std::size_t mpos = *group.begin();
0300     for (auto idx : group) {
0301       if (hits[mpos].getEnergy() < hits[idx].getEnergy()) {
0302         mpos = idx;
0303       }
0304     }
0305     if (hits[mpos].getEnergy() >= m_cfg.minClusterCenterEdep) {
0306       maxima.push_back(mpos);
0307     }
0308     return maxima;
0309   }
0310
0311   for (std::size_t idx1 : group) {
0312     // not a qualified center
0313     if (hits[idx1].getEnergy() < m_cfg.minClusterCenterEdep) {
0314       continue;
0315     }
0316
0317     bool maximum = true;
0318     for (std::size_t idx2 : group) {
0319       if (idx1 == idx2) {
0320         continue;
0321       }
0322
0323       if (is_maximum_neighbourhood(hits[idx1], hits[idx2]) &&
0324           (hits[idx2].getEnergy() > hits[idx1].getEnergy())) {
0325         maximum = false;
0326         break;
0327       }
0328     }
0329
0330     if (maximum) {
0331       maxima.push_back(idx1);
0332     }
0333   }
0334
0335   return maxima;
0336 }
0337
0338 // split a group of hits according to the local maxima
0339 //TODO: confirm protoclustering without protoclustercollection
0340 void CalorimeterIslandCluster::split_group(const edm4eic::CalorimeterHitCollection& hits,
0341                                            std::set<std::size_t>& group,
0342                                            const std::vector<std::size_t>& maxima,
0343                                            edm4eic::ProtoClusterCollection* protoClusters) const {
0344   // special cases
0345   if (maxima.empty()) {
0346     debug("No maxima found, not building any clusters");
0347     return;
0348   } else if (maxima.size() == 1) {
0349     edm4eic::MutableProtoCluster pcl = protoClusters->create();
0350     for (std::size_t idx : group) {
0351       pcl.addToHits(hits[idx]);
0352       pcl.addToWeights(1.);
0353     }
0354
0355     debug("A single maximum found, added one ProtoCluster");
0356
0357     return;
0358   }
0359
0360   // split between maxima
0361   // TODO, here we can implement iterations with profile, or even ML for better splits
0362   std::vector<double> weights(maxima.size(), 1.);
0363   std::vector<edm4eic::MutableProtoCluster> pcls;
0364   for (std::size_t k = 0; k < maxima.size(); ++k) {
0365     pcls.push_back(protoClusters->create());
0366   }
0367
0368   for (std::size_t idx : group) {
0369     std::size_t j = 0;
0370     // calculate weights for local maxima
0371     for (std::size_t cidx : maxima) {
0372       double energy = hits[cidx].getEnergy();
0373       double dist = edm4hep::utils::magnitude(transverseEnergyProfileMetric(hits[cidx], hits[idx]));
0374       weights[j] =
0375           std::exp(-dist * transverseEnergyProfileScaleUnits / m_cfg.transverseEnergyProfileScale) *
0376           energy;
0377       j += 1;
0378     }
0379
0380     // normalize weights
0381     vec_normalize(weights);
0382
0383     // ignore small weights
0384     for (auto& w : weights) {
0385       if (w < 0.02) {
0386         w = 0;
0387       }
0388     }
0389     vec_normalize(weights);
0390
0391     // split energy between local maxima
0392     for (std::size_t k = 0; k < maxima.size(); ++k) {
0393       double weight = weights[k];
0394       if (weight <= 1e-6) {
0395         continue;
0396       }
0397       pcls[k].addToHits(hits[idx]);
0398       pcls[k].addToWeights(weight);
0399     }
0400   }
0401   debug("Multiple ({}) maxima found, added a ProtoClusters for each maximum", maxima.size());
0402 }
0403
0404 } // namespace eicrecon