algorithms/calorimetry/CalorimeterClusterShape.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2025 Chao Peng, Dhevan Gangadharan, Sebouh Paul, Derek Anderson
0003
0004 #include "CalorimeterClusterShape.h"
0005
0006 #include <boost/algorithm/string/join.hpp>
0007 #include <boost/iterator/iterator_facade.hpp>
0008 #include <boost/range/adaptor/map.hpp>
0009 #include <edm4hep/Vector3f.h>
0010 #include <edm4hep/utils/vector_utils.h>
0011 #include <Eigen/Core>
0012 #include <Eigen/Eigenvalues>
0013 #include <Eigen/Householder>
0014 #include <cmath>
0015
0016 namespace eicrecon {
0017
0018   void CalorimeterClusterShape::init() {
0019
0020     // select weighting method
0021     std::string ew = m_cfg.energyWeight;
0022
0023     // make it case-insensitive
0024     std::transform(ew.begin(), ew.end(), ew.begin(), [](char s) { return std::tolower(s); });
0025     auto it = m_weightMethods.find(ew);
0026     if (it == m_weightMethods.end()) {
0027       error(
0028         "Cannot find energy weighting method {}, choose one from [{}]",
0029          m_cfg.energyWeight,
0030          boost::algorithm::join(m_weightMethods | boost::adaptors::map_keys, ", ")
0031       );
0032     } else {
0033       m_weightFunc = it->second;
0034     }
0035
0036   }  // end 'init()'
0037
0038   /*! Primary algorithm call: algorithm ingests a collection of clusters
0039    *  and computes their cluster shape parameters.  Clusters are copied
0040    *  onto output with computed shape parameters.  If associations are
0041    *  provided, they are copied to the output.
0042    *
0043    *  Parameters calculated:
0044    *    - radius,
0045    *    - dispersion (energy weighted radius),
0046    *    - theta-phi cluster widths (2D)
0047    *    - x-y-z cluster widths (3D)
0048    */
0049   void CalorimeterClusterShape::process(
0050     const CalorimeterClusterShape::Input& input,
0051     const CalorimeterClusterShape::Output& output
0052   ) const {
0053
0054     // grab inputs/outputs
0055     const auto [in_clusters, in_associations] = input;
0056     auto [out_clusters, out_associations] = output;
0057
0058     // exit if no clusters in collection
0059     if (in_clusters->size() == 0) {
0060       debug("No clusters in input collection.");
0061       return;
0062     }
0063
0064     // loop over input clusters
0065     for (const auto& in_clust : *in_clusters) {
0066
0067       // copy input cluster
0068       edm4eic::MutableCluster out_clust = in_clust.clone();
0069
0070       // set up base for weights
0071       double logWeightBase=m_cfg.logWeightBase;
0072       if (m_cfg.logWeightBaseCoeffs.size() != 0){
0073         double l=log(out_clust.getEnergy()/m_cfg.logWeightBase_Eref);
0074         logWeightBase=0;
0075         for(std::size_t i =0; i<m_cfg.logWeightBaseCoeffs.size(); i++){
0076           logWeightBase += m_cfg.logWeightBaseCoeffs[i]*pow(l,i);
0077         }
0078       }
0079
0080       // ----------------------------------------------------------------------
0081       // do shape parameter calculation
0082       // ----------------------------------------------------------------------
0083       {
0084
0085         // create addresses for quantities we'll need later
0086         float radius = 0, dispersion = 0, w_sum = 0;
0087
0088         // set up matrices/vectors
0089         Eigen::Matrix2f sum2_2D = Eigen::Matrix2f::Zero();
0090         Eigen::Matrix3f sum2_3D = Eigen::Matrix3f::Zero();
0091         Eigen::Vector2f sum1_2D = Eigen::Vector2f::Zero();
0092         Eigen::Vector3f sum1_3D = Eigen::Vector3f::Zero();
0093         Eigen::Vector2cf eigenValues_2D = Eigen::Vector2cf::Zero();
0094         Eigen::Vector3cf eigenValues_3D = Eigen::Vector3cf::Zero();
0095
0096         // the axis is the direction of the eigenvalue corresponding to the largest eigenvalue.
0097         edm4hep::Vector3f axis;
0098         if (out_clust.getNhits() > 1) {
0099           for (std::size_t iHit = 0; const auto& hit : out_clust.getHits()) {
0100
0101             // get weight of hit
0102             const double eTotal = out_clust.getEnergy() * m_cfg.sampFrac;
0103             const float w = m_weightFunc(hit.getEnergy(), eTotal, logWeightBase, 0);
0104
0105             // theta, phi
0106             Eigen::Vector2f pos2D( edm4hep::utils::anglePolar( hit.getPosition() ), edm4hep::utils::angleAzimuthal( hit.getPosition() ) );
0107             // x, y, z
0108             Eigen::Vector3f pos3D( hit.getPosition().x, hit.getPosition().y, hit.getPosition().z );
0109
0110             const auto delta = out_clust.getPosition() - hit.getPosition();
0111             radius     += delta * delta;
0112             dispersion += delta * delta * w;
0113
0114             // Weighted Sum x*x, x*y, x*z, y*y, etc.
0115             sum2_2D += w * pos2D * pos2D.transpose();
0116             sum2_3D += w * pos3D * pos3D.transpose();
0117
0118             // Weighted Sum x, y, z
0119             sum1_2D += w * pos2D;
0120             sum1_3D += w * pos3D;
0121
0122             w_sum += w;
0123             ++iHit;
0124           }  // end hit loop
0125
0126           radius = sqrt((1. / (out_clust.getNhits() - 1.)) * radius);
0127           if( w_sum > 0 ) {
0128             dispersion = sqrt( dispersion / w_sum );
0129
0130             // normalize matrices
0131             sum2_2D /= w_sum;
0132             sum2_3D /= w_sum;
0133             sum1_2D /= w_sum;
0134             sum1_3D /= w_sum;
0135
0136             // 2D and 3D covariance matrices
0137             Eigen::Matrix2f cov2 = sum2_2D - sum1_2D * sum1_2D.transpose();
0138             Eigen::Matrix3f cov3 = sum2_3D - sum1_3D * sum1_3D.transpose();
0139
0140             // Solve for eigenvalues.  Corresponds to out_cluster's 2nd moments (widths)
0141             Eigen::EigenSolver<Eigen::Matrix2f> es_2D(cov2, false); // set to true for eigenvector calculation
0142             Eigen::EigenSolver<Eigen::Matrix3f> es_3D(cov3, true); // set to true for eigenvector calculation
0143
0144             // eigenvalues of symmetric real matrix are always real
0145             eigenValues_2D = es_2D.eigenvalues();
0146             eigenValues_3D = es_3D.eigenvalues();
0147             //find the eigenvector corresponding to the largest eigenvalue
0148             auto eigenvectors= es_3D.eigenvectors();
0149             auto max_eigenvalue_it = std::max_element(
0150               eigenValues_3D.begin(),
0151               eigenValues_3D.end(),
0152               [](auto a, auto b) {
0153                   return std::real(a) < std::real(b);
0154               }
0155             );
0156             auto axis_eigen = eigenvectors.col(std::distance(
0157                   eigenValues_3D.begin(),
0158                   max_eigenvalue_it
0159               ));
0160             axis = {
0161               axis_eigen(0,0).real(),
0162               axis_eigen(1,0).real(),
0163               axis_eigen(2,0).real(),
0164             };
0165           }  // end if weight sum is nonzero
0166         }  // end if n hits > 1
0167
0168         // set shape parameters
0169         out_clust.addToShapeParameters( radius );
0170         out_clust.addToShapeParameters( dispersion );
0171         out_clust.addToShapeParameters( eigenValues_2D[0].real() ); // 2D theta-phi out_cluster width 1
0172         out_clust.addToShapeParameters( eigenValues_2D[1].real() ); // 2D theta-phi out_cluster width 2
0173         out_clust.addToShapeParameters( eigenValues_3D[0].real() ); // 3D x-y-z out_cluster width 1
0174         out_clust.addToShapeParameters( eigenValues_3D[1].real() ); // 3D x-y-z out_cluster width 2
0175         out_clust.addToShapeParameters( eigenValues_3D[2].real() ); // 3D x-y-z out_cluster width 3
0176
0177         // check axis orientation
0178         double dot_product = out_clust.getPosition() * axis;
0179         if (dot_product < 0) {
0180           axis = -1 * axis;
0181         }
0182
0183         // set intrinsic theta/phi
0184         if (m_cfg.longitudinalShowerInfoAvailable) {
0185           out_clust.setIntrinsicTheta( edm4hep::utils::anglePolar(axis) );
0186           out_clust.setIntrinsicPhi( edm4hep::utils::angleAzimuthal(axis) );
0187           // TODO intrinsicDirectionError
0188         } else {
0189           out_clust.setIntrinsicTheta(NAN);
0190           out_clust.setIntrinsicPhi(NAN);
0191         }
0192       }  // end shape parameter calculation
0193
0194       out_clusters->push_back(out_clust);
0195       trace("Completed shape calculation for cluster {}", in_clust.getObjectID().index);
0196
0197       // ----------------------------------------------------------------------
0198       // if provided, copy associations
0199       // ----------------------------------------------------------------------
0200       for (auto in_assoc : *in_associations) {
0201         if (in_assoc.getRec() == in_clust) {
0202           auto mc_par = in_assoc.getSim();
0203           auto out_assoc = out_associations->create();
0204           out_assoc.setRecID( out_clust.getObjectID().index );
0205           out_assoc.setSimID( mc_par.getObjectID().index );
0206           out_assoc.setRec( out_clust );
0207           out_assoc.setSim( mc_par );
0208           out_assoc.setWeight( in_assoc.getWeight() );
0209         }
0210       }  // end input association loop
0211     }  // end input cluster loop
0212     debug("Completed processing input clusters");
0213
0214   }  // end 'process(Input&, Output&)'
0215
0216 }  // end eicrecon namespace