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