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