Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2024-06-17 07:05:49

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Chao, Chao Peng, Wouter Deconinck, Jihee Kim, Whitney Armstrong
0003 
0004 /*
0005  *  Island Clustering Algorithm for Calorimeter Blocks
0006  *  1. group all the adjacent modules
0007  *  2. split the groups between their local maxima with the energy deposit above <minClusterCenterEdep>
0008  *
0009  *  Author: Chao Peng (ANL), 09/27/2020
0010  *  References:
0011  *      https://cds.cern.ch/record/687345/files/note01_034.pdf
0012  *      https://www.jlab.org/primex/weekly_meetings/primexII/slides_2012_01_20/island_algorithm.pdf
0013  */
0014 #include <algorithm>
0015 #include <functional>
0016 #include <sstream>
0017 #include <tuple>
0018 
0019 #include "fmt/format.h"
0020 
0021 #include "Gaudi/Property.h"
0022 #include "GaudiAlg/GaudiAlgorithm.h"
0023 #include "GaudiAlg/GaudiTool.h"
0024 #include "GaudiAlg/Transformer.h"
0025 #include "GaudiKernel/PhysicalConstants.h"
0026 #include "GaudiKernel/RndmGenerators.h"
0027 #include "GaudiKernel/ToolHandle.h"
0028 
0029 #include "DDRec/CellIDPositionConverter.h"
0030 #include "DDRec/Surface.h"
0031 #include "DDRec/SurfaceManager.h"
0032 
0033 #include "JugBase/DataHandle.h"
0034 #include "JugBase/IGeoSvc.h"
0035 
0036 // Event Model related classes
0037 #include "edm4eic/CalorimeterHitCollection.h"
0038 #include "edm4eic/ClusterCollection.h"
0039 #include "edm4eic/ProtoClusterCollection.h"
0040 #include "edm4hep/utils/vector_utils.h"
0041 #include "edm4hep/Vector3f.h"
0042 #include "edm4hep/Vector2f.h"
0043 
0044 // Expression evaluator from DD4hep
0045 #include <Evaluator/Evaluator.h>
0046 #include <Evaluator/detail/Evaluator.h>
0047 
0048 #if defined __has_include
0049 #  if __has_include ("edm4eic/Vector3f.h")
0050 #    include "edm4eic/Vector3f.h"
0051 #  endif
0052 #  if __has_include ("edm4eic/Vector2f.h")
0053 #    include "edm4eic/Vector2f.h"
0054 #  endif
0055 #endif
0056 
0057 namespace edm4eic {
0058   class Vector2f;
0059   class Vector3f;
0060 }
0061 
0062 using namespace Gaudi::Units;
0063 
0064 namespace {
0065 
0066 using CaloHit = edm4eic::CalorimeterHit;
0067 using CaloHitCollection = edm4eic::CalorimeterHitCollection;
0068 
0069 using Vector2f = std::conditional_t<
0070   std::is_same_v<decltype(edm4eic::CalorimeterHitData::position), edm4hep::Vector3f>,
0071   edm4hep::Vector2f,
0072   edm4eic::Vector2f
0073 >;
0074 
0075 // helper functions to get distance between hits
0076 static Vector2f localDistXY(const CaloHit& h1, const CaloHit& h2) {
0077   const auto delta = h1.getLocal() - h2.getLocal();
0078   return {delta.x, delta.y};
0079 }
0080 static Vector2f localDistXZ(const CaloHit& h1, const CaloHit& h2) {
0081   const auto delta = h1.getLocal() - h2.getLocal();
0082   return {delta.x, delta.z};
0083 }
0084 static Vector2f localDistYZ(const CaloHit& h1, const CaloHit& h2) {
0085   const auto delta = h1.getLocal() - h2.getLocal();
0086   return {delta.y, delta.z};
0087 }
0088 static Vector2f dimScaledLocalDistXY(const CaloHit& h1, const CaloHit& h2) {
0089   const auto delta = h1.getLocal() - h2.getLocal();
0090   const auto dimsum = h1.getDimension() + h2.getDimension();
0091   return {2 * delta.x / dimsum.x, 2 * delta.y / dimsum.y};
0092 }
0093 static Vector2f globalDistRPhi(const CaloHit& h1, const CaloHit& h2) {
0094   using vector_type = decltype(Vector2f::a);
0095   return {
0096     static_cast<vector_type>(
0097       edm4hep::utils::magnitude(h1.getPosition()) - edm4hep::utils::magnitude(h2.getPosition())
0098     ),
0099     static_cast<vector_type>(
0100       edm4hep::utils::angleAzimuthal(h1.getPosition()) - edm4hep::utils::angleAzimuthal(h2.getPosition())
0101     )
0102   };
0103 }
0104 static Vector2f globalDistEtaPhi(const CaloHit& h1,
0105                                        const CaloHit& h2) {
0106   using vector_type = decltype(Vector2f::a);
0107   return {
0108     static_cast<vector_type>(
0109       edm4hep::utils::eta(h1.getPosition()) - edm4hep::utils::eta(h2.getPosition())
0110     ),
0111     static_cast<vector_type>(
0112       edm4hep::utils::angleAzimuthal(h1.getPosition()) - edm4hep::utils::angleAzimuthal(h2.getPosition())
0113     )
0114   };
0115 }
0116 // name: {method, units}
0117 static std::map<std::string,
0118                 std::tuple<std::function<Vector2f(const CaloHit&, const CaloHit&)>, std::vector<double>>>
0119     distMethods{
0120         {"localDistXY", {localDistXY, {mm, mm}}},        {"localDistXZ", {localDistXZ, {mm, mm}}},
0121         {"localDistYZ", {localDistYZ, {mm, mm}}},        {"dimScaledLocalDistXY", {dimScaledLocalDistXY, {1., 1.}}},
0122         {"globalDistRPhi", {globalDistRPhi, {mm, rad}}}, {"globalDistEtaPhi", {globalDistEtaPhi, {1., rad}}},
0123     };
0124 
0125 } // namespace
0126 namespace Jug::Reco {
0127 
0128 /**
0129  *  Island Clustering Algorithm for Calorimeter Blocks.
0130  *
0131  *  1. group all the adjacent modules
0132  *  2. split the groups between their local maxima with the energy deposit above <minClusterCenterEdep>
0133  *
0134  *  References:
0135  *      https://cds.cern.ch/record/687345/files/note01_034.pdf
0136  *      https://www.jlab.org/primex/weekly_meetings/primexII/slides_2012_01_20/island_algorithm.pdf
0137  *
0138  * \ingroup reco
0139  */
0140 class CalorimeterIslandCluster : public GaudiAlgorithm {
0141 private:
0142   Gaudi::Property<bool> m_splitCluster{this, "splitCluster", true};
0143   Gaudi::Property<double> m_minClusterHitEdep{this, "minClusterHitEdep", 0.};
0144   Gaudi::Property<double> m_minClusterCenterEdep{this, "minClusterCenterEdep", 50.0 * MeV};
0145   DataHandle<CaloHitCollection> m_inputHitCollection{"inputHitCollection", Gaudi::DataHandle::Reader, this};
0146   DataHandle<edm4eic::ProtoClusterCollection> m_outputProtoCollection{"outputProtoClusterCollection",
0147                                                                   Gaudi::DataHandle::Writer, this};
0148 
0149   Gaudi::Property<std::string> m_geoSvcName{this, "geoServiceName", "GeoSvc"};
0150   Gaudi::Property<std::string> m_readout{this, "readoutClass", ""};
0151   Gaudi::Property<std::string> u_adjacencyMatrix{this, "adjacencyMatrix", ""};
0152   // neighbour checking distances
0153   Gaudi::Property<double> m_sectorDist{this, "sectorDist", 5.0 * cm};
0154   Gaudi::Property<std::vector<double>> u_localDistXY{this, "localDistXY", {}};
0155   Gaudi::Property<std::vector<double>> u_localDistXZ{this, "localDistXZ", {}};
0156   Gaudi::Property<std::vector<double>> u_localDistYZ{this, "localDistYZ", {}};
0157   Gaudi::Property<std::vector<double>> u_globalDistRPhi{this, "globalDistRPhi", {}};
0158   Gaudi::Property<std::vector<double>> u_globalDistEtaPhi{this, "globalDistEtaPhi", {}};
0159   Gaudi::Property<std::vector<double>> u_dimScaledLocalDistXY{this, "dimScaledLocalDistXY", {1.8, 1.8}};
0160   // neighbor checking function
0161   std::function<Vector2f(const CaloHit&, const CaloHit&)> hitsDist;
0162 
0163   // helper function to group hits
0164   std::function<bool(const CaloHit& h1, const CaloHit& h2)> is_neighbour;
0165 
0166   // unitless counterparts of the input parameters
0167   double minClusterHitEdep{0}, minClusterCenterEdep{0}, sectorDist{0};
0168   std::array<double, 2> neighbourDist = {0., 0.};
0169 
0170   /// Pointer to the geometry service
0171   SmartIF<IGeoSvc> m_geoSvc;
0172   dd4hep::IDDescriptor m_idSpec;
0173 
0174 public:
0175   CalorimeterIslandCluster(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0176     declareProperty("inputHitCollection", m_inputHitCollection, "");
0177     declareProperty("outputProtoClusterCollection", m_outputProtoCollection, "");
0178   }
0179 
0180   StatusCode initialize() override {
0181     if (GaudiAlgorithm::initialize().isFailure()) {
0182       return StatusCode::FAILURE;
0183     }
0184 
0185     // unitless conversion, keep consistency with juggler internal units (GeV, mm, ns, rad)
0186     minClusterHitEdep    = m_minClusterHitEdep.value() / GeV;
0187     minClusterCenterEdep = m_minClusterCenterEdep.value() / GeV;
0188     sectorDist           = m_sectorDist.value() / mm;
0189 
0190     // set coordinate system
0191     auto set_dist_method = [this](const Gaudi::Property<std::vector<double>>& uprop) {
0192       if (uprop.size() == 0) {
0193         return false;
0194       }
0195       auto& [method, units] = distMethods[uprop.name()];
0196       if (uprop.size() != units.size()) {
0197         info() << units.size() << endmsg;
0198         warning() << fmt::format("Expect {} values from {}, received {}: ({}), ignored it.", units.size(), uprop.name(),
0199                                  uprop.size(), fmt::join(uprop.value(), ", "))
0200                   << endmsg;
0201         return false;
0202       } else {
0203         for (size_t i = 0; i < units.size(); ++i) {
0204           neighbourDist[i] = uprop.value()[i] / units[i];
0205         }
0206         hitsDist = method;
0207         info() << fmt::format("Clustering uses {} with distances <= [{}]", uprop.name(), fmt::join(neighbourDist, ","))
0208                << endmsg;
0209       }
0210       return true;
0211     };
0212 
0213     std::vector<Gaudi::Property<std::vector<double>>> uprops{
0214         u_localDistXY,
0215         u_localDistXZ,
0216         u_localDistYZ,
0217         u_globalDistRPhi,
0218         u_globalDistEtaPhi,
0219         // default one should be the last one
0220         u_dimScaledLocalDistXY,
0221     };
0222 
0223     bool method_found = false;
0224     if (u_adjacencyMatrix.value() != "") {
0225       m_geoSvc = service(m_geoSvcName);
0226       // sanity checks
0227       if (!m_geoSvc) {
0228         error() << "Unable to locate Geometry Service. "
0229                 << "Make sure you have GeoSvc and SimSvc in the right order in the configuration."
0230                 << endmsg;
0231         return StatusCode::FAILURE;
0232       }
0233       if (m_readout.value().empty()) {
0234         error() << "readoutClass is not provided, it is needed to know the fields in readout ids"
0235                 << endmsg;
0236       }
0237       m_idSpec = m_geoSvc->detector()->readout(m_readout).idSpec();
0238 
0239       is_neighbour = [this](const CaloHit& h1, const CaloHit& h2) {
0240         dd4hep::tools::Evaluator::Object evaluator;
0241         for(const auto &p : m_idSpec.fields()) {
0242           const std::string &name = p.first;
0243           const dd4hep::IDDescriptor::Field* field = p.second;
0244           evaluator.setVariable((name + "_1").c_str(), field->value(h1.getCellID()));
0245           evaluator.setVariable((name + "_2").c_str(), field->value(h2.getCellID()));
0246           debug() << "setVariable(\"" << name  << "_1\", " << field->value(h1.getCellID()) << ");" << endmsg;
0247           debug() << "setVariable(\"" << name  << "_2\", " << field->value(h2.getCellID()) << ");" << endmsg;
0248         }
0249         dd4hep::tools::Evaluator::Object::EvalStatus eval = evaluator.evaluate(u_adjacencyMatrix.value().c_str());
0250         if (eval.status()) {
0251           // no known conversion from 'MsgStream' to 'std::ostream &'
0252           std::stringstream sstr;
0253           eval.print_error(sstr);
0254           error() << sstr.str() << endmsg;
0255         }
0256         debug() << "result = " << eval.result() << endmsg;
0257         return eval.result();
0258       };
0259       method_found = true;
0260     }
0261     if (!method_found)
0262     {
0263       for (auto& uprop : uprops) {
0264         if (set_dist_method(uprop)) {
0265           method_found = true;
0266 
0267           is_neighbour = [this](const CaloHit& h1, const CaloHit& h2) {
0268             // in the same sector
0269             if (h1.getSector() == h2.getSector()) {
0270               auto dist = hitsDist(h1, h2);
0271               return (dist.a <= neighbourDist[0]) && (dist.b <= neighbourDist[1]);
0272               // different sector, local coordinates do not work, using global coordinates
0273             } else {
0274               // sector may have rotation (barrel), so z is included
0275               return (edm4hep::utils::magnitude(h1.getPosition() - h2.getPosition()) <= sectorDist);
0276             }
0277           };
0278 
0279           break;
0280         }
0281       }
0282     }
0283     if (not method_found) {
0284       error() << "Cannot determine the clustering coordinates" << endmsg;
0285       return StatusCode::FAILURE;
0286     }
0287 
0288     return StatusCode::SUCCESS;
0289   }
0290 
0291   StatusCode execute() override {
0292     // input collections
0293     const auto& hits = *(m_inputHitCollection.get());
0294     // Create output collections
0295     auto& proto = *(m_outputProtoCollection.createAndPut());
0296 
0297     // group neighboring hits
0298     std::vector<std::vector<std::pair<uint32_t, CaloHit>>> groups;
0299 
0300     std::vector<bool> visits(hits.size(), false);
0301     for (size_t i = 0; i < hits.size(); ++i) {
0302       if (msgLevel(MSG::DEBUG)) {
0303         const auto& hit = hits[i];
0304         debug() << fmt::format("hit {:d}: energy = {:.4f} MeV, local = ({:.4f}, {:.4f}) mm, "
0305                                "global=({:.4f}, {:.4f}, {:.4f}) mm",
0306                                i, hit.getEnergy() * 1000., hit.getLocal().x, hit.getLocal().y, hit.getPosition().x,
0307                                hit.getPosition().y, hit.getPosition().z)
0308                 << endmsg;
0309       }
0310       // already in a group
0311       if (visits[i]) {
0312         continue;
0313       }
0314       groups.emplace_back();
0315       // create a new group, and group all the neighboring hits
0316       dfs_group(groups.back(), i, hits, visits);
0317     }
0318 
0319     for (auto& group : groups) {
0320       if (group.empty()) {
0321         continue;
0322       }
0323       auto maxima = find_maxima(group, !m_splitCluster.value());
0324       split_group(group, maxima, proto);
0325       if (msgLevel(MSG::DEBUG)) {
0326         debug() << "hits in a group: " << group.size() << ", "
0327                 << "local maxima: " << maxima.size() << endmsg;
0328       }
0329     }
0330 
0331     return StatusCode::SUCCESS;
0332   }
0333 
0334 private:
0335   // grouping function with Depth-First Search
0336   void dfs_group(std::vector<std::pair<uint32_t, CaloHit>>& group, int idx,
0337                  const CaloHitCollection& hits, std::vector<bool>& visits) const {
0338     // not a qualified hit to particpate clustering, stop here
0339     if (hits[idx].getEnergy() < minClusterHitEdep) {
0340       visits[idx] = true;
0341       return;
0342     }
0343 
0344     group.emplace_back(idx, hits[idx]);
0345     visits[idx] = true;
0346     for (size_t i = 0; i < hits.size(); ++i) {
0347       if (visits[i] || !is_neighbour(hits[idx], hits[i])) {
0348         continue;
0349       }
0350       dfs_group(group, i, hits, visits);
0351     }
0352   }
0353 
0354   // find local maxima that above a certain threshold
0355   std::vector<CaloHit>
0356   find_maxima(const std::vector<std::pair<uint32_t, CaloHit>>& group,
0357               bool global = false) const {
0358     std::vector<CaloHit> maxima;
0359     if (group.empty()) {
0360       return maxima;
0361     }
0362 
0363     if (global) {
0364       int mpos = 0;
0365       for (size_t i = 0; i < group.size(); ++i) {
0366         if (group[mpos].second.getEnergy() < group[i].second.getEnergy()) {
0367           mpos = i;
0368         }
0369       }
0370       if (group[mpos].second.getEnergy() >= minClusterCenterEdep) {
0371         maxima.push_back(group[mpos].second);
0372       }
0373       return maxima;
0374     }
0375 
0376     for (const auto& [idx, hit] : group) {
0377       // not a qualified center
0378       if (hit.getEnergy() < minClusterCenterEdep) {
0379         continue;
0380       }
0381 
0382       bool maximum = true;
0383       for (const auto& [idx2, hit2] : group) {
0384         if (hit == hit2) {
0385           continue;
0386         }
0387 
0388         if (is_neighbour(hit, hit2) && hit2.getEnergy() > hit.getEnergy()) {
0389           maximum = false;
0390           break;
0391         }
0392       }
0393 
0394       if (maximum) {
0395         maxima.push_back(hit);
0396       }
0397     }
0398 
0399     return maxima;
0400   }
0401 
0402   // helper function
0403   inline static void vec_normalize(std::vector<double>& vals) {
0404     double total = 0.;
0405     for (auto& val : vals) {
0406       total += val;
0407     }
0408     for (auto& val : vals) {
0409       val /= total;
0410     }
0411   }
0412 
0413   // split a group of hits according to the local maxima
0414   void split_group(std::vector<std::pair<uint32_t, CaloHit>>& group, const std::vector<CaloHit>& maxima,
0415                    edm4eic::ProtoClusterCollection& proto) const {
0416     // special cases
0417     if (maxima.empty()) {
0418       if (msgLevel(MSG::VERBOSE)) {
0419         verbose() << "No maxima found, not building any clusters" << endmsg;
0420       }
0421       return;
0422     } else if (maxima.size() == 1) {
0423       edm4eic::MutableProtoCluster pcl;
0424       for (auto& [idx, hit] : group) {
0425         pcl.addToHits(hit);
0426         pcl.addToWeights(1.);
0427       }
0428       proto.push_back(pcl);
0429       if (msgLevel(MSG::VERBOSE)) {
0430         verbose() << "A single maximum found, added one ProtoCluster" << endmsg;
0431       }
0432       return;
0433     }
0434 
0435     // split between maxima
0436     // TODO, here we can implement iterations with profile, or even ML for better splits
0437     std::vector<double> weights(maxima.size(), 1.);
0438     std::vector<edm4eic::MutableProtoCluster> pcls;
0439     for (size_t k = 0; k < maxima.size(); ++k) {
0440       pcls.emplace_back();
0441     }
0442 
0443     size_t i = 0;
0444     for (const auto& [idx, hit] : group) {
0445       size_t j = 0;
0446       // calculate weights for local maxima
0447       for (const auto& chit : maxima) {
0448         double dist_ref = chit.getDimension().x;
0449         double energy   = chit.getEnergy();
0450         double dist     = edm4hep::utils::magnitude(hitsDist(chit, hit));
0451         weights[j]      = std::exp(-dist / dist_ref) * energy;
0452         j += 1;
0453       }
0454 
0455       // normalize weights
0456       vec_normalize(weights);
0457 
0458       // ignore small weights
0459       for (auto& w : weights) {
0460         if (w < 0.02) {
0461           w = 0;
0462         }
0463       }
0464       vec_normalize(weights);
0465 
0466       // split energy between local maxima
0467       for (size_t k = 0; k < maxima.size(); ++k) {
0468         double weight = weights[k];
0469         if (weight <= 1e-6) {
0470           continue;
0471         }
0472         pcls[k].addToHits(hit);
0473         pcls[k].addToWeights(weight);
0474       }
0475       i += 1;
0476     }
0477     for (auto& pcl : pcls) {
0478       proto.push_back(pcl);
0479     }
0480     if (msgLevel(MSG::VERBOSE)) {
0481       verbose() << "Multiple (" << maxima.size() << ") maxima found, added a ProtoClusters for each maximum" << endmsg;
0482     }
0483   }
0484 };
0485 
0486 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0487 DECLARE_COMPONENT(CalorimeterIslandCluster)
0488 
0489 } // namespace Jug::Reco