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 <edm4eic/EDM4eicVersion.h>
0010 #include <edm4hep/MCParticle.h>
0011 #include <edm4hep/Vector3f.h>
0012 #include <edm4hep/utils/vector_utils.h>
0013 #include <podio/LinkNavigator.h>
0014 #include <podio/RelationRange.h>
0015 #include <podio/detail/LinkCollectionImpl.h>
0016 #include <Eigen/Core>
0017 #include <Eigen/Eigenvalues>
0018 #include <Eigen/Householder> // IWYU pragma: keep
0019 #include <Eigen/Jacobi>
0020 #include <algorithm>
0021 #include <cctype>
0022 #include <cmath>
0023 #include <cstddef>
0024 #include <gsl/pointers>
0025 #include <tuple>
0026 #include <utility>
0027 #include <vector>
0028
0029 #include "algorithms/calorimetry/CalorimeterClusterShapeConfig.h"
0030 #include "algorithms/interfaces/LinkTruthUtils.h"
0031
0032 namespace eicrecon {
0033
0034 void CalorimeterClusterShape::init() {
0035
0036   // select weighting method
0037   std::string ew = m_cfg.energyWeight;
0038
0039   // make it case-insensitive
0040   std::ranges::transform(ew, ew.begin(), [](char s) { return std::tolower(s); });
0041   auto it = m_weightMethods.find(ew);
0042   if (it == m_weightMethods.end()) {
0043     error("Cannot find energy weighting method {}, choose one from [{}]", m_cfg.energyWeight,
0044           boost::algorithm::join(m_weightMethods | boost::adaptors::map_keys, ", "));
0045   } else {
0046     m_weightFunc = it->second;
0047   }
0048
0049 } // end 'init()'
0050
0051 /*! Primary algorithm call: algorithm ingests a collection of clusters
0052    *  and computes their cluster shape parameters.  Clusters are copied
0053    *  onto output with computed shape parameters.  If associations are
0054    *  provided, they are copied to the output.
0055    *
0056    *  Parameters calculated:
0057    *    - radius,
0058    *    - dispersion (energy weighted radius),
0059    *    - theta-phi cluster widths (2D)
0060    *    - x-y-z cluster widths (3D)
0061    */
0062 void CalorimeterClusterShape::process(const CalorimeterClusterShape::Input& input,
0063                                       const CalorimeterClusterShape::Output& output) const {
0064
0065   // grab inputs/outputs
0066   const auto [in_clusters, in_links]               = input;
0067   auto [out_clusters, out_links, out_associations] = output;
0068
0069   const truth::EventLinkNavigator<edm4eic::MCRecoClusterParticleLinkCollection> link_nav(in_links);
0070
0071   // exit if no clusters in collection
0072   if (in_clusters->empty()) {
0073     debug("No clusters in input collection.");
0074     return;
0075   }
0076
0077   // loop over input clusters
0078   for (const auto& in_clust : *in_clusters) {
0079
0080     // copy input cluster
0081     edm4eic::MutableCluster out_clust = in_clust.clone();
0082
0083     // set up base for weights
0084     double logWeightBase = m_cfg.logWeightBase;
0085     if (!m_cfg.logWeightBaseCoeffs.empty()) {
0086       double l      = std::log(out_clust.getEnergy() / m_cfg.logWeightBase_Eref);
0087       logWeightBase = 0;
0088       for (std::size_t i = 0; i < m_cfg.logWeightBaseCoeffs.size(); i++) {
0089         logWeightBase += m_cfg.logWeightBaseCoeffs[i] * pow(l, i);
0090       }
0091     }
0092
0093     // ----------------------------------------------------------------------
0094     // do shape parameter calculation
0095     // ----------------------------------------------------------------------
0096     {
0097
0098       // create addresses for quantities we'll need later
0099       double radius     = 0;
0100       double dispersion = 0;
0101       double w_sum      = 0;
0102       // set up matrices/vectors
0103       Eigen::Matrix2d sum2_2D        = Eigen::Matrix2d::Zero();
0104       Eigen::Matrix3d sum2_3D        = Eigen::Matrix3d::Zero();
0105       Eigen::Vector2d sum1_2D        = Eigen::Vector2d::Zero();
0106       Eigen::Vector3d sum1_3D        = Eigen::Vector3d::Zero();
0107       Eigen::Vector2d eigenValues_2D = Eigen::Vector2d::Zero();
0108       Eigen::Vector3d eigenValues_3D = Eigen::Vector3d::Zero();
0109
0110       // the axis is the direction of the eigenvalue corresponding to the largest eigenvalue.
0111       edm4hep::Vector3f axis;
0112       if (out_clust.getNhits() > 1) {
0113         for (const auto& hit : out_clust.getHits()) {
0114
0115           // get weight of hit
0116           const double eTotal = out_clust.getEnergy() * m_cfg.sampFrac;
0117           const double w      = m_weightFunc(hit.getEnergy(), eTotal, logWeightBase, 0);
0118
0119           // theta, phi
0120           Eigen::Vector2d pos2D(edm4hep::utils::anglePolar(hit.getPosition()),
0121                                 edm4hep::utils::angleAzimuthal(hit.getPosition()));
0122           // x, y, z
0123           Eigen::Vector3d pos3D(hit.getPosition().x, hit.getPosition().y, hit.getPosition().z);
0124           const auto delta = out_clust.getPosition() - hit.getPosition();
0125           radius += delta * delta;
0126           dispersion += delta * delta * w;
0127
0128           // Weighted Sum x*x, x*y, x*z, y*y, etc.
0129           sum2_2D += w * pos2D * pos2D.transpose();
0130           sum2_3D += w * pos3D * pos3D.transpose();
0131
0132           // Weighted Sum x, y, z
0133           sum1_2D += w * pos2D;
0134           sum1_3D += w * pos3D;
0135
0136           w_sum += w;
0137         } // end hit loop
0138
0139         radius = sqrt((1. / (out_clust.getNhits() - 1.)) * radius);
0140         if (w_sum > 0) {
0141           dispersion = sqrt(dispersion / w_sum);
0142
0143           // normalize matrices
0144           sum2_2D /= w_sum;
0145           sum2_3D /= w_sum;
0146           sum1_2D /= w_sum;
0147           sum1_3D /= w_sum;
0148
0149           // 2D and 3D covariance matrices
0150           Eigen::Matrix2d cov2 = sum2_2D - sum1_2D * sum1_2D.transpose();
0151           Eigen::Matrix3d cov3 = sum2_3D - sum1_3D * sum1_3D.transpose();
0152
0153           // Use SelfAdjointEigenSolver for symmetric covariance matrices.
0154           // More accurate than EigenSolver for symmetric matrices, guarantees
0155           // real eigenvalues, and returns them sorted ascending: [0]=smallest.
0156           Eigen::SelfAdjointEigenSolver<Eigen::Matrix2d> es_2D(cov2);
0157           Eigen::SelfAdjointEigenSolver<Eigen::Matrix3d> es_3D(cov3);
0158
0159           // eigenvalues of symmetric real matrix are always real
0160           // Store descending: [0]=largest, [1/2]=smaller
0161           auto ev2          = es_2D.eigenvalues(); // ascending real double
0162           eigenValues_2D[0] = ev2[1];              // largest
0163           eigenValues_2D[1] = ev2[0];              // smallest
0164
0165           auto ev3          = es_3D.eigenvalues(); // ascending real double
0166           eigenValues_3D[0] = ev3[2];              // largest
0167           eigenValues_3D[1] = ev3[1];
0168           eigenValues_3D[2] = ev3[0]; // smallest (0 for flat-z detectors)
0169
0170           // eigenvector for largest eigenvalue (index 2 in ascending order)
0171           auto axis_eigen = es_3D.eigenvectors().col(2);
0172           axis            = {
0173               static_cast<float>(axis_eigen(0)),
0174               static_cast<float>(axis_eigen(1)),
0175               static_cast<float>(axis_eigen(2)),
0176           };
0177         } // end if weight sum is nonzero
0178       } // end if n hits > 1
0179
0180       // set shape parameters
0181       // NOTE: shapeParameters stores raw covariance-matrix values ([mm^2], [rad^2]),
0182       // kept as-is for backward compatibility. shapeParameters will go away eventually.
0183       out_clust.addToShapeParameters(radius);
0184       out_clust.addToShapeParameters(dispersion);
0185       out_clust.addToShapeParameters(eigenValues_2D[0]); // 2D theta-phi out_cluster width 1 [rad^2]
0186       out_clust.addToShapeParameters(eigenValues_2D[1]); // 2D theta-phi out_cluster width 2 [rad^2]
0187       out_clust.addToShapeParameters(eigenValues_3D[0]); // 3D x-y-z out_cluster width 1 [mm^2]
0188       out_clust.addToShapeParameters(eigenValues_3D[1]); // 3D x-y-z out_cluster width 2 [mm^2]
0189       out_clust.addToShapeParameters(eigenValues_3D[2]); // 3D x-y-z out_cluster width 3 [mm^2]
0190
0191 #if EDM4EIC_BUILD_VERSION >= EDM4EIC_VERSION(8, 10, 0)
0192       // set dedicated shape variables (sqrt of covariance-matrix values, [mm]/[rad] units)
0193       out_clust.setRadius(static_cast<float>(radius));
0194       out_clust.setDispersion(static_cast<float>(dispersion));
0195       out_clust.setPrincipalAxesLengthsXYZ({
0196           static_cast<float>(std::sqrt(std::abs(eigenValues_3D[0]))),
0197           static_cast<float>(std::sqrt(std::abs(eigenValues_3D[1]))),
0198           static_cast<float>(std::sqrt(std::abs(eigenValues_3D[2]))),
0199       });
0200       out_clust.setPrincipalAxesLengthsThetaPhi({
0201           static_cast<float>(std::sqrt(std::abs(eigenValues_2D[0]))),
0202           static_cast<float>(std::sqrt(std::abs(eigenValues_2D[1]))),
0203       });
0204 #endif
0205
0206       // check axis orientation
0207       double dot_product = out_clust.getPosition() * axis;
0208       if (dot_product < 0) {
0209         axis = -1 * axis;
0210       }
0211
0212       // set intrinsic theta/phi from 3D principal axis
0213       float intrinsicTheta = edm4hep::utils::anglePolar(axis);
0214       float intrinsicPhi   = edm4hep::utils::angleAzimuthal(axis);
0215       out_clust.setIntrinsicTheta(intrinsicTheta);
0216       out_clust.setIntrinsicPhi(intrinsicPhi);
0217       // TODO intrinsicDirectionError
0218
0219       trace("ClusterShape: radius={:.3f} [mm] dispersion={:.3f} [mm] "
0220             "2D w1={:.4f} w2={:.4f} [rad] "
0221             "3D w1={:.3f} w2={:.3f} w3={:.3f} [mm] "
0222             "intrinsicTheta={:.4f} Phi={:.4f} [rad]",
0223             radius, dispersion, std::sqrt(std::abs(eigenValues_2D[0])),
0224             std::sqrt(std::abs(eigenValues_2D[1])), std::sqrt(std::abs(eigenValues_3D[0])),
0225             std::sqrt(std::abs(eigenValues_3D[1])), std::sqrt(std::abs(eigenValues_3D[2])),
0226             intrinsicTheta, intrinsicPhi);
0227     } // end shape parameter calculation
0228
0229     out_clusters->push_back(out_clust);
0230
0231     // ----------------------------------------------------------------------
0232     // if provided, copy links and associations
0233     // ----------------------------------------------------------------------
0234     if (link_nav.enabled()) {
0235       for (const auto& [mc_par, weight] : link_nav.linked(in_clust)) {
0236         truth::addWeightedRelation(
0237             out_clust, mc_par, weight,
0238             gsl::not_null<edm4eic::MCRecoClusterParticleLinkCollection*>{out_links},
0239             gsl::not_null<edm4eic::MCRecoClusterParticleAssociationCollection*>{out_associations});
0240       }
0241     } // end input link loop
0242   } // end input cluster loop
0243   debug("Completed processing input clusters");
0244
0245 } // end 'process(Input&, Output&)'
0246
0247 } // namespace eicrecon