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 <algorithm>
0007 #include <boost/algorithm/string/join.hpp>
0008 #include <boost/range/adaptor/map.hpp>
0009 #include <edm4eic/CalorimeterHitCollection.h>
0010 #include <edm4hep/MCParticleCollection.h>
0011 #include <edm4hep/Vector3f.h>
0012 #include <edm4hep/utils/vector_utils.h>
0013 #include <fmt/core.h>
0014 #include <podio/ObjectID.h>
0015 #include <podio/RelationRange.h>
0016 #include <Eigen/Core>
0017 #include <Eigen/Eigenvalues>
0018 #include <Eigen/Householder> // IWYU pragma: keep
0019 #include <Eigen/Jacobi>
0020 #include <cctype>
0021 #include <cmath>
0022 #include <complex>
0023 #include <cstddef>
0024 #include <gsl/pointers>
0025 #include <iterator>
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   auto [out_clusters, out_associations]     = output;
0067
0068   // exit if no clusters in collection
0069   if (in_clusters->empty()) {
0070     debug("No clusters in input collection.");
0071     return;
0072   }
0073
0074   // loop over input clusters
0075   for (const auto& in_clust : *in_clusters) {
0076
0077     // copy input cluster
0078     edm4eic::MutableCluster out_clust = in_clust.clone();
0079
0080     // set up base for weights
0081     double logWeightBase = m_cfg.logWeightBase;
0082     if (!m_cfg.logWeightBaseCoeffs.empty()) {
0083       double l      = std::log(out_clust.getEnergy() / m_cfg.logWeightBase_Eref);
0084       logWeightBase = 0;
0085       for (std::size_t i = 0; i < m_cfg.logWeightBaseCoeffs.size(); i++) {
0086         logWeightBase += m_cfg.logWeightBaseCoeffs[i] * pow(l, i);
0087       }
0088     }
0089
0090     // ----------------------------------------------------------------------
0091     // do shape parameter calculation
0092     // ----------------------------------------------------------------------
0093     {
0094
0095       // create addresses for quantities we'll need later
0096       float radius     = 0;
0097       float dispersion = 0;
0098       float w_sum      = 0;
0099
0100       // set up matrices/vectors
0101       Eigen::Matrix2f sum2_2D         = Eigen::Matrix2f::Zero();
0102       Eigen::Matrix3f sum2_3D         = Eigen::Matrix3f::Zero();
0103       Eigen::Vector2f sum1_2D         = Eigen::Vector2f::Zero();
0104       Eigen::Vector3f sum1_3D         = Eigen::Vector3f::Zero();
0105       Eigen::Vector2cf eigenValues_2D = Eigen::Vector2cf::Zero();
0106       Eigen::Vector3cf eigenValues_3D = Eigen::Vector3cf::Zero();
0107
0108       // the axis is the direction of the eigenvalue corresponding to the largest eigenvalue.
0109       edm4hep::Vector3f axis;
0110       if (out_clust.getNhits() > 1) {
0111         for (const auto& hit : out_clust.getHits()) {
0112
0113           // get weight of hit
0114           const double eTotal = out_clust.getEnergy() * m_cfg.sampFrac;
0115           const float w       = m_weightFunc(hit.getEnergy(), eTotal, logWeightBase, 0);
0116
0117           // theta, phi
0118           Eigen::Vector2f pos2D(edm4hep::utils::anglePolar(hit.getPosition()),
0119                                 edm4hep::utils::angleAzimuthal(hit.getPosition()));
0120           // x, y, z
0121           Eigen::Vector3f pos3D(hit.getPosition().x, hit.getPosition().y, hit.getPosition().z);
0122
0123           const auto delta = out_clust.getPosition() - hit.getPosition();
0124           radius += delta * delta;
0125           dispersion += delta * delta * w;
0126
0127           // Weighted Sum x*x, x*y, x*z, y*y, etc.
0128           sum2_2D += w * pos2D * pos2D.transpose();
0129           sum2_3D += w * pos3D * pos3D.transpose();
0130
0131           // Weighted Sum x, y, z
0132           sum1_2D += w * pos2D;
0133           sum1_3D += w * pos3D;
0134
0135           w_sum += w;
0136         } // end hit loop
0137
0138         radius = sqrt((1. / (out_clust.getNhits() - 1.)) * radius);
0139         if (w_sum > 0) {
0140           dispersion = sqrt(dispersion / w_sum);
0141
0142           // normalize matrices
0143           sum2_2D /= w_sum;
0144           sum2_3D /= w_sum;
0145           sum1_2D /= w_sum;
0146           sum1_3D /= w_sum;
0147
0148           // 2D and 3D covariance matrices
0149           Eigen::Matrix2f cov2 = sum2_2D - sum1_2D * sum1_2D.transpose();
0150           Eigen::Matrix3f cov3 = sum2_3D - sum1_3D * sum1_3D.transpose();
0151
0152           // Solve for eigenvalues.  Corresponds to out_cluster's 2nd moments (widths)
0153           Eigen::EigenSolver<Eigen::Matrix2f> es_2D(
0154               cov2, false); // set to true for eigenvector calculation
0155           Eigen::EigenSolver<Eigen::Matrix3f> es_3D(
0156               cov3, true); // set to true for eigenvector calculation
0157
0158           // eigenvalues of symmetric real matrix are always real
0159           eigenValues_2D = es_2D.eigenvalues();
0160           eigenValues_3D = es_3D.eigenvalues();
0161           //find the eigenvector corresponding to the largest eigenvalue
0162           auto eigenvectors      = es_3D.eigenvectors();
0163           auto max_eigenvalue_it = std::ranges::max_element(
0164               eigenValues_3D, [](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