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