src/components/ImagingClusterReco.cpp

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Chao Peng, Wouter Deconinck
0003
0004 /*
0005  *  Reconstruct the cluster/layer info for imaging calorimeter
0006  *  Logarithmic weighting is used to describe energy deposit in transverse direction
0007  *
0008  *  Author: Chao Peng (ANL), 06/02/2021
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 // Event Model related classes
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 /** Imaging cluster reconstruction.
0044  *
0045  *  Reconstruct the cluster/layer info for imaging calorimeter
0046  *  Logarithmic weighting is used to describe energy deposit in transverse direction
0047  *
0048  *  \ingroup reco
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   // Collection for MC hits when running on MC
0059   Gaudi::Property<std::string> m_mcHits{this, "mcHits", ""};
0060   // Optional handle to MC hits
0061   std::unique_ptr<DataHandle<edm4hep::SimCalorimeterHitCollection>> m_mcHits_ptr;
0062
0063   // Collection for associations when running on MC
0064   Gaudi::Property<std::string> m_outputAssociations{this, "outputAssociations", ""};
0065   // Optional handle to MC hits
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     // Initialize the optional MC input hit collection if requested
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     // Initialize the optional association collection if requested
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     // input collections
0099     const auto& proto = *m_inputProtoClusters.get();
0100     // output collections
0101     auto& layers   = *m_outputLayers.createAndPut();
0102     auto& clusters = *m_outputClusters.createAndPut();
0103
0104     // Optional input MC data
0105     const edm4hep::SimCalorimeterHitCollection* mcHits = nullptr;
0106     if (m_mcHits_ptr) {
0107       mcHits = m_mcHits_ptr->get();
0108     }
0109
0110     // Optional output associations
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       // get cluster and associated layers
0122       auto cl        = reconstruct_cluster(pcl);
0123       auto cl_layers = reconstruct_cluster_layers(pcl);
0124
0125       // Get cluster direction from the layer profile
0126       auto [theta, phi] = fit_track(cl_layers);
0127       cl.setIntrinsicTheta(theta);
0128       cl.setIntrinsicPhi(phi);
0129       // no error on the intrinsic direction TODO
0130
0131       // store layer and clusters on the datastore
0132       for (auto& layer : cl_layers) {
0133         layers.push_back(layer);
0134         cl.addToClusters(layer);
0135       }
0136       clusters.push_back(cl);
0137
0138
0139       // If mcHits are available, associate cluster with MCParticle
0140       if (m_mcHits_ptr.get() != nullptr && m_outputAssociations_ptr.get() != nullptr) {
0141
0142         // 1. find pclhit with largest energy deposition
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         // 2. find mchit with same CellID
0153         auto mchit = mcHits->begin();
0154         for ( ; mchit != mcHits->end(); ++mchit) {
0155           // break loop when CellID match found
0156           if (mchit->getCellID() == pclhit->getCellID()) {
0157             break;
0158           }
0159         }
0160         if (!(mchit != mcHits->end())) {
0161           // break if no matching hit found for this CellID
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         // 3. find mchit's MCParticle
0168         const auto& mcp = mchit->getContributions(0).getParticle();
0169
0170         // set association
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         //clusterassoc.setSim(mcp);
0177         associations->push_back(clusterassoc);
0178       }
0179
0180     }
0181
0182     // debug output
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     // using map to have hits sorted by layer
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     // create layers
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     // Calculate averages
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     // positionError not set
0247     // Intrinsic direction meaningless in a cluster layer --> not set
0248
0249     // Calculate radius as the standard deviation of the hits versus the cluster center
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     // TODO Skewedness
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       // energy weighting for the other variables
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     // shower radius estimate (eta-phi plane)
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     // Skewedness not calculated TODO
0303
0304     // Optionally store the MC truth associated with the first hit in this cluster
0305     // FIXME no connection between cluster and truth in edm4hep
0306     // if (mcHits) {
0307     //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0308     //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0309     //}
0310
0311     return cluster;
0312   }
0313
0314   std::pair<double /* polar */, double /* azimuthal */> 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     // cannot fit
0324     if (nrows < 2) {
0325       return {};
0326     }
0327
0328     mean_pos = mean_pos / nrows;
0329     // fill position data
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     // theta and phi
0345     return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0346   }
0347 };
0348
0349 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0350 DECLARE_COMPONENT(ImagingClusterReco)
0351
0352 } // namespace Jug::Reco