algorithms/calorimetry/CalorimeterClusterRecoCoG.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 Sylvester Joosten, Chao, Chao Peng, Whitney Armstrong, Dhevan Gangadharan
0003
0004 /*
0005  *  Reconstruct the cluster with Center of Gravity method
0006  *  Logarithmic weighting is used for mimicking energy deposit in transverse direction
0007  *
0008  *  Author: Chao Peng (ANL), 09/27/2020
0009  */
0010
0011 #include <Evaluator/DD4hepUnits.h>
0012 #include <boost/algorithm/string/join.hpp>
0013 #include <boost/iterator/iterator_facade.hpp>
0014 #include <boost/range/adaptor/map.hpp>
0015 #include <edm4eic/CalorimeterHitCollection.h>
0016 #include <edm4hep/CaloHitContributionCollection.h>
0017 #include <edm4hep/MCParticleCollection.h>
0018 #include <edm4hep/Vector3f.h>
0019 #include <edm4hep/utils/vector_utils.h>
0020 #include <fmt/core.h>
0021 #include <podio/ObjectID.h>
0022 #include <podio/RelationRange.h>
0023 #include <Eigen/Core>
0024 #include <Eigen/Eigenvalues>
0025 #include <Eigen/Householder> // IWYU pragma: keep
0026 #include <cctype>
0027 #include <complex>
0028 #include <cstddef>
0029 #include <gsl/pointers>
0030 #include <iterator>
0031 #include <limits>
0032 #include <map>
0033 #include <optional>
0034 #include <vector>
0035
0036 #include "CalorimeterClusterRecoCoG.h"
0037 #include "algorithms/calorimetry/CalorimeterClusterRecoCoGConfig.h"
0038
0039 namespace eicrecon {
0040
0041   using namespace dd4hep;
0042
0043   void CalorimeterClusterRecoCoG::init() {
0044
0045     // select weighting method
0046     std::string ew = m_cfg.energyWeight;
0047     // make it case-insensitive
0048     std::transform(ew.begin(), ew.end(), ew.begin(), [](char s) { return std::tolower(s); });
0049     auto it = weightMethods.find(ew);
0050     if (it == weightMethods.end()) {
0051       error("Cannot find energy weighting method {}, choose one from [{}]", m_cfg.energyWeight, boost::algorithm::join(weightMethods | boost::adaptors::map_keys, ", "));
0052       return;
0053     }
0054     weightFunc = it->second;
0055   }
0056
0057   void CalorimeterClusterRecoCoG::process(
0058       const CalorimeterClusterRecoCoG::Input& input,
0059       const CalorimeterClusterRecoCoG::Output& output) const {
0060
0061     const auto [proto, mchits] = input;
0062     auto [clusters, associations] = output;
0063
0064     for (const auto& pcl : *proto) {
0065
0066       // skip protoclusters with no hits
0067       if (pcl.hits_size() == 0) {
0068         continue;
0069       }
0070
0071       auto cl_opt = reconstruct(pcl);
0072       if (! cl_opt.has_value()) {
0073         continue;
0074       }
0075       auto cl = *std::move(cl_opt);
0076
0077       debug("{} hits: {} GeV, ({}, {}, {})", cl.getNhits(), cl.getEnergy() / dd4hep::GeV, cl.getPosition().x / dd4hep::mm, cl.getPosition().y / dd4hep::mm, cl.getPosition().z / dd4hep::mm);
0078       clusters->push_back(cl);
0079
0080       // If mcHits are available, associate cluster with MCParticle
0081       // 1. find proto-cluster hit with largest energy deposition
0082       // 2. find first mchit with same CellID
0083       // 3. assign mchit's MCParticle as cluster truth
0084       if (mchits->size() > 0) {
0085
0086         // 1. find pclhit with largest energy deposition
0087         auto pclhits = pcl.getHits();
0088         auto pclhit = std::max_element(
0089           pclhits.begin(),
0090           pclhits.end(),
0091           [](const auto& pclhit1, const auto& pclhit2) {
0092             return pclhit1.getEnergy() < pclhit2.getEnergy();
0093           }
0094         );
0095
0096         // 2. find mchit with same CellID
0097         // find_if not working, https://github.com/AIDASoft/podio/pull/273
0098         //auto mchit = std::find_if(
0099         //  mchits->begin(),
0100         //  mchits->end(),
0101         //  [&pclhit](const auto& mchit1) {
0102         //    return mchit1.getCellID() == pclhit->getCellID();
0103         //  }
0104         //);
0105         auto mchit = mchits->begin();
0106         for ( ; mchit != mchits->end(); ++mchit) {
0107           // break loop when CellID match found
0108           if ( mchit->getCellID() == pclhit->getCellID()) {
0109             break;
0110           }
0111         }
0112         if (!(mchit != mchits->end())) {
0113           // break if no matching hit found for this CellID
0114           warning("Proto-cluster has highest energy in CellID {}, but no mc hit with that CellID was found.", pclhit->getCellID());
0115           trace("Proto-cluster hits: ");
0116           for (const auto& pclhit1: pclhits) {
0117             trace("{}: {}", pclhit1.getCellID(), pclhit1.getEnergy());
0118           }
0119           trace("MC hits: ");
0120           for (const auto& mchit1: *mchits) {
0121             trace("{}: {}", mchit1.getCellID(), mchit1.getEnergy());
0122           }
0123           break;
0124         }
0125
0126         // 3. find mchit's MCParticle
0127         const auto& mcp = mchit->getContributions(0).getParticle();
0128
0129         debug("cluster has largest energy in cellID: {}", pclhit->getCellID());
0130         debug("pcl hit with highest energy {} at index {}", pclhit->getEnergy(), pclhit->getObjectID().index);
0131         debug("corresponding mc hit energy {} at index {}", mchit->getEnergy(), mchit->getObjectID().index);
0132         debug("from MCParticle index {}, PDG {}, {}", mcp.getObjectID().index, mcp.getPDG(), edm4hep::utils::magnitude(mcp.getMomentum()));
0133
0134         // set association
0135         auto clusterassoc = associations->create();
0136         clusterassoc.setRecID(cl.getObjectID().index); // if not using collection, this is always set to -1
0137         clusterassoc.setSimID(mcp.getObjectID().index);
0138         clusterassoc.setWeight(1.0);
0139         clusterassoc.setRec(cl);
0140         clusterassoc.setSim(mcp);
0141       } else {
0142         debug("No mcHitCollection was provided, so no truth association will be performed.");
0143       }
0144     }
0145 }
0146
0147 //------------------------------------------------------------------------
0148 std::optional<edm4eic::MutableCluster> CalorimeterClusterRecoCoG::reconstruct(const edm4eic::ProtoCluster& pcl) const {
0149   edm4eic::MutableCluster cl;
0150   cl.setNhits(pcl.hits_size());
0151
0152   debug("hit size = {}", pcl.hits_size());
0153
0154   // no hits
0155   if (pcl.hits_size() == 0) {
0156     return {};
0157   }
0158
0159   // calculate total energy, find the cell with the maximum energy deposit
0160   float totalE = 0.;
0161   // Used to optionally constrain the cluster eta to those of the contributing hits
0162   float minHitEta = std::numeric_limits<float>::max();
0163   float maxHitEta = std::numeric_limits<float>::min();
0164   auto time       = pcl.getHits()[0].getTime();
0165   auto timeError  = pcl.getHits()[0].getTimeError();
0166   for (unsigned i = 0; i < pcl.getHits().size(); ++i) {
0167     const auto& hit   = pcl.getHits()[i];
0168     const auto weight = pcl.getWeights()[i];
0169     debug("hit energy = {} hit weight: {}", hit.getEnergy(), weight);
0170     auto energy = hit.getEnergy() * weight;
0171     totalE += energy;
0172     cl.addToHits(hit);
0173     cl.addToHitContributions(energy);
0174     const float eta = edm4hep::utils::eta(hit.getPosition());
0175     if (eta < minHitEta) {
0176       minHitEta = eta;
0177     }
0178     if (eta > maxHitEta) {
0179       maxHitEta = eta;
0180     }
0181   }
0182   cl.setEnergy(totalE / m_cfg.sampFrac);
0183   cl.setEnergyError(0.);
0184   cl.setTime(time);
0185   cl.setTimeError(timeError);
0186
0187   // center of gravity with logarithmic weighting
0188   float tw = 0.;
0189   auto v   = cl.getPosition();
0190
0191   double logWeightBase=m_cfg.logWeightBase;
0192   if (m_cfg.logWeightBaseCoeffs.size() != 0){
0193     double l=log(cl.getEnergy()/m_cfg.logWeightBase_Eref);
0194     logWeightBase=0;
0195     for(std::size_t i =0; i<m_cfg.logWeightBaseCoeffs.size(); i++){
0196       logWeightBase += m_cfg.logWeightBaseCoeffs[i]*pow(l,i);
0197     }
0198   }
0199
0200   for (unsigned i = 0; i < pcl.getHits().size(); ++i) {
0201     const auto& hit   = pcl.getHits()[i];
0202     const auto weight = pcl.getWeights()[i];
0203     //      _DBG_<<" -- weight = " << weight << "  E=" << hit.getEnergy() << " totalE=" <<totalE << " log(E/totalE)=" << std::log(hit.getEnergy()/totalE) << std::endl;
0204     float w           = weightFunc(hit.getEnergy() * weight, totalE, logWeightBase, 0);
0205     tw += w;
0206     v = v + (hit.getPosition() * w);
0207   }
0208   if (tw == 0.) {
0209     warning("zero total weights encountered, you may want to adjust your weighting parameter.");
0210     return {};
0211   }
0212   cl.setPosition(v / tw);
0213   cl.setPositionError({}); // @TODO: Covariance matrix
0214
0215   // Optionally constrain the cluster to the hit eta values
0216   if (m_cfg.enableEtaBounds) {
0217     const bool overflow  = (edm4hep::utils::eta(cl.getPosition()) > maxHitEta);
0218     const bool underflow = (edm4hep::utils::eta(cl.getPosition()) < minHitEta);
0219     if (overflow || underflow) {
0220       const double newEta   = overflow ? maxHitEta : minHitEta;
0221       const double newTheta = edm4hep::utils::etaToAngle(newEta);
0222       const double newR     = edm4hep::utils::magnitude(cl.getPosition());
0223       const double newPhi   = edm4hep::utils::angleAzimuthal(cl.getPosition());
0224       cl.setPosition(edm4hep::utils::sphericalToVector(newR, newTheta, newPhi));
0225       debug("Bound cluster position to contributing hits due to {}", (overflow ? "overflow" : "underflow"));
0226     }
0227   }
0228
0229   // Additional convenience variables
0230
0231   // best estimate on the cluster direction is the cluster position
0232   // for simple 2D CoG clustering
0233   cl.setIntrinsicTheta(edm4hep::utils::anglePolar(cl.getPosition()));
0234   cl.setIntrinsicPhi(edm4hep::utils::angleAzimuthal(cl.getPosition()));
0235   // TODO errors
0236
0237   //_______________________________________
0238   // Calculate cluster profile:
0239   //    radius,
0240   //    dispersion (energy weighted radius),
0241   //    theta-phi cluster widths (2D)
0242   //    x-y-z cluster widths (3D)
0243   float radius = 0, dispersion = 0, w_sum = 0;
0244
0245   Eigen::Matrix2f sum2_2D = Eigen::Matrix2f::Zero();
0246   Eigen::Matrix3f sum2_3D = Eigen::Matrix3f::Zero();
0247   Eigen::Vector2f sum1_2D = Eigen::Vector2f::Zero();
0248   Eigen::Vector3f sum1_3D = Eigen::Vector3f::Zero();
0249   Eigen::Vector2cf eigenValues_2D = Eigen::Vector2cf::Zero();
0250   Eigen::Vector3cf eigenValues_3D = Eigen::Vector3cf::Zero();
0251   //the axis is the direction of the eigenvalue corresponding to the largest eigenvalue.
0252   double axis_x=0, axis_y=0, axis_z=0;
0253   if (cl.getNhits() > 1) {
0254
0255     for (const auto& hit : pcl.getHits()) {
0256
0257       float w = weightFunc(hit.getEnergy(), totalE, logWeightBase, 0);
0258
0259       // theta, phi
0260       Eigen::Vector2f pos2D( edm4hep::utils::anglePolar( hit.getPosition() ), edm4hep::utils::angleAzimuthal( hit.getPosition() ) );
0261       // x, y, z
0262       Eigen::Vector3f pos3D( hit.getPosition().x, hit.getPosition().y, hit.getPosition().z );
0263
0264       const auto delta = cl.getPosition() - hit.getPosition();
0265       radius          += delta * delta;
0266       dispersion      += delta * delta * w;
0267
0268       // Weighted Sum x*x, x*y, x*z, y*y, etc.
0269       sum2_2D += w * pos2D * pos2D.transpose();
0270       sum2_3D += w * pos3D * pos3D.transpose();
0271
0272       // Weighted Sum x, y, z
0273       sum1_2D += w * pos2D;
0274       sum1_3D += w * pos3D;
0275
0276       w_sum += w;
0277     }
0278
0279     radius     = sqrt((1. / (cl.getNhits() - 1.)) * radius);
0280     if( w_sum > 0 ) {
0281       dispersion = sqrt( dispersion / w_sum );
0282
0283       // normalize matrices
0284       sum2_2D /= w_sum;
0285       sum2_3D /= w_sum;
0286       sum1_2D /= w_sum;
0287       sum1_3D /= w_sum;
0288
0289       // 2D and 3D covariance matrices
0290       Eigen::Matrix2f cov2 = sum2_2D - sum1_2D * sum1_2D.transpose();
0291       Eigen::Matrix3f cov3 = sum2_3D - sum1_3D * sum1_3D.transpose();
0292
0293       // Solve for eigenvalues.  Corresponds to cluster's 2nd moments (widths)
0294       Eigen::EigenSolver<Eigen::Matrix2f> es_2D(cov2, false); // set to true for eigenvector calculation
0295       Eigen::EigenSolver<Eigen::Matrix3f> es_3D(cov3, true); // set to true for eigenvector calculation
0296
0297       // eigenvalues of symmetric real matrix are always real
0298       eigenValues_2D = es_2D.eigenvalues();
0299       eigenValues_3D = es_3D.eigenvalues();
0300       //find the eigenvector corresponding to the largest eigenvalue
0301       auto eigenvectors= es_3D.eigenvectors();
0302       auto max_eigenvalue_it = std::max_element(
0303         eigenValues_3D.begin(),
0304         eigenValues_3D.end(),
0305         [](auto a, auto b) {
0306             return std::real(a) < std::real(b);
0307         }
0308       );
0309       auto axis = eigenvectors.col(std::distance(
0310             eigenValues_3D.begin(),
0311             max_eigenvalue_it
0312         ));
0313       axis_x=axis(0,0).real();
0314       axis_y=axis(1,0).real();
0315       axis_z=axis(2,0).real();
0316       double norm=sqrt(axis_x*axis_x+axis_y*axis_y+axis_z*axis_z);
0317       if (norm!=0){
0318         axis_x/=norm;
0319         axis_y/=norm;
0320         axis_z/=norm;
0321       }
0322     }
0323   }
0324
0325   cl.addToShapeParameters( radius );
0326   cl.addToShapeParameters( dispersion );
0327   cl.addToShapeParameters( eigenValues_2D[0].real() ); // 2D theta-phi cluster width 1
0328   cl.addToShapeParameters( eigenValues_2D[1].real() ); // 2D theta-phi cluster width 2
0329   cl.addToShapeParameters( eigenValues_3D[0].real() ); // 3D x-y-z cluster width 1
0330   cl.addToShapeParameters( eigenValues_3D[1].real() ); // 3D x-y-z cluster width 2
0331   cl.addToShapeParameters( eigenValues_3D[2].real() ); // 3D x-y-z cluster width 3
0332   //last 3 shape parameters are the components of the axis direction
0333   cl.addToShapeParameters( axis_x );
0334   cl.addToShapeParameters( axis_y );
0335   cl.addToShapeParameters( axis_z );
0336   return std::move(cl);
0337 }
0338
0339 } // eicrecon