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