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/SVD>
0025 #include <algorithm>
0026 #include <cmath>
0027 #include <gsl/pointers>
0028 #include <map>
0029 #include <new>
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   const auto [proto, mchitassociations] = input;
0039   auto [clusters, associations, layers] = output;
0040
0041   for (const auto& pcl : *proto) {
0042     if (!pcl.getHits().empty() && !pcl.getHits(0).isAvailable()) {
0043       warning("Protocluster hit relation is invalid, skipping protocluster");
0044       continue;
0045     }
0046     // get cluster and associated layers
0047     auto cl        = reconstruct_cluster(pcl);
0048     auto cl_layers = reconstruct_cluster_layers(pcl);
0049
0050     // Get cluster direction from the layer profile
0051     auto [theta, phi] = fit_track(cl_layers);
0052     cl.setIntrinsicTheta(theta);
0053     cl.setIntrinsicPhi(phi);
0054     // no error on the intrinsic direction TODO
0055
0056     // store layer and clusters on the datastore
0057     for (const auto& layer : cl_layers) {
0058       layers->push_back(layer);
0059       cl.addToClusters(layer);
0060     }
0061     clusters->push_back(cl);
0062
0063     // If sim hits are available, associate cluster with MCParticle
0064     if (mchitassociations->empty()) {
0065       debug("Provided MCRecoCalorimeterHitAssociation collection is empty. No truth associations "
0066             "will be performed.");
0067       continue;
0068     }
0069     associate_mc_particles(cl, mchitassociations, associations);
0070   }
0071
0072   // debug output
0073   for (const auto& cl : *clusters) {
0074     debug("Cluster {:d}: Edep = {:.3f} MeV, Dir = ({:.3f}, {:.3f}) deg", cl.getObjectID().index,
0075           cl.getEnergy() * 1000., cl.getIntrinsicTheta() / M_PI * 180.,
0076           cl.getIntrinsicPhi() / M_PI * 180.);
0077   }
0078 }
0079
0080 std::vector<edm4eic::MutableCluster>
0081 ImagingClusterReco::reconstruct_cluster_layers(const edm4eic::ProtoCluster& pcl) const {
0082   const auto& hits    = pcl.getHits();
0083   const auto& weights = pcl.getWeights();
0084   // using map to have hits sorted by layer
0085   std::map<int, std::vector<std::pair<const edm4eic::CalorimeterHit, float>>> layer_map;
0086   for (unsigned i = 0; i < hits.size(); ++i) {
0087     const auto hit = hits[i];
0088     auto lid       = hit.getLayer();
0089     //            if (layer_map.count(lid) == 0) {
0090     //                std::vector<std::pair<const edm4eic::CalorimeterHit, float>> v;
0091     //                layer_map[lid] = {};
0092     //            }
0093     layer_map[lid].emplace_back(hit, weights[i]);
0094   }
0095
0096   // create layers
0097   std::vector<edm4eic::MutableCluster> cl_layers;
0098   for (const auto& [lid, layer_hits] : layer_map) {
0099     auto layer = reconstruct_layer(layer_hits);
0100     cl_layers.push_back(layer);
0101   }
0102   return cl_layers;
0103 }
0104
0105 edm4eic::MutableCluster ImagingClusterReco::reconstruct_layer(
0106     const std::vector<std::pair<const edm4eic::CalorimeterHit, float>>& hits) const {
0107   edm4eic::MutableCluster layer;
0108   layer.setType(Jug::Reco::ClusterType::kClusterSlice);
0109   // Calculate averages
0110   double energy{0};
0111   double energyError{0};
0112   double time{0};
0113   double timeError{0};
0114   double sumOfWeights{0};
0115   auto pos = layer.getPosition();
0116   for (const auto& [hit, weight] : hits) {
0117     energy += hit.getEnergy() * weight;
0118     energyError += std::pow(hit.getEnergyError() * weight, 2);
0119     time += hit.getTime() * weight;
0120     timeError += std::pow(hit.getTimeError() * weight, 2);
0121     pos = pos + hit.getPosition() * weight;
0122     sumOfWeights += weight;
0123     layer.addToHits(hit);
0124   }
0125   layer.setEnergy(energy);
0126   layer.setEnergyError(std::sqrt(energyError));
0127   layer.setTime(time / sumOfWeights);
0128   layer.setTimeError(std::sqrt(timeError) / sumOfWeights);
0129   layer.setNhits(hits.size());
0130   layer.setPosition(pos / sumOfWeights);
0131   // positionError not set
0132   // Intrinsic direction meaningless in a cluster layer --> not set
0133
0134   // Shape parameters are calculated separately by CalorimeterClusterShape algorithm
0135
0136   return layer;
0137 }
0138
0139 edm4eic::MutableCluster
0140 ImagingClusterReco::reconstruct_cluster(const edm4eic::ProtoCluster& pcl) const {
0141   edm4eic::MutableCluster cluster;
0142
0143   const auto& hits    = pcl.getHits();
0144   const auto& weights = pcl.getWeights();
0145
0146   cluster.setType(Jug::Reco::ClusterType::kCluster3D);
0147   double energy      = 0.;
0148   double energyError = 0.;
0149   double time        = 0.;
0150   double timeError   = 0.;
0151   double meta        = 0.;
0152   double mphi        = 0.;
0153   double r           = 9999 * dd4hep::cm;
0154   for (unsigned i = 0; i < hits.size(); ++i) {
0155     const auto& hit    = hits[i];
0156     const auto& weight = weights[i];
0157     energy += hit.getEnergy() * weight;
0158     energyError += std::pow(hit.getEnergyError() * weight, 2);
0159     // energy weighting for the other variables
0160     const double energyWeight = hit.getEnergy() * weight;
0161     time += hit.getTime() * energyWeight;
0162     timeError += std::pow(hit.getTimeError() * energyWeight, 2);
0163     meta += edm4hep::utils::eta(hit.getPosition()) * energyWeight;
0164     mphi += edm4hep::utils::angleAzimuthal(hit.getPosition()) * energyWeight;
0165     r = std::min(edm4hep::utils::magnitude(hit.getPosition()), r);
0166     cluster.addToHits(hit);
0167   }
0168   cluster.setEnergy(energy);
0169   cluster.setEnergyError(std::sqrt(energyError));
0170   cluster.setTime(time / energy);
0171   cluster.setTimeError(std::sqrt(timeError) / energy);
0172   cluster.setNhits(hits.size());
0173   cluster.setPosition(edm4hep::utils::sphericalToVector(
0174       r, edm4hep::utils::etaToAngle(meta / energy), mphi / energy));
0175
0176   // Shape parameters are calculated separately by CalorimeterClusterShape algorithm
0177
0178   // Optionally store the MC truth associated with the first hit in this cluster
0179   // FIXME no connection between cluster and truth in edm4hep
0180   // if (mcHits) {
0181   //  const auto& mc_hit    = (*mcHits)[pcl.getHits(0).ID.value];
0182   //  cluster.mcID({mc_hit.truth().trackID, m_kMonteCarloSource});
0183   //}
0184
0185   return cluster;
0186 }
0187
0188 std::pair<double /* polar */, double /* azimuthal */>
0189 ImagingClusterReco::fit_track(const std::vector<edm4eic::MutableCluster>& layers) const {
0190   int nrows = 0;
0191   decltype(edm4eic::ClusterData::position) mean_pos{0, 0, 0};
0192   for (const auto& layer : layers) {
0193     if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0194       mean_pos = mean_pos + layer.getPosition();
0195       nrows += 1;
0196     }
0197   }
0198
0199   // cannot fit
0200   if (nrows < 2) {
0201     return {};
0202   }
0203
0204   mean_pos = mean_pos / nrows;
0205   // fill position data
0206   Eigen::MatrixXd pos(nrows, 3);
0207   int ir = 0;
0208   for (const auto& layer : layers) {
0209     if ((layer.getNhits() > 0) && (layer.getHits(0).getLayer() <= m_cfg.trackStopLayer)) {
0210       auto delta = layer.getPosition() - mean_pos;
0211       pos(ir, 0) = delta.x;
0212       pos(ir, 1) = delta.y;
0213       pos(ir, 2) = delta.z;
0214       ir += 1;
0215     }
0216   }
0217
0218   Eigen::JacobiSVD<Eigen::MatrixXd> svd(pos, Eigen::ComputeThinU | Eigen::ComputeThinV);
0219   const auto dir = svd.matrixV().col(0);
0220   // theta and phi
0221   return {std::acos(dir(2)), std::atan2(dir(1), dir(0))};
0222 }
0223
0224 void ImagingClusterReco::associate_mc_particles(
0225     const edm4eic::Cluster& cl,
0226     const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
0227     edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const {
0228   // --------------------------------------------------------------------------
0229   // Association Logic
0230   // --------------------------------------------------------------------------
0231   /*  1. identify all sim hits associated with a given protocluster, and sum
0232          *     the energy of the sim hits.
0233          *  2. for each sim hit
0234          *     - identify parents of each contributing particles; and
0235          *     - if parent is a primary particle, add to list of contributors
0236          *       and sum the energy contributed by the parent.
0237          *  3. create an association for each contributing primary with a weight
0238          *     of contributed energy over total sim hit energy.
0239          */
0240
0241   // lambda to compare MCParticles
0242   auto compare = [](const edm4hep::MCParticle& lhs, const edm4hep::MCParticle& rhs) {
0243     if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
0244       return (lhs.getObjectID().index < rhs.getObjectID().index);
0245     } else {
0246       return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
0247     }
0248   };
0249
0250   // bookkeeping maps for associated primaries
0251   std::map<edm4hep::MCParticle, double, decltype(compare)> mapMCParToContrib(compare);
0252
0253   // --------------------------------------------------------------------------
0254   // 1. get associated sim hits and sum energy
0255   // --------------------------------------------------------------------------
0256   double eSimHitSum = 0.;
0257   for (auto clhit : cl.getHits()) {
0258     // vector to hold associated sim hits
0259     std::vector<edm4hep::SimCalorimeterHit> vecAssocSimHits;
0260
0261     for (const auto& hitAssoc : *mchitassociations) {
0262       // if found corresponding raw hit, add sim hit to vector
0263       // and increment energy sum
0264       if (clhit.getRawHit() == hitAssoc.getRawHit()) {
0265         vecAssocSimHits.push_back(hitAssoc.getSimHit());
0266         eSimHitSum += vecAssocSimHits.back().getEnergy();
0267       }
0268     }
0269     debug("{} associated sim hits found for reco hit (cell ID = {})", vecAssocSimHits.size(),
0270           clhit.getCellID());
0271
0272     // ------------------------------------------------------------------------
0273     // 2. loop through associated sim hits
0274     // ------------------------------------------------------------------------
0275     for (const auto& simHit : vecAssocSimHits) {
0276       for (const auto& contrib : simHit.getContributions()) {
0277         // --------------------------------------------------------------------
0278         // grab primary responsible for contribution & increment relevant sum
0279         // --------------------------------------------------------------------
0280         edm4hep::MCParticle primary = get_primary(contrib);
0281         mapMCParToContrib[primary] += contrib.getEnergy();
0282
0283         trace("Identified primary: id = {}, pid = {}, total energy = {}, contributed = {}",
0284               primary.getObjectID().index, primary.getPDG(), primary.getEnergy(),
0285               mapMCParToContrib[primary]);
0286       }
0287     }
0288   }
0289   debug("Found {} primaries contributing a total of {} GeV", mapMCParToContrib.size(), eSimHitSum);
0290
0291   // --------------------------------------------------------------------------
0292   // 3. create association for each contributing primary
0293   // --------------------------------------------------------------------------
0294   for (auto [part, contribution] : mapMCParToContrib) {
0295     // calculate weight
0296     const double weight = contribution / eSimHitSum;
0297
0298     // set association
0299     auto assoc = assocs->create();
0300     assoc.setRecID(cl.getObjectID().index); // if not using collection, this is always set to -1
0301     assoc.setSimID(part.getObjectID().index);
0302     assoc.setWeight(weight);
0303     assoc.setRec(cl);
0304     assoc.setSim(part);
0305     debug("Associated cluster #{} to MC Particle #{} (pid = {}, status = {}, energy = {}) with "
0306           "weight ({})",
0307           cl.getObjectID().index, part.getObjectID().index, part.getPDG(),
0308           part.getGeneratorStatus(), part.getEnergy(), weight);
0309   }
0310 }
0311
0312 edm4hep::MCParticle
0313 ImagingClusterReco::get_primary(const edm4hep::CaloHitContribution& contrib) const {
0314   // get contributing particle
0315   const auto contributor = contrib.getParticle();
0316
0317   // walk back through parents to find primary
0318   //   - TODO finalize primary selection. This
0319   //     can be improved!!
0320   edm4hep::MCParticle primary = contributor;
0321   while (primary.parents_size() > 0) {
0322     if (primary.getGeneratorStatus() != 0) {
0323       break;
0324     }
0325     primary = primary.getParents(0);
0326   }
0327   return primary;
0328 }
0329
0330 } // namespace eicrecon