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