algorithms/calorimetry/ImagingTopoCluster.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Chao Peng, Sylvester Joosten, Whitney Armstrong, Wouter Deconinck
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  *  Original reference: https://arxiv.org/pdf/1603.02934.pdf
0010  *
0011  *  Modifications:
0012  *
0013  *  Wouter Deconinck (Manitoba), 08/24/2024
0014  *  - converted hit storage model from std::vector to std::set sorted on layer
0015  *    where only hits remaining to be assigned to a group are in the set
0016  *  - erase hits that are too low in energy to be part of a cluster
0017  *  - converted group storage model from std::set to std::list to allow adding
0018  *    hits while keeping iterators valid
0019  *
0020  */
0021
0022 #include "algorithms/calorimetry/ImagingTopoCluster.h"
0023
0024 #include <DD4hep/Handle.h>
0025 #include <Evaluator/DD4hepUnits.h>
0026 #include <edm4hep/Vector3f.h>
0027 #include <edm4hep/utils/vector_utils.h>
0028 #include <fmt/core.h>
0029 #include <cmath>
0030 #include <cstdlib>
0031 #include <gsl/pointers>
0032 #include <stdexcept>
0033 #include <utility>
0034 #include <variant>
0035 #include <vector>
0036
0037 #include "algorithms/calorimetry/ImagingTopoClusterConfig.h"
0038
0039 namespace eicrecon {
0040 template <typename... L> struct multilambda : L... {
0041   using L::operator()...;
0042   constexpr multilambda(L... lambda) : L(std::move(lambda))... {}
0043 };
0044
0045 void ImagingTopoCluster::init() {
0046
0047   multilambda _toDouble = {
0048       [](const std::string& v) { return dd4hep::_toDouble(v); },
0049       [](const double& v) { return v; },
0050   };
0051
0052   // unitless conversion
0053   // sanity checks
0054   if (m_cfg.minClusterCenterEdep < m_cfg.minClusterHitEdep) {
0055     const std::string msg =
0056         "minClusterCenterEdep must be greater than or equal to minClusterHitEdep";
0057     error(msg);
0058     throw std::runtime_error(msg);
0059   }
0060
0061   // using juggler internal units (GeV, dd4hep::mm, dd4hep::ns, dd4hep::rad)
0062   sameLayerDistXY[0]     = std::visit(_toDouble, m_cfg.sameLayerDistXY[0]) / dd4hep::mm;
0063   sameLayerDistXY[1]     = std::visit(_toDouble, m_cfg.sameLayerDistXY[1]) / dd4hep::mm;
0064   diffLayerDistXY[0]     = std::visit(_toDouble, m_cfg.diffLayerDistXY[0]) / dd4hep::mm;
0065   diffLayerDistXY[1]     = std::visit(_toDouble, m_cfg.diffLayerDistXY[1]) / dd4hep::mm;
0066   sameLayerDistEtaPhi[0] = m_cfg.sameLayerDistEtaPhi[0];
0067   sameLayerDistEtaPhi[1] = m_cfg.sameLayerDistEtaPhi[1] / dd4hep::rad;
0068   diffLayerDistEtaPhi[0] = m_cfg.diffLayerDistEtaPhi[0];
0069   diffLayerDistEtaPhi[1] = m_cfg.diffLayerDistEtaPhi[1] / dd4hep::rad;
0070   sameLayerDistTZ[0]     = m_cfg.sameLayerDistTZ[0] / dd4hep::mm;
0071   sameLayerDistTZ[1]     = m_cfg.sameLayerDistTZ[1] / dd4hep::mm;
0072   diffLayerDistTZ[0]     = m_cfg.diffLayerDistTZ[0] / dd4hep::mm;
0073   diffLayerDistTZ[1]     = m_cfg.diffLayerDistTZ[1] / dd4hep::mm;
0074
0075   sectorDist           = m_cfg.sectorDist / dd4hep::mm;
0076   minClusterHitEdep    = m_cfg.minClusterHitEdep / dd4hep::GeV;
0077   minClusterCenterEdep = m_cfg.minClusterCenterEdep / dd4hep::GeV;
0078   minClusterEdep       = m_cfg.minClusterEdep / dd4hep::GeV;
0079
0080   // same layer clustering parameters
0081   switch (m_cfg.sameLayerMode) {
0082   case ImagingTopoClusterConfig::ELayerMode::xy:
0083     if (m_cfg.sameLayerDistXY.size() != 2) {
0084       const std::string msg = "Expected 2 values (x_dist, y_dist) for sameLayerDistXY";
0085       error(msg);
0086       throw std::runtime_error(msg);
0087     }
0088     info("Same-layer clustering (same sector and same layer): "
0089          "Local [x, y] distance between hits <= [{:.4f} mm, {:.4f} mm].",
0090          sameLayerDistXY[0], sameLayerDistXY[1]);
0091     break;
0092   case ImagingTopoClusterConfig::ELayerMode::etaphi:
0093     if (m_cfg.sameLayerDistEtaPhi.size() != 2) {
0094       const std::string msg = "Expected 2 values (eta_dist, phi_dist) for sameLayerDistEtaPhi";
0095       error(msg);
0096       throw std::runtime_error(msg);
0097     }
0098     info("Same-layer clustering (same sector and same layer): "
0099          "Global [eta, phi] distance between hits <= [{:.4f}, {:.4f} rad].",
0100          sameLayerDistEtaPhi[0], sameLayerDistEtaPhi[1]);
0101     break;
0102   case ImagingTopoClusterConfig::ELayerMode::tz:
0103     if (m_cfg.sameLayerDistTZ.size() != 2) {
0104       const std::string msg = "Expected 2 values (t_dist, z_dist) for sameLayerDistTZ";
0105       error(msg);
0106       throw std::runtime_error(msg);
0107     }
0108     info("Same-layer clustering (same sector and same layer): "
0109          "Global [t, z] distance between hits <= [{:.4f} mm, {:.4f} mm].",
0110          sameLayerDistTZ[0], sameLayerDistTZ[1]);
0111     break;
0112   default:
0113     throw std::runtime_error("Unknown same-layer mode.");
0114   }
0115
0116   // different layer clustering parameters
0117   switch (m_cfg.diffLayerMode) {
0118   case ImagingTopoClusterConfig::ELayerMode::etaphi:
0119     if (m_cfg.diffLayerDistEtaPhi.size() != 2) {
0120       const std::string msg = "Expected 2 values (eta_dist, phi_dist) for diffLayerDistEtaPhi";
0121       error(msg);
0122       throw std::runtime_error(msg);
0123     }
0124     info("Neighbour layers clustering (same sector and layer id within +- {:d}): "
0125          "Global [eta, phi] distance between hits <= [{:.4f}, {:.4f} rad].",
0126          m_cfg.neighbourLayersRange, diffLayerDistEtaPhi[0], diffLayerDistEtaPhi[1]);
0127     break;
0128   case ImagingTopoClusterConfig::ELayerMode::xy:
0129     if (m_cfg.diffLayerDistXY.size() != 2) {
0130       const std::string msg = "Expected 2 values (x_dist, y_dist) for diffLayerDistXY";
0131       error(msg);
0132       throw std::runtime_error(msg);
0133     }
0134     info("Neighbour layers clustering (same sector and layer id within +- {:d}): "
0135          "Global [x, y] distance between hits <= [{:.4f} mm, {:.4f} mm].",
0136          m_cfg.neighbourLayersRange, diffLayerDistXY[0], diffLayerDistXY[1]);
0137     break;
0138   case ImagingTopoClusterConfig::ELayerMode::tz:
0139     if (m_cfg.diffLayerDistTZ.size() != 2) {
0140       const std::string msg = "Expected 2 values (t_dist, z_dist) for diffLayerDistTZ";
0141       error(msg);
0142       throw std::runtime_error(msg);
0143     }
0144     info("Neighbour layers clustering (same sector and layer id within +- {:d}): "
0145          "Global [t, z] distance between hits <= [{:.4f} mm, {:.4f} mm].",
0146          m_cfg.neighbourLayersRange, diffLayerDistTZ[0], diffLayerDistTZ[1]);
0147     break;
0148   default:
0149     error("Unknown different-layer mode.");
0150     throw std::runtime_error("Unknown different-layer mode.");
0151   }
0152   info("Neighbour sectors clustering (different sector): "
0153        "Global distance between hits <= {:.4f} mm.",
0154        sectorDist);
0155 }
0156
0157 void ImagingTopoCluster::process(const Input& input, const Output& output) const {
0158
0159   const auto [hits] = input;
0160   auto [proto]      = output;
0161
0162   // Sort hit indices (podio collections do not support std::sort)
0163   auto compare = [&hits](const auto& a, const auto& b) {
0164     // if !(a < b) and !(b < a), then a and b are equivalent
0165     // and only one of them will be allowed in a set
0166     if ((*hits)[a].getLayer() == (*hits)[b].getLayer()) {
0167       return (*hits)[a].getObjectID().index < (*hits)[b].getObjectID().index;
0168     }
0169     return (*hits)[a].getLayer() < (*hits)[b].getLayer();
0170   };
0171   // indices contains the remaining hit indices that have not
0172   // been assigned to a group yet
0173   std::set<std::size_t, decltype(compare)> indices(compare);
0174   // set does not have a size yet, so cannot fill with iota
0175   for (std::size_t i = 0; i < hits->size(); ++i) {
0176     indices.insert(i);
0177   }
0178   // ensure no hits were dropped due to equivalency in set
0179   if (hits->size() != indices.size()) {
0180     error("equivalent hits were dropped: #hits {:d}, #indices {:d}", hits->size(), indices.size());
0181   }
0182
0183   // Group neighbouring hits
0184   std::vector<std::list<std::size_t>> groups;
0185   // because indices changes, the loop over indices requires some care:
0186   // - we must use iterators instead of range-for
0187   // - erase returns an incremented iterator and therefore acts as idx++
0188   // - when the set becomes empty on erase, idx is invalid and idx++ will be too
0189   // (also applies to loop in bfs_group below)
0190   for (auto idx = indices.begin(); idx != indices.end();
0191        indices.empty() ? idx = indices.end() : idx) {
0192
0193     trace("hit {:d}: local position = ({}, {}, {}), global position = ({}, {}, {}), energy = {}",
0194           *idx, (*hits)[*idx].getLocal().x, (*hits)[*idx].getLocal().y, (*hits)[*idx].getLocal().z,
0195           (*hits)[*idx].getPosition().x, (*hits)[*idx].getPosition().y,
0196           (*hits)[*idx].getPosition().z, (*hits)[*idx].getEnergy());
0197
0198     // not energetic enough for cluster center, but could still be cluster hit
0199     if ((*hits)[*idx].getEnergy() < minClusterCenterEdep) {
0200       idx++;
0201       continue;
0202     }
0203
0204     // create a new group, and group all the neighbouring hits
0205     groups.emplace_back(std::list{*idx});
0206     bfs_group(*hits, indices, groups.back(), *idx);
0207
0208     // wait with erasing until after bfs_group to ensure iterator is not invalidated in bfs_group
0209     idx = indices.erase(idx); // takes role of idx++
0210   }
0211   debug("found {} potential clusters (groups of hits)", groups.size());
0212   for (std::size_t i = 0; i < groups.size(); ++i) {
0213     debug("group {}: {} hits", i, groups[i].size());
0214     for (auto idx : groups[i]) {
0215       const auto& hit = (*hits)[idx];
0216       trace("  hit {} -> energy = {:.6f}, layer = {}, sector = {}, local = ({:.2f}, {:.2f}, "
0217             "{:.2f}), global = ({:.2f}, {:.2f}, {:.2f})",
0218             idx, hit.getEnergy(), hit.getLayer(), hit.getSector(), hit.getLocal().x,
0219             hit.getLocal().y, hit.getLocal().z, hit.getPosition().x, hit.getPosition().y,
0220             hit.getPosition().z);
0221     }
0222   }
0223
0224   // form clusters
0225   for (const auto& group : groups) {
0226     if (group.size() < m_cfg.minClusterNhits) {
0227       continue;
0228     }
0229     double energy = 0.;
0230     for (std::size_t idx : group) {
0231       energy += (*hits)[idx].getEnergy();
0232     }
0233     if (energy < minClusterEdep) {
0234       continue;
0235     }
0236     auto pcl = proto->create();
0237     for (std::size_t idx : group) {
0238       pcl.addToHits((*hits)[idx]);
0239       pcl.addToWeights(1);
0240     }
0241   }
0242 }
0243
0244 // helper function to group hits
0245 bool ImagingTopoCluster::is_neighbour(const edm4eic::CalorimeterHit& h1,
0246                                       const edm4eic::CalorimeterHit& h2) const {
0247   // different sectors, simple distance check
0248   if (h1.getSector() != h2.getSector()) {
0249     return std::hypot((h1.getPosition().x - h2.getPosition().x),
0250                       (h1.getPosition().y - h2.getPosition().y),
0251                       (h1.getPosition().z - h2.getPosition().z)) <= sectorDist;
0252   }
0253
0254   // layer check
0255   int ldiff = std::abs(h1.getLayer() - h2.getLayer());
0256   // same layer, check local positions
0257   if (ldiff == 0) {
0258     switch (m_cfg.sameLayerMode) {
0259     case ImagingTopoClusterConfig::ELayerMode::xy:
0260       return (std::abs(h1.getLocal().x - h2.getLocal().x) <= sameLayerDistXY[0]) &&
0261              (std::abs(h1.getLocal().y - h2.getLocal().y) <= sameLayerDistXY[1]);
0262
0263     case ImagingTopoClusterConfig::ELayerMode::etaphi:
0264       return (std::abs(edm4hep::utils::eta(h1.getPosition()) -
0265                        edm4hep::utils::eta(h2.getPosition())) <= sameLayerDistEtaPhi[0]) &&
0266              (std::abs(edm4hep::utils::angleAzimuthal(h1.getPosition()) -
0267                        edm4hep::utils::angleAzimuthal(h2.getPosition())) <= sameLayerDistEtaPhi[1]);
0268
0269     case ImagingTopoClusterConfig::ELayerMode::tz: {
0270       // Layer mode 'tz' uses the average phi of the hits to define a rotated direction. The coordinate is a distance, not an angle.
0271       auto phi  = 0.5 * (edm4hep::utils::angleAzimuthal(h1.getPosition()) +
0272                         edm4hep::utils::angleAzimuthal(h2.getPosition()));
0273       auto h1_t = (h1.getPosition().x * sin(phi)) - (h1.getPosition().y * cos(phi));
0274       auto h2_t = (h2.getPosition().x * sin(phi)) - (h2.getPosition().y * cos(phi));
0275       auto h1_z = h1.getPosition().z;
0276       auto h2_z = h2.getPosition().z;
0277
0278       return (std::abs(h1_t - h2_t) <= sameLayerDistTZ[0]) &&
0279              (std::abs(h1_z - h2_z) <= sameLayerDistTZ[1]);
0280     }
0281
0282     default:
0283       error("Unknown layer mode for same-layer clustering.");
0284       return false;
0285     }
0286   } else if (ldiff <= m_cfg.neighbourLayersRange) {
0287     switch (m_cfg.diffLayerMode) {
0288     case eicrecon::ImagingTopoClusterConfig::ELayerMode::etaphi:
0289       return (std::abs(edm4hep::utils::eta(h1.getPosition()) -
0290                        edm4hep::utils::eta(h2.getPosition())) <= diffLayerDistEtaPhi[0]) &&
0291              (std::abs(edm4hep::utils::angleAzimuthal(h1.getPosition()) -
0292                        edm4hep::utils::angleAzimuthal(h2.getPosition())) <= diffLayerDistEtaPhi[1]);
0293
0294     case eicrecon::ImagingTopoClusterConfig::ELayerMode::xy:
0295       // Here, the xy layer mode is based on global XY positions rather than local XY positions, and thus it only works for endcap detectors.
0296       return (std::abs(h1.getPosition().x - h2.getPosition().x) <= diffLayerDistXY[0]) &&
0297              (std::abs(h1.getPosition().y - h2.getPosition().y) <= diffLayerDistXY[1]);
0298
0299     case eicrecon::ImagingTopoClusterConfig::ELayerMode::tz: {
0300       auto phi  = 0.5 * (edm4hep::utils::angleAzimuthal(h1.getPosition()) +
0301                         edm4hep::utils::angleAzimuthal(h2.getPosition()));
0302       auto h1_t = (h1.getPosition().x * sin(phi)) - (h1.getPosition().y * cos(phi));
0303       auto h2_t = (h2.getPosition().x * sin(phi)) - (h2.getPosition().y * cos(phi));
0304       auto h1_z = h1.getPosition().z;
0305       auto h2_z = h2.getPosition().z;
0306
0307       return (std::abs(h1_t - h2_t) <= diffLayerDistTZ[0]) &&
0308              (std::abs(h1_z - h2_z) <= diffLayerDistTZ[1]);
0309     }
0310
0311     default:
0312       error("Unknown layer mode for different-layer clustering.");
0313       return false;
0314     }
0315   }
0316
0317   // not in adjacent layers
0318   return false;
0319 }
0320
0321 } // namespace eicrecon