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