src/todo/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 "JugBase/DataHandle.h"
0027 #include "JugBase/IGeoSvc.h"
0028 #include "JugBase/Utilities/Utils.hpp"
0029 #include "JugReco/ClusterTypes.h"
0030
0031 // Event Model related classes
0032 #include "edm4hep/MCParticleCollection.h"
0033 #include "edm4hep/SimCalorimeterHitCollection.h"
0034 #include "edm4eic/CalorimeterHitCollection.h"
0035 #include "edm4eic/ClusterCollection.h"
0036 #include "edm4eic/MCRecoClusterParticleAssociationCollection.h"
0037 #include "edm4eic/ProtoClusterCollection.h"
0038 #include "edm4hep/utils/vector_utils.h"
0039
0040 using namespace Gaudi::Units;
0041 using namespace Eigen;
0042
0043 namespace Jug::Reco {
0044
0045 /** Imaging cluster reconstruction.
0046  *
0047  *  Reconstruct the cluster/layer info for imaging calorimeter
0048  *  Logarithmic weighting is used to describe energy deposit in transverse direction
0049  *
0050  *  \ingroup reco
0051  */
0052 class ImagingClusterReco : public GaudiAlgorithm {
0053 private:
0054   Gaudi::Property<int> m_trackStopLayer{this, "trackStopLayer", 9};
0055
0056   DataHandle<edm4eic::ProtoClusterCollection> m_inputProtoClusters{"inputProtoClusters", Gaudi::DataHandle::Reader, this};
0057   DataHandle<edm4eic::ClusterCollection> m_outputLayers{"outputLayers", Gaudi::DataHandle::Writer, this};
0058   DataHandle<edm4eic::ClusterCollection> m_outputClusters{"outputClusters", Gaudi::DataHandle::Reader, this};
0059
0060   // Collection for MC hits when running on MC
0061   Gaudi::Property<std::string> m_mcHits{this, "mcHits", ""};
0062   // Optional handle to MC hits
0063   std::unique_ptr<DataHandle<edm4hep::SimCalorimeterHitCollection>> m_mcHits_ptr;
0064
0065   // Collection for associations when running on MC
0066   Gaudi::Property<std::string> m_outputAssociations{this, "outputAssociations", ""};
0067   // Optional handle to MC hits
0068   std::unique_ptr<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>> m_outputAssociations_ptr;
0069
0070 public:
0071   ImagingClusterReco(const std::string& name, ISvcLocator* svcLoc) : GaudiAlgorithm(name, svcLoc) {
0072     declareProperty("inputProtoClusters", m_inputProtoClusters, "");
0073     declareProperty("outputLayers", m_outputLayers, "");
0074     declareProperty("outputClusters", m_outputClusters, "");
0075   }
0076
0077   StatusCode initialize() override {
0078     if (GaudiAlgorithm::initialize().isFailure()) {
0079       return StatusCode::FAILURE;
0080     }
0081
0082     // Initialize the optional MC input hit collection if requested
0083     if (m_mcHits != "") {
0084       m_mcHits_ptr =
0085         std::make_unique<DataHandle<edm4hep::SimCalorimeterHitCollection>>(m_mcHits, Gaudi::DataHandle::Reader,
0086         this);
0087     }
0088
0089     // Initialize the optional association collection if requested
0090     if (m_outputAssociations != "") {
0091       m_outputAssociations_ptr =
0092         std::make_unique<DataHandle<edm4eic::MCRecoClusterParticleAssociationCollection>>(m_outputAssociations, Gaudi::DataHandle::Writer,
0093         this);
0094     }
0095
0096     return StatusCode::SUCCESS;
0097   }
0098
0099   StatusCode execute() override {
0100     // input collections
0101     const auto& proto = *m_inputProtoClusters.get();
0102     // output collections
0103     auto& layers   = *m_outputLayers.createAndPut();
0104     auto& clusters = *m_outputClusters.createAndPut();
0105
0106     // Optional input MC data
0107     const edm4hep::SimCalorimeterHitCollection* mcHits = nullptr;
0108     if (m_mcHits_ptr) {
0109       mcHits = m_mcHits_ptr->get();
0110     }
0111
0112     // Optional output associations
0113     edm4eic::MCRecoClusterParticleAssociationCollection* associations = nullptr;
0114     if (m_outputAssociations_ptr) {
0115       associations = m_outputAssociations_ptr->createAndPut();
0116     }
0117
0118     for (const auto& pcl : proto) {
0119       if (!pcl.getHits().empty() && !pcl.getHits(0).isAvailable()) {
0120         warning() << "Protocluster hit relation is invalid, skipping protocluster" << endmsg;
0121         continue;
0122       }
0123       // get cluster and associated layers
0124       auto cl        = reconstruct_cluster(pcl);
0125       auto cl_layers = reconstruct_cluster_layers(pcl);
0126
0127       // Get cluster direction from the layer profile
0128       auto [theta, phi] = fit_track(cl_layers);
0129       cl.setIntrinsicTheta(theta);
0130       cl.setIntrinsicPhi(phi);
0131       // no error on the intrinsic direction TODO
0132
0133       // store layer and clusters on the datastore
0134       for (auto& layer : cl_layers) {
0135         layers.push_back(layer);
0136         cl.addToClusters(layer);
0137       }
0138       clusters.push_back(cl);
0139
0140
0141       // If mcHits are available, associate cluster with MCParticle
0142       if (m_mcHits_ptr.get() != nullptr && m_outputAssociations_ptr.get() != nullptr) {
0143
0144         // 1. find pclhit with largest energy deposition
0145         auto pclhits = pcl.getHits();
0146         auto pclhit = std::max_element(
0147           pclhits.begin(),
0148           pclhits.end(),
0149           [](const auto& pclhit1, const auto& pclhit2) {
0150             return pclhit1.getEnergy() < pclhit2.getEnergy();
0151           }
0152         );
0153
0154         // 2. find mchit with same CellID
0155         auto mchit = mcHits->begin();
0156         for ( ; mchit != mcHits->end(); ++mchit) {
0157           // break loop when CellID match found
0158           if (mchit->getCellID() == pclhit->getCellID()) {
0159             break;
0160           }
0161         }
0162         if (!(mchit != mcHits->end())) {
0163           // break if no matching hit found for this CellID
0164           warning() << "Proto-cluster has highest energy in CellID " << pclhit->getCellID()
0165                     << ", but no mc hit with that CellID was found." << endmsg;
0166           break;
0167         }
0168
0169         // 3. find mchit's MCParticle
0170         const auto& mcp = mchit->getContributions(0).getParticle();
0171
0172         // set association
0173         edm4eic::MutableMCRecoClusterParticleAssociation clusterassoc;
0174         clusterassoc.setRecID(cl.getObjectID().index);
0175         clusterassoc.setSimID(mcp.getObjectID().index);
0176         clusterassoc.setWeight(1.0);
0177         clusterassoc.setRec(cl);
0178         //clusterassoc.setSim(mcp);
0179         associations->push_back(clusterassoc);
0180       }
0181
0182     }
0183
0184     // debug output
0185     if (msgLevel(MSG::DEBUG)) {
0186       for (const auto& cl : clusters) {
0187         debug() << fmt::format("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.id(),
0188                                cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0189                                cl.getIntrinsicPhi() / M_PI * 180.)
0190                 << endmsg;
0191       }
0192     }
0193
0194     return StatusCode::SUCCESS;
0195   }
0196
0197 private:
0198   template <typename T> static inline T pow2(const T& x) { return x * x; }
0199
0200   static std::vector<edm4eic::Cluster> reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) {
0201     const auto& hits    = pcl.getHits();
0202     const auto& weights = pcl.getWeights();
0203     // using map to have hits sorted by layer
0204     std::map<int, std::vector<std::pair<edm4eic::CalorimeterHit, float>>> layer_map;
0205     for (unsigned i = 0; i < hits.size(); ++i) {
0206       const auto hit = hits[i];
0207       auto lid       = hit.getLayer();
0208       if (layer_map.count(lid) == 0) {
0209         layer_map[lid] = {};
0210       }
0211       layer_map[lid].push_back({hit, weights[i]});
0212     }
0213
0214     // create layers
0215     std::vector<edm4eic::Cluster> cl_layers;
0216     for (const auto& [lid, layer_hits] : layer_map) {
0217       auto layer = reconstruct_layer(layer_hits);
0218       cl_layers.push_back(layer);
0219     }
0220     return cl_layers;
0221   }
0222
0223   static edm4eic::Cluster reconstruct_layer(const std::vector<std::pair<edm4eic::CalorimeterHit, float>>& hits) {
0224     edm4eic::MutableCluster layer;
0225     layer.setType(ClusterType::kClusterSlice);
0226     // Calculate averages
0227     double energy{0};
0228     double energyError{0};
0229     double time{0};
0230     double timeError{0};
0231     double sumOfWeights{0};
0232     auto pos            = layer.getPosition();
0233     for (const auto& [hit, weight] : hits) {
0234       energy += hit.getEnergy() * weight;
0235       energyError += std::pow(hit.getEnergyError() * weight, 2);
0236       time += hit.getTime() * weight;
0237       timeError += std::pow(hit.getTimeError() * weight, 2);
0238       pos = pos + hit.getPosition() * weight;
0239       sumOfWeights += weight;
0240       layer.addToHits(hit);
0241     }
0242     layer.setEnergy(energy);
0243     layer.setEnergyError(std::sqrt(energyError));
0244     layer.setTime(time / sumOfWeights);
0245     layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0246     layer.setNhits(hits.size());
0247     layer.setPosition(pos / sumOfWeights);
0248     // positionError not set
0249     // Intrinsic direction meaningless in a cluster layer --> not set
0250
0251     // Calculate radius as the standard deviation of the hits versus the cluster center
0252     double radius = 0.;
0253     for (const auto& [hit, weight] : hits) {
0254       radius += std::pow(edm4hep::utils::magnitude(hit.getPosition() - layer.getPosition()), 2);
0255     }
0256     layer.addToShapeParameters(std::sqrt(radius / layer.getNhits()));
0257     // TODO Skewedness
0258
0259     return layer;
0260   }
0261
0262   edm4eic::MutableCluster reconstruct_cluster(const edm4eic::ProtoCluster& pcl) {
0263     edm4eic::MutableCluster cluster;
0264
0265     const auto& hits    = pcl.getHits();
0266     const auto& weights = pcl.getWeights();
0267
0268     cluster.setType(ClusterType::kCluster3D);
0269     double energy      = 0.;
0270     double energyError = 0.;
0271     double time        = 0.;
0272     double timeError   = 0.;
0273     double meta        = 0.;
0274     double mphi        = 0.;
0275     double r           = 9999 * cm;
0276     for (unsigned i = 0; i < hits.size(); ++i) {
0277       const auto& hit    = hits[i];
0278       const auto& weight = weights[i];
0279       energy += hit.getEnergy() * weight;
0280       energyError += std::pow(hit.getEnergyError() * weight, 2);
0281       // energy weighting for the other variables
0282       const double energyWeight = hit.getEnergy() * weight;
0283       time += hit.getTime() * energyWeight;
0284       timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0285       meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0286       mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0287       r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0288       cluster.addToHits(hit);
0289     }
0290     cluster.setEnergy(energy);
0291     cluster.setEnergyError(std::sqrt(energyError));
0292     cluster.setTime(time / energy);
0293     cluster.setTimeError(std::sqrt(timeError) / energy);
0294     cluster.setNhits(hits.size());
0295     cluster.setPosition(edm4hep::utils::sphericalToVector(r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0296
0297     // shower radius estimate (eta-phi plane)
0298     double radius = 0.;
0299     for (const auto& hit : hits) {
0300       radius += pow2(edm4hep::utils::eta(hit.getPosition()) - edm4hep::utils::eta(cluster.getPosition())) +
0301                 pow2(edm4hep::utils::angleAzimuthal(hit.getPosition()) - edm4hep::utils::angleAzimuthal(cluster.getPosition()));
0302     }
0303     cluster.addToShapeParameters(std::sqrt(radius / cluster.getNhits()));
0304     // Skewedness not calculated TODO
0305
0306     // Optionally store the MC truth associated with the first hit in this cluster
0307     // FIXME no connection between cluster and truth in edm4hep
0308     // if (mcHits) {
0309     //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0310     //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0311     //}
0312
0313     return cluster;
0314   }
0315
0316   std::pair<double /* polar */, double /* azimuthal */> fit_track(const std::vector<edm4eic::Cluster>& layers) const {
0317     int nrows = 0;
0318     decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0319     for (const auto& layer : layers) {
0320       if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0321         mean_pos = mean_pos + layer.getPosition();
0322         nrows += 1;
0323       }
0324     }
0325     // cannot fit
0326     if (nrows < 2) {
0327       return {};
0328     }
0329
0330     mean_pos = mean_pos / nrows;
0331     // fill position data
0332     MatrixXd pos(nrows, 3);
0333     int ir = 0;
0334     for (const auto& layer : layers) {
0335       if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0336         auto delta = layer.getPosition() - mean_pos;
0337         pos(ir, 0) = delta.x;
0338         pos(ir, 1) = delta.y;
0339         pos(ir, 2) = delta.z;
0340         ir += 1;
0341       }
0342     }
0343
0344     JacobiSVD<MatrixXd> svd(pos, ComputeThinU | ComputeThinV);
0345     const auto dir = svd.matrixV().col(0);
0346     // theta and phi
0347     return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0348   }
0349 };
0350
0351 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0352 DECLARE_COMPONENT(ImagingClusterReco)
0353
0354 } // namespace Jug::Reco