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 k4FWCore::DataHandle<edm4eic::ProtoClusterCollection> m_inputProtoClusters{"inputProtoClusters", Gaudi::DataHandle::Reader, this};
0053   mutable k4FWCore::DataHandle<edm4eic::ClusterCollection> m_outputLayers{"outputLayers", Gaudi::DataHandle::Writer, this};
0054   mutable k4FWCore::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<k4FWCore::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<k4FWCore::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<k4FWCore::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<k4FWCore::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.setWeight(1.0);
0171         clusterassoc.setRec(cl);
0172         //clusterassoc.setSim(mcp);
0173         associations->push_back(clusterassoc);
0174       }
0175
0176     }
0177
0178     // debug output
0179     if (msgLevel(MSG::DEBUG)) {
0180       for (const auto& cl : clusters) {
0181         debug() << fmt::format("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.getObjectID().index,
0182                                cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0183                                cl.getIntrinsicPhi() / M_PI * 180.)
0184                 << endmsg;
0185       }
0186     }
0187
0188     return StatusCode::SUCCESS;
0189   }
0190
0191 private:
0192   template <typename T> static inline T pow2(const T& x) { return x * x; }
0193
0194   std::vector<edm4eic::MutableCluster> reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) const {
0195     const auto& hits    = pcl.getHits();
0196     const auto& weights = pcl.getWeights();
0197     // using map to have hits sorted by layer
0198     std::map<int, std::vector<std::pair<edm4eic::CalorimeterHit, float>>> layer_map;
0199     for (unsigned i = 0; i < hits.size(); ++i) {
0200       const auto hit = hits[i];
0201       auto lid       = hit.getLayer();
0202       if (layer_map.count(lid) == 0) {
0203         layer_map[lid] = {};
0204       }
0205       layer_map[lid].push_back({hit, weights[i]});
0206     }
0207
0208     // create layers
0209     std::vector<edm4eic::MutableCluster> cl_layers;
0210     for (const auto& [lid, layer_hits] : layer_map) {
0211       auto layer = reconstruct_layer(layer_hits);
0212       cl_layers.push_back(layer);
0213     }
0214     return cl_layers;
0215   }
0216
0217  edm4eic::MutableCluster reconstruct_layer(const std::vector<std::pair<edm4eic::CalorimeterHit, float>>& hits) const {
0218     edm4eic::MutableCluster layer;
0219     layer.setType(ClusterType::kClusterSlice);
0220     // Calculate averages
0221     double energy{0};
0222     double energyError{0};
0223     double time{0};
0224     double timeError{0};
0225     double sumOfWeights{0};
0226     auto pos            = layer.getPosition();
0227     for (const auto& [hit, weight] : hits) {
0228       energy += hit.getEnergy() * weight;
0229       energyError += std::pow(hit.getEnergyError() * weight, 2);
0230       time += hit.getTime() * weight;
0231       timeError += std::pow(hit.getTimeError() * weight, 2);
0232       pos = pos + hit.getPosition() * weight;
0233       sumOfWeights += weight;
0234       layer.addToHits(hit);
0235     }
0236     layer.setEnergy(energy);
0237     layer.setEnergyError(std::sqrt(energyError));
0238     layer.setTime(time / sumOfWeights);
0239     layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0240     layer.setNhits(hits.size());
0241     layer.setPosition(pos / sumOfWeights);
0242     // positionError not set
0243     // Intrinsic direction meaningless in a cluster layer --> not set
0244
0245     // Calculate radius as the standard deviation of the hits versus the cluster center
0246     double radius = 0.;
0247     for (const auto& [hit, weight] : hits) {
0248       radius += std::pow(edm4hep::utils::magnitude(hit.getPosition() - layer.getPosition()), 2);
0249     }
0250     layer.addToShapeParameters(std::sqrt(radius / layer.getNhits()));
0251     // TODO Skewedness
0252
0253     return layer;
0254   }
0255
0256   edm4eic::MutableCluster reconstruct_cluster(const edm4eic::ProtoCluster& pcl) const {
0257     edm4eic::MutableCluster cluster;
0258
0259     const auto& hits    = pcl.getHits();
0260     const auto& weights = pcl.getWeights();
0261
0262     cluster.setType(ClusterType::kCluster3D);
0263     double energy      = 0.;
0264     double energyError = 0.;
0265     double time        = 0.;
0266     double timeError   = 0.;
0267     double meta        = 0.;
0268     double mphi        = 0.;
0269     double r           = 9999 * cm;
0270     for (unsigned i = 0; i < hits.size(); ++i) {
0271       const auto& hit    = hits[i];
0272       const auto& weight = weights[i];
0273       energy += hit.getEnergy() * weight;
0274       energyError += std::pow(hit.getEnergyError() * weight, 2);
0275       // energy weighting for the other variables
0276       const double energyWeight = hit.getEnergy() * weight;
0277       time += hit.getTime() * energyWeight;
0278       timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0279       meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0280       mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0281       r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0282       cluster.addToHits(hit);
0283     }
0284     cluster.setEnergy(energy);
0285     cluster.setEnergyError(std::sqrt(energyError));
0286     cluster.setTime(time / energy);
0287     cluster.setTimeError(std::sqrt(timeError) / energy);
0288     cluster.setNhits(hits.size());
0289     cluster.setPosition(edm4hep::utils::sphericalToVector(r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0290
0291     // shower radius estimate (eta-phi plane)
0292     double radius = 0.;
0293     for (const auto& hit : hits) {
0294       radius += pow2(edm4hep::utils::eta(hit.getPosition()) - edm4hep::utils::eta(cluster.getPosition())) +
0295                 pow2(edm4hep::utils::angleAzimuthal(hit.getPosition()) - edm4hep::utils::angleAzimuthal(cluster.getPosition()));
0296     }
0297     cluster.addToShapeParameters(std::sqrt(radius / cluster.getNhits()));
0298     // Skewedness not calculated TODO
0299
0300     // Optionally store the MC truth associated with the first hit in this cluster
0301     // FIXME no connection between cluster and truth in edm4hep
0302     // if (mcHits) {
0303     //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0304     //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0305     //}
0306
0307     return cluster;
0308   }
0309
0310   std::pair<double /* polar */, double /* azimuthal */> fit_track(const std::vector<edm4eic::MutableCluster>& layers) const {
0311     int nrows = 0;
0312     decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0313     for (const auto& layer : layers) {
0314       if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0315         mean_pos = mean_pos + layer.getPosition();
0316         nrows += 1;
0317       }
0318     }
0319     // cannot fit
0320     if (nrows < 2) {
0321       return {};
0322     }
0323
0324     mean_pos = mean_pos / nrows;
0325     // fill position data
0326     Eigen::MatrixXd pos(nrows, 3);
0327     int ir = 0;
0328     for (const auto& layer : layers) {
0329       if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_trackStopLayer)) {
0330         auto delta = layer.getPosition() - mean_pos;
0331         pos(ir, 0) = delta.x;
0332         pos(ir, 1) = delta.y;
0333         pos(ir, 2) = delta.z;
0334         ir += 1;
0335       }
0336     }
0337
0338     Eigen::JacobiSVD<Eigen::MatrixXd> svd(pos, Eigen::ComputeThinU | Eigen::ComputeThinV);
0339     const auto dir = svd.matrixV().col(0);
0340     // theta and phi
0341     return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0342   }
0343 };
0344
0345 // NOLINTNEXTLINE(cppcoreguidelines-avoid-non-const-global-variables)
0346 DECLARE_COMPONENT(ImagingClusterReco)
0347
0348 } // namespace Jug::Reco