Skip to content
Snippets Groups Projects
ImagingTopoCluster.cpp 8.26 KiB
Newer Older
  • Learn to ignore specific revisions
  • Chao Peng's avatar
    Chao Peng committed
    /*
     *  Topological Cell Clustering Algorithm for Imaging Calorimetry
     *  1. group all the adjacent pixels
     *
     *  Author: Chao Peng (ANL), 06/02/2021
     *  References: https://arxiv.org/pdf/1603.02934.pdf
     *
     */
    #include "fmt/format.h"
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include <algorithm>
    
    Chao Peng's avatar
    Chao Peng committed
    
    #include "Gaudi/Property.h"
    #include "GaudiAlg/GaudiAlgorithm.h"
    #include "GaudiAlg/GaudiTool.h"
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include "GaudiAlg/Transformer.h"
    
    Chao Peng's avatar
    Chao Peng committed
    #include "GaudiKernel/PhysicalConstants.h"
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include "GaudiKernel/RndmGenerators.h"
    #include "GaudiKernel/ToolHandle.h"
    
    Chao Peng's avatar
    Chao Peng committed
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include "DD4hep/BitFieldCoder.h"
    
    Chao Peng's avatar
    Chao Peng committed
    #include "DDRec/CellIDPositionConverter.h"
    #include "DDRec/Surface.h"
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include "DDRec/SurfaceManager.h"
    
    Chao Peng's avatar
    Chao Peng committed
    
    // FCCSW
    #include "JugBase/DataHandle.h"
    #include "JugBase/IGeoSvc.h"
    
    // Event Model related classes
    #include "eicd/ImagingClusterCollection.h"
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    #include "eicd/ImagingPixelCollection.h"
    
    Chao Peng's avatar
    Chao Peng committed
    
    using namespace Gaudi::Units;
    
    namespace Jug::Reco {
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
      class ImagingTopoCluster : public GaudiAlgorithm {
      public:
    
    Chao Peng's avatar
    Chao Peng committed
        // maximum difference in layer numbers that can be considered as neighbours
    
        Gaudi::Property<int> m_neighbourLayersRange{this, "neighbourLayersRange", 1};
    
    Chao Peng's avatar
    Chao Peng committed
        // maximum distance of local (x, y) to be considered as neighbors at the same layer
    
        Gaudi::Property<std::vector<double>> u_localDistXY{this, "localDistXY", {1.0 * mm, 1.0 * mm}};
    
    Chao Peng's avatar
    Chao Peng committed
        // maximum distance of global (eta, phi) to be considered as neighbors at different layers
    
        Gaudi::Property<std::vector<double>> u_layerDistEtaPhi{this, "layerDistEtaPhi", {0.01, 0.01}};
    
    Chao Peng's avatar
    Chao Peng committed
        // maximum global distance to be considered as neighbors in different sectors
    
        Gaudi::Property<double> m_sectorDist{this, "sectorDist", 1.0 * cm};
    
        // minimum hit energy to participate clustering
        Gaudi::Property<double> m_minClusterHitEdep{this, "minClusterHitEdep", 0.};
    
    Chao Peng's avatar
    Chao Peng committed
        // minimum cluster center energy (to be considered as a seed for cluster)
        Gaudi::Property<double> m_minClusterCenterEdep{this, "minClusterCenterEdep", 0.};
        // minimum cluster energy (to save this cluster)
    
        Gaudi::Property<double> m_minClusterEdep{this, "minClusterEdep", 0.5 * MeV};
    
    Chao Peng's avatar
    Chao Peng committed
        // minimum number of hits (to save this cluster)
    
        Gaudi::Property<int> m_minClusterNhits{this, "minClusterNhits", 10};
    
    Chao Peng's avatar
    Chao Peng committed
        // input hits collection
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
        DataHandle<eic::ImagingPixelCollection> m_inputHitCollection{"inputHitCollection",
                                                                     Gaudi::DataHandle::Reader, this};
        // output clustered hits
        DataHandle<eic::ImagingPixelCollection> m_outputHitCollection{"outputHitCollection",
                                                                      Gaudi::DataHandle::Writer, this};
    
    
        // unitless counterparts of the input parameters
        double localDistXY[2], layerDistEtaPhi[2], sectorDist;
        double minClusterHitEdep, minClusterCenterEdep, minClusterEdep, minClusterNhits;
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
        ImagingTopoCluster(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc)
    
    Chao Peng's avatar
    Chao Peng committed
        {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          declareProperty("inputHitCollection", m_inputHitCollection, "");
          declareProperty("outputHitCollection", m_outputHitCollection, "");
    
    Chao Peng's avatar
    Chao Peng committed
        }
    
        StatusCode initialize() override
        {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          if (GaudiAlgorithm::initialize().isFailure()) {
            return StatusCode::FAILURE;
          }
    
    
          // unitless conversion
          // sanity checks
          if (u_localDistXY.size() != 2) {
            error() << "Expected 2 values (x_dist, y_dist) for localDistXY" << endmsg;
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
            return StatusCode::FAILURE;
          }
    
          if (u_layerDistEtaPhi.size() != 2) {
            error() << "Expected 2 values (eta_dist, phi_dist) for layerDistEtaPhi" << endmsg;
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
            return StatusCode::FAILURE;
          }
    
    
          // using juggler internal units (GeV, mm, ns, rad)
          localDistXY[0] = u_localDistXY.value()[0]/mm;
          localDistXY[1] = u_localDistXY.value()[1]/mm;
          layerDistEtaPhi[0] = u_layerDistEtaPhi.value()[0];
          layerDistEtaPhi[1] = u_layerDistEtaPhi.value()[1]/rad;
          sectorDist = m_sectorDist.value()/mm;
          minClusterHitEdep = m_minClusterHitEdep.value()/GeV;
          minClusterCenterEdep = m_minClusterCenterEdep.value()/GeV;
          minClusterEdep = m_minClusterEdep.value()/GeV;
    
          // summarize the clustering parameters
          info() << fmt::format("Local clustering (same sector and same layer): "
                                "Local [x, y] distance between hits <= [{:.4f} mm, {:.4f} mm].",
                                localDistXY[0], localDistXY[1]) << endmsg;
          info() << fmt::format("Neighbour layers clustering (same sector and layer id within +- {:d}: "
                                "Global [eta, phi] distance between hits <= [{:.4f}, {:.4f} rad].",
                                m_neighbourLayersRange.value(), layerDistEtaPhi[0], layerDistEtaPhi[1]) << endmsg;
          info() << fmt::format("Neighbour sectors clustering (different sector): "
                                "Global distance between hits <= {:.4f} mm.",
                                sectorDist) << endmsg;
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
    
          return StatusCode::SUCCESS;
    
    Chao Peng's avatar
    Chao Peng committed
        }
    
        StatusCode execute() override
        {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          // input collections
          const auto& hits = *m_inputHitCollection.get();
          // Create output collections
          auto& clustered_hits = *m_outputHitCollection.createAndPut();
    
          // group neighboring hits
          std::vector<bool>                           visits(hits.size(), false);
          std::vector<std::vector<eic::ImagingPixel>> groups;
          for (size_t i = 0; i < hits.size(); ++i) {
            // already in a group, or not energetic enough to form a cluster
    
            if (visits[i] || hits[i].edep() < minClusterCenterEdep) {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
              continue;
            }
            groups.emplace_back();
            // create a new group, and group all the neighboring hits
            dfs_group(groups.back(), i, hits, visits);
          }
          debug() << "we have " << groups.size() << " groups of hits" << endmsg;
    
          size_t clusterID = 0;
          for (const auto& group : groups) {
    
            if (static_cast<int>(group.size()) < m_minClusterNhits.value()) {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
              continue;
            }
            double edep = 0.;
            for (const auto& hit : group) {
              edep += hit.edep();
            }
    
            if (edep < minClusterEdep) {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
              continue;
    
    Chao Peng's avatar
    Chao Peng committed
            }
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
            for (auto hit : group) {
              hit.clusterID(clusterID);
              clustered_hits.push_back(hit);
    
    Chao Peng's avatar
    Chao Peng committed
            }
    
    Chao Peng's avatar
    Chao Peng committed
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          return StatusCode::SUCCESS;
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
      private:
        template <typename T>
        static inline T pow2(const T& x)
        {
          return x * x;
        }
    
    Chao Peng's avatar
    Chao Peng committed
    
        // helper function to group hits
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
        bool is_neighbor(const eic::ConstImagingPixel& h1, const eic::ConstImagingPixel& h2) const
    
    Chao Peng's avatar
    Chao Peng committed
        {
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          // different sectors, simple distance check
          if (h1.sectorID() != h2.sectorID()) {
            return std::sqrt(pow2(h1.x() - h2.x()) + pow2(h1.y() - h2.y()) + pow2(h1.z() - h2.z())) <=
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          }
    
          // layer check
          int ldiff = std::abs(h1.layerID() - h2.layerID());
          // same layer, check local positions
          if (!ldiff) {
    
            return (std::abs(h1.local_x() - h2.local_x()) <= localDistXY[0]) &&
                   (std::abs(h1.local_y() - h2.local_y()) <= localDistXY[1]);
          } else if (ldiff <= m_neighbourLayersRange) {
            return (std::abs(h1.eta() - h2.eta()) <= layerDistEtaPhi[0]) &&
                   (std::abs(h1.phi() - h2.phi()) <= layerDistEtaPhi[1]);
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          }
    
          // not in adjacent layers
          return false;
    
    Chao Peng's avatar
    Chao Peng committed
        }
    
        // grouping function with Depth-First Search
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
        void dfs_group(std::vector<eic::ImagingPixel>& group, int idx,
                       const eic::ImagingPixelCollection& hits, std::vector<bool>& visits) const
    
    Chao Peng's avatar
    Chao Peng committed
        {
    
          // not a qualified hit to participate in clustering, stop here
          if (hits[idx].edep() < minClusterHitEdep) {
            visits[idx] = true;
            return;
          }
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
          eic::ImagingPixel hit{hits[idx].clusterID(), hits[idx].layerID(), hits[idx].sectorID(),
                                hits[idx].hitID(),     hits[idx].edep(),    hits[idx].time(),
                                hits[idx].eta(),       hits[idx].local(),   hits[idx].position(),
                                hits[idx].polar()};
          group.push_back(hit);
          visits[idx] = true;
          for (size_t i = 0; i < hits.size(); ++i) {
            // visited, or not a neighbor
            if (visits[i] || !is_neighbor(hit, hits[i])) {
              continue;
    
    Chao Peng's avatar
    Chao Peng committed
            }
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
            dfs_group(group, i, hits, visits);
          }
    
    Chao Peng's avatar
    Chao Peng committed
        }
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
      }; // namespace Jug::Reco
    
    Chao Peng's avatar
    Chao Peng committed
    
    
    Sylvester Joosten's avatar
    Sylvester Joosten committed
      DECLARE_COMPONENT(ImagingTopoCluster)
    
    Chao Peng's avatar
    Chao Peng committed
    
    } // namespace Jug::Reco