File indexing completed on 2024-09-28 07:03:47
0001
0002
0003
0004
0005
0006
0007
0008
0009
0010 #include "fmt/format.h"
0011 #include <Eigen/Dense>
0012 #include <algorithm>
0013
0014 #include "Gaudi/Property.h"
0015 #include "GaudiAlg/GaudiAlgorithm.h"
0016 #include "GaudiAlg/GaudiTool.h"
0017 #include "GaudiAlg/Transformer.h"
0018 #include "GaudiKernel/PhysicalConstants.h"
0019 #include "GaudiKernel/RndmGenerators.h"
0020 #include "GaudiKernel/ToolHandle.h"
0021
0022 #include "DDRec/CellIDPositionConverter.h"
0023 #include "DDRec/Surface.h"
0024 #include "DDRec/SurfaceManager.h"
0025
0026 #include <k4FWCore/DataHandle.h>
0027 #include <k4Interface/IGeoSvc.h>
0028 #include "JugReco/ClusterTypes.h"
0029
0030
0031 #include "edm4hep/MCParticleCollection.h"
0032 #include "edm4hep/SimCalorimeterHitCollection.h"
0033 #include "edm4eic/CalorimeterHitCollection.h"
0034 #include "edm4eic/ClusterCollection.h"
0035 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0036 #include "edm4eic/ProtoClusterCollection.h"
0037 #include "edm4hep/utils/vector_utils.h"
0038
0039 using namespace Gaudi::Units;
0040
0041 namespace Jug::Reco {
0042
0043
0044
0045
0046
0047
0048
0049
0050 class ImagingClusterReco : public GaudiAlgorithm {
0051 private:
0052 Gaudi::Property<int> m_trackStopLayer{this, "trackStopLayer", 9};
0053
0054 DataHandle<edm4eic::ProtoClusterCollection> m_inputProtoClusters{"inputProtoClusters", Gaudi::DataHandle::Reader, this};
0055 DataHandle<edm4eic::ClusterCollection> m_outputLayers{"outputLayers", Gaudi::DataHandle::Writer, this};
0056 DataHandle<edm4eic::ClusterCollection> m_outputClusters{"outputClusters", Gaudi::DataHandle::Reader, this};
0057
0058
0059 Gaudi::Property<std::string> m_mcHits{this, "mcHits", ""};
0060
0061 std::unique_ptr<DataHandle<edm4hep::SimCalorimeterHitCollection>> m_mcHits_ptr;
0062
0063
0064 Gaudi::Property<std::string> m_outputAssociations{this, "outputAssociations", ""};
0065
0066 std::unique_ptr<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>> m_outputAssociations_ptr;
0067
0068 public:
0069 ImagingClusterReco(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0070 declareProperty("inputProtoClusters", m_inputProtoClusters, "");
0071 declareProperty("outputLayers", m_outputLayers, "");
0072 declareProperty("outputClusters", m_outputClusters, "");
0073 }
0074
0075 StatusCode initialize() override {
0076 if (GaudiAlgorithm::initialize().isFailure()) {
0077 return StatusCode::FAILURE;
0078 }
0079
0080
0081 if (m_mcHits != "") {
0082 m_mcHits_ptr =
0083 std::make_unique<DataHandle<edm4hep::SimCalorimeterHitCollection>>(m_mcHits, Gaudi::DataHandle::Reader,
0084 this);
0085 }
0086
0087
0088 if (m_outputAssociations != "") {
0089 m_outputAssociations_ptr =
0090 std::make_unique<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>>(m_outputAssociations, Gaudi::DataHandle::Writer,
0091 this);
0092 }
0093
0094 return StatusCode::SUCCESS;
0095 }
0096
0097 StatusCode execute() override {
0098
0099 const auto& proto = *m_inputProtoClusters.get();
0100
0101 auto& layers = *m_outputLayers.createAndPut();
0102 auto& clusters = *m_outputClusters.createAndPut();
0103
0104
0105 const edm4hep::SimCalorimeterHitCollection* mcHits = nullptr;
0106 if (m_mcHits_ptr) {
0107 mcHits = m_mcHits_ptr->get();
0108 }
0109
0110
0111 edm4eic::MCRecoClusterParticleAssociationCollection* associations = nullptr;
0112 if (m_outputAssociations_ptr) {
0113 associations = m_outputAssociations_ptr->createAndPut();
0114 }
0115
0116 for (const auto& pcl : proto) {
0117 if (!pcl.getHits().empty() && !pcl.getHits(0).isAvailable()) {
0118 warning() << "Protocluster hit relation is invalid, skipping protocluster" << endmsg;
0119 continue;
0120 }
0121
0122 auto cl = reconstruct_cluster(pcl);
0123 auto cl_layers = reconstruct_cluster_layers(pcl);
0124
0125
0126 auto [theta, phi] = fit_track(cl_layers);
0127 cl.setIntrinsicTheta(theta);
0128 cl.setIntrinsicPhi(phi);
0129
0130
0131
0132 for (auto& layer : cl_layers) {
0133 layers.push_back(layer);
0134 cl.addToClusters(layer);
0135 }
0136 clusters.push_back(cl);
0137
0138
0139
0140 if (m_mcHits_ptr.get() != nullptr && m_outputAssociations_ptr.get() != nullptr) {
0141
0142
0143 auto pclhits = pcl.getHits();
0144 auto pclhit = std::max_element(
0145 pclhits.begin(),
0146 pclhits.end(),
0147 [](const auto& pclhit1, const auto& pclhit2) {
0148 return pclhit1.getEnergy() < pclhit2.getEnergy();
0149 }
0150 );
0151
0152
0153 auto mchit = mcHits->begin();
0154 for ( ; mchit != mcHits->end(); ++mchit) {
0155
0156 if (mchit->getCellID() == pclhit->getCellID()) {
0157 break;
0158 }
0159 }
0160 if (!(mchit != mcHits->end())) {
0161
0162 warning() << "Proto-cluster has highest energy in CellID " << pclhit->getCellID()
0163 << ", but no mc hit with that CellID was found." << endmsg;
0164 break;
0165 }
0166
0167
0168 const auto& mcp = mchit->getContributions(0).getParticle();
0169
0170
0171 edm4eic::MutableMCRecoClusterParticleAssociation clusterassoc;
0172 clusterassoc.setRecID(cl.getObjectID().index);
0173 clusterassoc.setSimID(mcp.getObjectID().index);
0174 clusterassoc.setWeight(1.0);
0175 clusterassoc.setRec(cl);
0176
0177 associations->push_back(clusterassoc);
0178 }
0179
0180 }
0181
0182
0183 if (msgLevel(MSG::DEBUG)) {
0184 for (const auto& cl : clusters) {
0185 debug() << fmt::format("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.getObjectID().index,
0186 cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0187 cl.getIntrinsicPhi() / M_PI * 180.)
0188 << endmsg;
0189 }
0190 }
0191
0192 return StatusCode::SUCCESS;
0193 }
0194
0195 private:
0196 template <typename T> static inline T pow2(const T& x) { return x * x; }
0197
0198 static std::vector<edm4eic::MutableCluster> reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) {
0199 const auto& hits = pcl.getHits();
0200 const auto& weights = pcl.getWeights();
0201
0202 std::map<int, std::vector<std::pair<edm4eic::CalorimeterHit, float>>> layer_map;
0203 for (unsigned i = 0; i < hits.size(); ++i) {
0204 const auto hit = hits[i];
0205 auto lid = hit.getLayer();
0206 if (layer_map.count(lid) == 0) {
0207 layer_map[lid] = {};
0208 }
0209 layer_map[lid].push_back({hit, weights[i]});
0210 }
0211
0212
0213 std::vector<edm4eic::MutableCluster> cl_layers;
0214 for (const auto& [lid, layer_hits] : layer_map) {
0215 auto layer = reconstruct_layer(layer_hits);
0216 cl_layers.push_back(layer);
0217 }
0218 return cl_layers;
0219 }
0220
0221 static edm4eic::MutableCluster reconstruct_layer(const std::vector<std::pair<edm4eic::CalorimeterHit, float>>& hits) {
0222 edm4eic::MutableCluster layer;
0223 layer.setType(ClusterType::kClusterSlice);
0224
0225 double energy{0};
0226 double energyError{0};
0227 double time{0};
0228 double timeError{0};
0229 double sumOfWeights{0};
0230 auto pos = layer.getPosition();
0231 for (const auto& [hit, weight] : hits) {
0232 energy += hit.getEnergy() * weight;
0233 energyError += std::pow(hit.getEnergyError() * weight, 2);
0234 time += hit.getTime() * weight;
0235 timeError += std::pow(hit.getTimeError() * weight, 2);
0236 pos = pos + hit.getPosition() * weight;
0237 sumOfWeights += weight;
0238 layer.addToHits(hit);
0239 }
0240 layer.setEnergy(energy);
0241 layer.setEnergyError(std::sqrt(energyError));
0242 layer.setTime(time / sumOfWeights);
0243 layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0244 layer.setNhits(hits.size());
0245 layer.setPosition(pos / sumOfWeights);
0246
0247
0248
0249
0250 double radius = 0.;
0251 for (const auto& [hit, weight] : hits) {
0252 radius += std::pow(edm4hep::utils::magnitude(hit.getPosition() - layer.getPosition()), 2);
0253 }
0254 layer.addToShapeParameters(std::sqrt(radius / layer.getNhits()));
0255
0256
0257 return layer;
0258 }
0259
0260 edm4eic::MutableCluster reconstruct_cluster(const edm4eic::ProtoCluster& pcl) {
0261 edm4eic::MutableCluster cluster;
0262
0263 const auto& hits = pcl.getHits();
0264 const auto& weights = pcl.getWeights();
0265
0266 cluster.setType(ClusterType::kCluster3D);
0267 double energy = 0.;
0268 double energyError = 0.;
0269 double time = 0.;
0270 double timeError = 0.;
0271 double meta = 0.;
0272 double mphi = 0.;
0273 double r = 9999 * cm;
0274 for (unsigned i = 0; i < hits.size(); ++i) {
0275 const auto& hit = hits[i];
0276 const auto& weight = weights[i];
0277 energy += hit.getEnergy() * weight;
0278 energyError += std::pow(hit.getEnergyError() * weight, 2);
0279
0280 const double energyWeight = hit.getEnergy() * weight;
0281 time += hit.getTime() * energyWeight;
0282 timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0283 meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0284 mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0285 r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0286 cluster.addToHits(hit);
0287 }
0288 cluster.setEnergy(energy);
0289 cluster.setEnergyError(std::sqrt(energyError));
0290 cluster.setTime(time / energy);
0291 cluster.setTimeError(std::sqrt(timeError) / energy);
0292 cluster.setNhits(hits.size());
0293 cluster.setPosition(edm4hep::utils::sphericalToVector(r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0294
0295
0296 double radius = 0.;
0297 for (const auto& hit : hits) {
0298 radius += pow2(edm4hep::utils::eta(hit.getPosition()) - edm4hep::utils::eta(cluster.getPosition())) +
0299 pow2(edm4hep::utils::angleAzimuthal(hit.getPosition()) - edm4hep::utils::angleAzimuthal(cluster.getPosition()));
0300 }
0301 cluster.addToShapeParameters(std::sqrt(radius / cluster.getNhits()));
0302
0303
0304
0305
0306
0307
0308
0309
0310
0311 return cluster;
0312 }
0313
0314 std::pair<double , double > fit_track(const std::vector<edm4eic::MutableCluster>& layers) const {
0315 int nrows = 0;
0316 decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0317 for (const auto& layer : layers) {
0318 if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0319 mean_pos = mean_pos + layer.getPosition();
0320 nrows += 1;
0321 }
0322 }
0323
0324 if (nrows < 2) {
0325 return {};
0326 }
0327
0328 mean_pos = mean_pos / nrows;
0329
0330 Eigen::MatrixXd pos(nrows, 3);
0331 int ir = 0;
0332 for (const auto& layer : layers) {
0333 if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0334 auto delta = layer.getPosition() - mean_pos;
0335 pos(ir, 0) = delta.x;
0336 pos(ir, 1) = delta.y;
0337 pos(ir, 2) = delta.z;
0338 ir += 1;
0339 }
0340 }
0341
0342 Eigen::JacobiSVD<Eigen::MatrixXd> svd(pos, Eigen::ComputeThinU | Eigen::ComputeThinV);
0343 const auto dir = svd.matrixV().col(0);
0344
0345 return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0346 }
0347 };
0348
0349
0350 DECLARE_COMPONENT(ImagingClusterReco)
0351
0352 }