ImagingTopoCluster.cpp 9.18 KB
Newer Older
Chao Peng's avatar
Chao Peng committed
1
2
3
4
5
6
7
8
9
/*
 *  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
10
#include <algorithm>
Chao Peng's avatar
Chao Peng committed
11
12
13
14

#include "Gaudi/Property.h"
#include "GaudiAlg/GaudiAlgorithm.h"
#include "GaudiAlg/GaudiTool.h"
Sylvester Joosten's avatar
Sylvester Joosten committed
15
#include "GaudiAlg/Transformer.h"
Chao Peng's avatar
Chao Peng committed
16
#include "GaudiKernel/PhysicalConstants.h"
Sylvester Joosten's avatar
Sylvester Joosten committed
17
18
#include "GaudiKernel/RndmGenerators.h"
#include "GaudiKernel/ToolHandle.h"
Chao Peng's avatar
Chao Peng committed
19

Sylvester Joosten's avatar
Sylvester Joosten committed
20
#include "DD4hep/BitFieldCoder.h"
Chao Peng's avatar
Chao Peng committed
21
22
#include "DDRec/CellIDPositionConverter.h"
#include "DDRec/Surface.h"
Sylvester Joosten's avatar
Sylvester Joosten committed
23
#include "DDRec/SurfaceManager.h"
Chao Peng's avatar
Chao Peng committed
24
25
26
27
28
29
30

// FCCSW
#include "JugBase/DataHandle.h"
#include "JugBase/IGeoSvc.h"

// Event Model related classes
#include "eicd/ImagingClusterCollection.h"
Sylvester Joosten's avatar
Sylvester Joosten committed
31
#include "eicd/ImagingPixelCollection.h"
Chao Peng's avatar
Chao Peng committed
32
33
34
35
36

using namespace Gaudi::Units;

namespace Jug::Reco {

Whitney Armstrong's avatar
Whitney Armstrong committed
37
38
39
40
41
42
43
44
45
46
  /** Topological Cell Clustering Algorithm.
   *
   * 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
   *
   * \ingroup reco
   */
Sylvester Joosten's avatar
Sylvester Joosten committed
47
48
  class ImagingTopoCluster : public GaudiAlgorithm {
  public:
Chao Peng's avatar
Chao Peng committed
49
    // maximum difference in layer numbers that can be considered as neighbours
50
    Gaudi::Property<int> m_neighbourLayersRange{this, "neighbourLayersRange", 1};
Chao Peng's avatar
Chao Peng committed
51
    // maximum distance of local (x, y) to be considered as neighbors at the same layer
52
    Gaudi::Property<std::vector<double>> u_localDistXY{this, "localDistXY", {1.0 * mm, 1.0 * mm}};
Chao Peng's avatar
Chao Peng committed
53
    // maximum distance of global (eta, phi) to be considered as neighbors at different layers
54
    Gaudi::Property<std::vector<double>> u_layerDistEtaPhi{this, "layerDistEtaPhi", {0.01, 0.01}};
Chao Peng's avatar
Chao Peng committed
55
    // maximum global distance to be considered as neighbors in different sectors
56
57
58
59
    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
60
61
62
    // 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)
63
    Gaudi::Property<double> m_minClusterEdep{this, "minClusterEdep", 0.5 * MeV};
Chao Peng's avatar
Chao Peng committed
64
    // minimum number of hits (to save this cluster)
65
    Gaudi::Property<int> m_minClusterNhits{this, "minClusterNhits", 10};
Chao Peng's avatar
Chao Peng committed
66
    // input hits collection
Sylvester Joosten's avatar
Sylvester Joosten committed
67
68
69
70
71
72
    DataHandle<eic::ImagingPixelCollection> m_inputHitCollection{"inputHitCollection",
                                                                 Gaudi::DataHandle::Reader, this};
    // output clustered hits
    DataHandle<eic::ImagingPixelCollection> m_outputHitCollection{"outputHitCollection",
                                                                  Gaudi::DataHandle::Writer, this};

73
74
75
76
    // unitless counterparts of the input parameters
    double localDistXY[2], layerDistEtaPhi[2], sectorDist;
    double minClusterHitEdep, minClusterCenterEdep, minClusterEdep, minClusterNhits;

Sylvester Joosten's avatar
Sylvester Joosten committed
77
    ImagingTopoCluster(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc)
Chao Peng's avatar
Chao Peng committed
78
    {
Sylvester Joosten's avatar
Sylvester Joosten committed
79
80
      declareProperty("inputHitCollection", m_inputHitCollection, "");
      declareProperty("outputHitCollection", m_outputHitCollection, "");
Chao Peng's avatar
Chao Peng committed
81
82
83
84
    }

    StatusCode initialize() override
    {
Sylvester Joosten's avatar
Sylvester Joosten committed
85
86
87
88
      if (GaudiAlgorithm::initialize().isFailure()) {
        return StatusCode::FAILURE;
      }

89
90
91
92
      // 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
93
94
        return StatusCode::FAILURE;
      }
95
96
      if (u_layerDistEtaPhi.size() != 2) {
        error() << "Expected 2 values (eta_dist, phi_dist) for layerDistEtaPhi" << endmsg;
Sylvester Joosten's avatar
Sylvester Joosten committed
97
98
        return StatusCode::FAILURE;
      }
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

      // 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
120
121

      return StatusCode::SUCCESS;
Chao Peng's avatar
Chao Peng committed
122
123
124
125
    }

    StatusCode execute() override
    {
Sylvester Joosten's avatar
Sylvester Joosten committed
126
127
128
129
130
131
132
133
134
      // 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) {
135
136
137
        /*
        debug() << fmt::format("hit {:d}: local position = ({}, {}, {}), global position = ({}, {}, {})",
                i + 1,
138
139
                hits[i].local().local_x, hits[i].local().local_y, hits[i].local_z(),
                hits[i].position().x, hits[i].position().y, hits[i].position().z)
140
141
            << endmsg;
        */
Sylvester Joosten's avatar
Sylvester Joosten committed
142
        // already in a group, or not energetic enough to form a cluster
143
        if (visits[i] || hits[i].edep() < minClusterCenterEdep) {
Sylvester Joosten's avatar
Sylvester Joosten committed
144
145
146
147
148
149
          continue;
        }
        groups.emplace_back();
        // create a new group, and group all the neighboring hits
        dfs_group(groups.back(), i, hits, visits);
      }
