algorithms/calorimetry/ImagingClusterReco.cc

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2022 - 2024, Sylvester Joosten, Chao Peng, Wouter Deconinck, David Lawrence, Derek Anderson
0003
0004 /*
0005  *  Reconstruct the cluster/layer info for imaging calorimeter
0006  *  Logarithmic weighting is used to describe energy deposit in transverse direction
0007  *
0008  *  Author: Chao Peng (ANL), 06/02/2021
0009  */
0010
0011 #include "algorithms/calorimetry/ImagingClusterReco.h"
0012
0013 #include <Evaluator/DD4hepUnits.h>
0014 #include <edm4hep/RawCalorimeterHit.h>
0015 #include <edm4hep/SimCalorimeterHit.h>
0016 #include <edm4hep/Vector3f.h>
0017 #include <edm4hep/utils/vector_utils.h>
0018 #include <fmt/core.h>
0019 #include <podio/ObjectID.h>
0020 #include <podio/RelationRange.h>
0021 #include <Eigen/Core>
0022 #include <Eigen/Householder> // IWYU pragma: keep
0023 #include <Eigen/Jacobi>
0024 #include <Eigen/QR>
0025 #include <Eigen/SVD>
0026 #include <algorithm>
0027 #include <cmath>
0028 #include <gsl/pointers>
0029 #include <map>
0030
0031 #include "algorithms/calorimetry/ClusterTypes.h"
0032 #include "algorithms/calorimetry/ImagingClusterRecoConfig.h"
0033
0034 namespace eicrecon {
0035
0036 void ImagingClusterReco::process(const Input& input, const Output& output) const {
0037
0038 #if EDM4EIC_VERSION_MAJOR >= 7
0039   const auto [proto, mchitassociations] = input;
0040 #else
0041   const auto [proto, mchits] = input;
0042 #endif
0043   auto [clusters, associations, layers] = output;
0044
0045   for (const auto& pcl : *proto) {
0046     if (!pcl.getHits().empty() && !pcl.getHits(0).isAvailable()) {
0047       warning("Protocluster hit relation is invalid, skipping protocluster");
0048       continue;
0049     }
0050     // get cluster and associated layers
0051     auto cl        = reconstruct_cluster(pcl);
0052     auto cl_layers = reconstruct_cluster_layers(pcl);
0053
0054     // Get cluster direction from the layer profile
0055     auto [theta, phi] = fit_track(cl_layers);
0056     cl.setIntrinsicTheta(theta);
0057     cl.setIntrinsicPhi(phi);
0058     // no error on the intrinsic direction TODO
0059
0060     // store layer and clusters on the datastore
0061     for (const auto& layer : cl_layers) {
0062       layers->push_back(layer);
0063       cl.addToClusters(layer);
0064     }
0065     clusters->push_back(cl);
0066
0067     // If sim hits are available, associate cluster with MCParticle
0068 #if EDM4EIC_VERSION_MAJOR >= 7
0069     if (mchitassociations->size() == 0) {
0070       debug("Provided MCRecoCalorimeterHitAssociation collection is empty. No truth associations "
0071             "will be performed.");
0072       continue;
0073     } else {
0074       associate_mc_particles(cl, mchitassociations, associations);
0075     }
0076 #else
0077     if (mchits->size() == 0) {
0078       debug("Provided SimCalorimeterHitCollection is empty. No truth association will be "
0079             "performed.");
0080       continue;
0081     } else {
0082       associate_mc_particles(cl, mchits, associations);
0083     }
0084 #endif
0085   }
0086
0087   // debug output
0088   for (const auto& cl : *clusters) {
0089     debug("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.getObjectID().index,
0090           cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0091           cl.getIntrinsicPhi() / M_PI * 180.);
0092   }
0093 }
0094
0095 std::vector<edm4eic::MutableCluster>
0096 ImagingClusterReco::reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) const {
0097   const auto& hits    = pcl.getHits();
0098   const auto& weights = pcl.getWeights();
0099   // using map to have hits sorted by layer
0100   std::map<int, std::vector<std::pair<const edm4eic::CalorimeterHit, float>>> layer_map;
0101   for (unsigned i = 0; i < hits.size(); ++i) {
0102     const auto hit = hits[i];
0103     auto lid       = hit.getLayer();
0104     //            if (layer_map.count(lid) == 0) {
0105     //                std::vector<std::pair<const edm4eic::CalorimeterHit, float>> v;
0106     //                layer_map[lid] = {};
0107     //            }
0108     layer_map[lid].push_back({hit, weights[i]});
0109   }
0110
0111   // create layers
0112   std::vector<edm4eic::MutableCluster> cl_layers;
0113   for (const auto& [lid, layer_hits] : layer_map) {
0114     auto layer = reconstruct_layer(layer_hits);
0115     cl_layers.push_back(layer);
0116   }
0117   return cl_layers;
0118 }
0119
0120 edm4eic::MutableCluster ImagingClusterReco::reconstruct_layer(
0121     const std::vector<std::pair<const edm4eic::CalorimeterHit, float>>& hits) const {
0122   edm4eic::MutableCluster layer;
0123   layer.setType(Jug::Reco::ClusterType::kClusterSlice);
0124   // Calculate averages
0125   double energy{0};
0126   double energyError{0};
0127   double time{0};
0128   double timeError{0};
0129   double sumOfWeights{0};
0130   auto pos = layer.getPosition();
0131   for (const auto& [hit, weight] : hits) {
0132     energy += hit.getEnergy() * weight;
0133     energyError += std::pow(hit.getEnergyError() * weight, 2);
0134     time += hit.getTime() * weight;
0135     timeError += std::pow(hit.getTimeError() * weight, 2);
0136     pos = pos + hit.getPosition() * weight;
0137     sumOfWeights += weight;
0138     layer.addToHits(hit);
0139   }
0140   layer.setEnergy(energy);
0141   layer.setEnergyError(std::sqrt(energyError));
0142   layer.setTime(time / sumOfWeights);
0143   layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0144   layer.setNhits(hits.size());
0145   layer.setPosition(pos / sumOfWeights);
0146   // positionError not set
0147   // Intrinsic direction meaningless in a cluster layer --> not set
0148
0149   // Calculate radius as the standard deviation of the hits versus the cluster center
0150   double radius = 0.;
0151   for (const auto& [hit, weight] : hits) {
0152     radius += std::pow(edm4hep::utils::magnitude(hit.getPosition() - layer.getPosition()), 2);
0153   }
0154   layer.addToShapeParameters(std::sqrt(radius / layer.getNhits()));
0155   // TODO Skewedness
0156
0157   return layer;
0158 }
0159
0160 edm4eic::MutableCluster
0161 ImagingClusterReco::reconstruct_cluster(const edm4eic::ProtoCluster& pcl) const {
0162   edm4eic::MutableCluster cluster;
0163
0164   const auto& hits    = pcl.getHits();
0165   const auto& weights = pcl.getWeights();
0166
0167   cluster.setType(Jug::Reco::ClusterType::kCluster3D);
0168   double energy      = 0.;
0169   double energyError = 0.;
0170   double time        = 0.;
0171   double timeError   = 0.;
0172   double meta        = 0.;
0173   double mphi        = 0.;
0174   double r           = 9999 * dd4hep::cm;
0175   for (unsigned i = 0; i < hits.size(); ++i) {
0176     const auto& hit    = hits[i];
0177     const auto& weight = weights[i];
0178     energy += hit.getEnergy() * weight;
0179     energyError += std::pow(hit.getEnergyError() * weight, 2);
0180     // energy weighting for the other variables
0181     const double energyWeight = hit.getEnergy() * weight;
0182     time += hit.getTime() * energyWeight;
0183     timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0184     meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0185     mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0186     r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0187     cluster.addToHits(hit);
0188   }
0189   cluster.setEnergy(energy);
0190   cluster.setEnergyError(std::sqrt(energyError));
0191   cluster.setTime(time / energy);
0192   cluster.setTimeError(std::sqrt(timeError) / energy);
0193   cluster.setNhits(hits.size());
0194   cluster.setPosition(edm4hep::utils::sphericalToVector(
0195       r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0196
0197   // shower radius estimate (eta-phi plane)
0198   double radius = 0.;
0199   for (const auto& hit : hits) {
0200     radius += std::pow(std::hypot((edm4hep::utils::eta(hit.getPosition()) -
0201                                    edm4hep::utils::eta(cluster.getPosition())),
0202                                   (edm4hep::utils::angleAzimuthal(hit.getPosition()) -
0203                                    edm4hep::utils::angleAzimuthal(cluster.getPosition()))),
0204                        2.0);
0205   }
0206   cluster.addToShapeParameters(std::sqrt(radius / cluster.getNhits()));
0207   // Skewedness not calculated TODO
0208
0209   // Optionally store the MC truth associated with the first hit in this cluster
0210   // FIXME no connection between cluster and truth in edm4hep
0211   // if (mcHits) {
0212   //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0213   //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0214   //}
0215
0216   return cluster;
0217 }
0218
0219 std::pair<double /* polar */, double /* azimuthal */>
0220 ImagingClusterReco::fit_track(const std::vector<edm4eic::MutableCluster>& layers) const {
0221   int nrows = 0;
0222   decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0223   for (const auto& layer : layers) {
0224     if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0225       mean_pos = mean_pos + layer.getPosition();
0226       nrows += 1;
0227     }
0228   }
0229
0230   // cannot fit
0231   if (nrows < 2) {
0232     return {};
0233   }
0234
0235   mean_pos = mean_pos / nrows;
0236   // fill position data
0237   Eigen::MatrixXd pos(nrows, 3);
0238   int ir = 0;
0239   for (const auto& layer : layers) {
0240     if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0241       auto delta = layer.getPosition() - mean_pos;
0242       pos(ir, 0) = delta.x;
0243       pos(ir, 1) = delta.y;
0244       pos(ir, 2) = delta.z;
0245       ir += 1;
0246     }
0247   }
0248
0249   Eigen::JacobiSVD<Eigen::MatrixXd> svd(pos, Eigen::ComputeThinU | Eigen::ComputeThinV);
0250   const auto dir = svd.matrixV().col(0);
0251   // theta and phi
0252   return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0253 }
0254
0255 void ImagingClusterReco::associate_mc_particles(
0256     const edm4eic::Cluster& cl,
0257 #if EDM4EIC_VERSION_MAJOR >= 7
0258     const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
0259 #else
0260     const edm4hep::SimCalorimeterHitCollection* mchits,
0261 #endif
0262     edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const {
0263   // --------------------------------------------------------------------------
0264   // Association Logic
0265   // --------------------------------------------------------------------------
0266   /*  1. identify all sim hits associated with a given protocluster, and sum
0267          *     the energy of the sim hits.
0268          *  2. for each sim hit
0269          *     - identify parents of each contributing particles; and
0270          *     - if parent is a primary particle, add to list of contributors
0271          *       and sum the energy contributed by the parent.
0272          *  3. create an association for each contributing primary with a weight
0273          *     of contributed energy over total sim hit energy.
0274          */
0275
0276   // lambda to compare MCParticles
0277   auto compare = [](const edm4hep::MCParticle& lhs, const edm4hep::MCParticle& rhs) {
0278     if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
0279       return (lhs.getObjectID().index < rhs.getObjectID().index);
0280     } else {
0281       return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
0282     }
0283   };
0284
0285   // bookkeeping maps for associated primaries
0286   std::map<edm4hep::MCParticle, double, decltype(compare)> mapMCParToContrib(compare);
0287
0288   // --------------------------------------------------------------------------
0289   // 1. get associated sim hits and sum energy
0290   // --------------------------------------------------------------------------
0291   double eSimHitSum = 0.;
0292   for (auto clhit : cl.getHits()) {
0293     // vector to hold associated sim hits
0294     std::vector<edm4hep::SimCalorimeterHit> vecAssocSimHits;
0295
0296 #if EDM4EIC_VERSION_MAJOR >= 7
0297     for (const auto& hitAssoc : *mchitassociations) {
0298       // if found corresponding raw hit, add sim hit to vector
0299       // and increment energy sum
0300       if (clhit.getRawHit() == hitAssoc.getRawHit()) {
0301         vecAssocSimHits.push_back(hitAssoc.getSimHit());
0302         eSimHitSum += vecAssocSimHits.back().getEnergy();
0303       }
0304     }
0305 #else
0306     for (const auto& mchit : *mchits) {
0307       if (mchit.getCellID() == clhit.getCellID()) {
0308         vecAssocSimHits.push_back(mchit);
0309         eSimHitSum += vecAssocSimHits.back().getEnergy();
0310         break;
0311       }
0312     }
0313
0314     // if no matching cell ID found, continue
0315     // otherwise increment sum
0316     if (vecAssocSimHits.empty()) {
0317       debug("No matching SimHit for hit {}", clhit.getCellID());
0318       continue;
0319     }
0320 #endif
0321     debug("{} associated sim hits found for reco hit (cell ID = {})", vecAssocSimHits.size(),
0322           clhit.getCellID());
0323
0324     // ------------------------------------------------------------------------
0325     // 2. loop through associated sim hits
0326     // ------------------------------------------------------------------------
0327     for (const auto& simHit : vecAssocSimHits) {
0328       for (const auto& contrib : simHit.getContributions()) {
0329         // --------------------------------------------------------------------
0330         // grab primary responsible for contribution & increment relevant sum
0331         // --------------------------------------------------------------------
0332         edm4hep::MCParticle primary = get_primary(contrib);
0333         mapMCParToContrib[primary] += contrib.getEnergy();
0334
0335         trace("Identified primary: id = {}, pid = {}, total energy = {}, contributed = {}",
0336               primary.getObjectID().index, primary.getPDG(), primary.getEnergy(),
0337               mapMCParToContrib[primary]);
0338       }
0339     }
0340   }
0341   debug("Found {} primaries contributing a total of {} GeV", mapMCParToContrib.size(), eSimHitSum);
0342
0343   // --------------------------------------------------------------------------
0344   // 3. create association for each contributing primary
0345   // --------------------------------------------------------------------------
0346   for (auto [part, contribution] : mapMCParToContrib) {
0347     // calculate weight
0348     const double weight = contribution / eSimHitSum;
0349
0350     // set association
0351     auto assoc = assocs->create();
0352     assoc.setRecID(cl.getObjectID().index); // if not using collection, this is always set to -1
0353     assoc.setSimID(part.getObjectID().index);
0354     assoc.setWeight(weight);
0355     assoc.setRec(cl);
0356     assoc.setSim(part);
0357     debug("Associated cluster #{} to MC Particle #{} (pid = {}, status = {}, energy = {}) with "
0358           "weight ({})",
0359           cl.getObjectID().index, part.getObjectID().index, part.getPDG(),
0360           part.getGeneratorStatus(), part.getEnergy(), weight);
0361   }
0362 }
0363
0364 edm4hep::MCParticle
0365 ImagingClusterReco::get_primary(const edm4hep::CaloHitContribution& contrib) const {
0366   // get contributing particle
0367   const auto contributor = contrib.getParticle();
0368
0369   // walk back through parents to find primary
0370   //   - TODO finalize primary selection. This
0371   //     can be improved!!
0372   edm4hep::MCParticle primary = contributor;
0373   while (primary.parents_size() > 0) {
0374     if (primary.getGeneratorStatus() != 0)
0375       break;
0376     primary = primary.getParents(0);
0377   }
0378   return primary;
0379 }
0380
0381 } // namespace eicrecon