150
151
152
153
154
155
      debug() << "found " << groups.size() << " potential clusters (groups of hits)" << endmsg;
      /*
      for (size_t i = 0; i < groups.size(); ++i) {
        debug() << fmt::format("group {}: {} hits", i, groups[i].size()) << endmsg;
      }
      */
Sylvester Joosten's avatar
Sylvester Joosten committed
156

157
      // form clusters
Sylvester Joosten's avatar
Sylvester Joosten committed
158
159
      size_t clusterID = 0;
      for (const auto& group : groups) {
160
        if (static_cast<int>(group.size()) < m_minClusterNhits.value()) {
Sylvester Joosten's avatar
Sylvester Joosten committed
161
162
163
164
165
166
          continue;
        }
        double edep = 0.;
        for (const auto& hit : group) {
          edep += hit.edep();
        }
167
        if (edep < minClusterEdep) {
Sylvester Joosten's avatar
Sylvester Joosten committed
168
          continue;
Chao Peng's avatar
Chao Peng committed
169
        }
Sylvester Joosten's avatar
Sylvester Joosten committed
170
171
172
        for (auto hit : group) {
          hit.clusterID(clusterID);
          clustered_hits.push_back(hit);
Chao Peng's avatar
Chao Peng committed
173
        }
174
        clusterID += 1;
Sylvester Joosten's avatar
Sylvester Joosten committed
175
      }
Chao Peng's avatar
Chao Peng committed
176

Sylvester Joosten's avatar
Sylvester Joosten committed
177
      return StatusCode::SUCCESS;
Chao Peng's avatar
Chao Peng committed
178
179
    }

Sylvester Joosten's avatar
Sylvester Joosten committed
180
181
182
183
184
185
  private:
    template <typename T>
    static inline T pow2(const T& x)
    {
      return x * x;
    }
Chao Peng's avatar
Chao Peng committed
186
187

    // helper function to group hits
Sylvester Joosten's avatar
Sylvester Joosten committed
188
    bool is_neighbor(const eic::ConstImagingPixel& h1, const eic::ConstImagingPixel& h2) const
Chao Peng's avatar
Chao Peng committed
189
    {
Sylvester Joosten's avatar
Sylvester Joosten committed
190
191
      // different sectors, simple distance check
      if (h1.sectorID() != h2.sectorID()) {
192
        return std::sqrt(pow2(h1.position().x - h2.position().x) + pow2(h1.position().y - h2.position().y) + pow2(h1.position().z - h2.position().z)) <=
193
               sectorDist;
Sylvester Joosten's avatar
Sylvester Joosten committed
194
195
196
197
198
199
      }

      // layer check
      int ldiff = std::abs(h1.layerID() - h2.layerID());
      // same layer, check local positions
      if (!ldiff) {
200
201
        return (std::abs(h1.local().local_x - h2.local().local_x) <= localDistXY[0]) &&
               (std::abs(h1.local().local_y - h2.local().local_y) <= localDistXY[1]);
202
203
      } else if (ldiff <= m_neighbourLayersRange) {
        return (std::abs(h1.eta() - h2.eta()) <= layerDistEtaPhi[0]) &&
204
               (std::abs(h1.polar().phi - h2.polar().phi) <= layerDistEtaPhi[1]);
Sylvester Joosten's avatar
Sylvester Joosten committed
205
206
207
208
      }

      // not in adjacent layers
      return false;
Chao Peng's avatar
Chao Peng committed
209
210
211
    }

    // grouping function with Depth-First Search
Sylvester Joosten's avatar
Sylvester Joosten committed
212
213
    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
214
    {
215
216
217
218
219
220
      // 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
221
222
223
224
225
226
227
228
229
230
      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
231
        }
Sylvester Joosten's avatar
Sylvester Joosten committed
232
233
        dfs_group(group, i, hits, visits);
      }
Chao Peng's avatar
Chao Peng committed
234
    }
Sylvester Joosten's avatar
Sylvester Joosten committed
235
  }; // namespace Jug::Reco
Chao Peng's avatar
Chao Peng committed
236

Sylvester Joosten's avatar
Sylvester Joosten committed
237
  DECLARE_COMPONENT(ImagingTopoCluster)
Chao Peng's avatar
Chao Peng committed
238
239
240

} // namespace Jug::Reco