Back to home page

EIC code displayed by LXR

 
 

    


File indexing completed on 2025-06-08 07:53:21

0001 // SPDX-License-Identifier: LGPL-3.0-or-later
0002 // Copyright (C) 2025 Minho Kim, Sylvester Joosten, Derek Anderson, Wouter Deconinck
0003 
0004 #include "SimCalorimeterHitProcessor.h"
0005 
0006 #include <DD4hep/Detector.h>
0007 #include <DD4hep/IDDescriptor.h>
0008 #include <DD4hep/Readout.h>
0009 #include <DD4hep/config.h>
0010 #include <DDSegmentation/BitFieldCoder.h>
0011 #include <Evaluator/DD4hepUnits.h>
0012 #include <edm4eic/unit_system.h>
0013 #include <edm4hep/utils/vector_utils.h>
0014 #include <edm4hep/MCParticleCollection.h>
0015 #include <edm4hep/Vector3f.h>
0016 #include <fmt/core.h>
0017 #include <podio/ObjectID.h>
0018 #include <podio/RelationRange.h>
0019 #include <podio/podioVersion.h>
0020 #include <cmath>
0021 #include <cstddef>
0022 #include <gsl/pointers>
0023 #include <limits>
0024 #include <stdexcept>
0025 #include <tuple>
0026 #include <unordered_map>
0027 #include <variant>
0028 #include <vector>
0029 
0030 #include "algorithms/calorimetry/SimCalorimeterHitProcessorConfig.h"
0031 
0032 using namespace dd4hep;
0033 
0034 // Define necessary hash functions
0035 namespace std {
0036 
0037 #if defined(podio_VERSION_MAJOR) && defined(podio_VERSION_MINOR)
0038 #if podio_VERSION <= PODIO_VERSION(1, 2, 0)
0039 // Hash for podio::ObjectID
0040 template <> struct hash<podio::ObjectID> {
0041   size_t operator()(const podio::ObjectID& id) const noexcept {
0042     size_t h1 = std::hash<uint32_t>{}(id.collectionID);
0043     size_t h2 = std::hash<int>{}(id.index);
0044     return h1 ^ (h2 << 1);
0045   }
0046 };
0047 #endif // podio version check
0048 #endif // defined(podio_VERSION_MAJOR) && defined(podio_VERSION_MINOR)
0049 
0050 // Necessary to make MCParticle hashable
0051 // @TODO maybe this could be added to podio?
0052 template <> struct hash<edm4hep::MCParticle> {
0053   size_t operator()(const edm4hep::MCParticle& p) const noexcept {
0054     const auto& id = p.getObjectID();
0055     return std::hash<podio::ObjectID>()(id);
0056   }
0057 };
0058 // Hash for tuple<edm4hep::MCParticle, uint64_t>
0059 // --> not yet supported by any compiler at the moment
0060 template <> struct hash<std::tuple<edm4hep::MCParticle, uint64_t>> {
0061   size_t operator()(const std::tuple<edm4hep::MCParticle, uint64_t>& key) const noexcept {
0062     const auto& [particle, cellID] = key;
0063     size_t h1                      = hash<edm4hep::MCParticle>{}(particle);
0064     size_t h2                      = hash<uint64_t>{}(cellID);
0065     return h1 ^ (h2 << 1);
0066   }
0067 };
0068 
0069 } // namespace std
0070 
0071 // unnamed namespace for internal utility
0072 namespace {
0073 // Lookup primary MCParticle @TODO this should be a shared utiliy function in the edm4xxx
0074 // libraries
0075 edm4hep::MCParticle lookup_primary(const edm4hep::CaloHitContribution& contrib) {
0076   const auto contributor = contrib.getParticle();
0077 
0078   edm4hep::MCParticle primary = contributor;
0079   while (primary.parents_size() > 0) {
0080     if (primary.getGeneratorStatus() != 0)
0081       break;
0082     primary = primary.getParents(0);
0083   }
0084   return primary;
0085 }
0086 class HitContributionAccumulator {
0087 private:
0088   float m_energy{0};
0089   float m_avg_time{0};
0090   float m_min_time{std::numeric_limits<float>::max()};
0091   edm4hep::Vector3f m_avg_position{0, 0, 0};
0092 
0093 public:
0094   void add(const float energy, const float time, const edm4hep::Vector3f& pos) {
0095     m_energy += energy;
0096     m_avg_time += energy * time;
0097     m_avg_position = m_avg_position + energy * pos;
0098     m_min_time     = (time < m_min_time) ? time : m_min_time;
0099   }
0100   float getEnergy() const { return m_energy; }
0101   float getAvgTime() const { return m_energy > 0 ? m_avg_time / m_energy : 0; }
0102   float getMinTime() const { return m_min_time; }
0103   edm4hep::Vector3f getAvgPosition() const {
0104     return m_energy > 0 ? m_avg_position / m_energy : edm4hep::Vector3f{0, 0, 0};
0105   }
0106 };
0107 
0108 } // namespace
0109 
0110 namespace eicrecon {
0111 
0112 void SimCalorimeterHitProcessor::init() {
0113 
0114   // readout checks
0115   if (m_cfg.readout.empty()) {
0116     error("readoutClass is not provided, it is needed to know the fields in readout ids");
0117     throw std::runtime_error("readoutClass is not provided");
0118   }
0119 
0120   // get decoders
0121   try {
0122     m_id_spec = m_geo.detector()->readout(m_cfg.readout).idSpec();
0123   } catch (...) {
0124     debug("Failed to load ID decoder for {}", m_cfg.readout);
0125     throw std::runtime_error(fmt::format("Failed to load ID decoder for {}", m_cfg.readout));
0126   }
0127 
0128   // get m_hit_id_mask for adding up hits with the same dimensions that are merged over
0129   {
0130     uint64_t id_inverse_mask = 0;
0131     if (!m_cfg.hitMergeFields.empty()) {
0132       for (auto& field : m_cfg.hitMergeFields) {
0133         id_inverse_mask |= m_id_spec.field(field)->mask();
0134       }
0135     }
0136     m_hit_id_mask = ~id_inverse_mask;
0137     debug("ID mask in {:s}: {:#064b}", m_cfg.readout, m_hit_id_mask.value());
0138   }
0139 
0140   // get m_contribution_id_mask for adding up contributions with the same dimensions that are merged over
0141   {
0142     uint64_t id_inverse_mask = 0;
0143     if (!m_cfg.contributionMergeFields.empty()) {
0144       for (auto& field : m_cfg.contributionMergeFields) {
0145         id_inverse_mask |= m_id_spec.field(field)->mask();
0146       }
0147       m_contribution_id_mask = ~id_inverse_mask;
0148     }
0149     debug("ID mask in {:s}: {:#064b}", m_cfg.readout, m_contribution_id_mask.value());
0150   }
0151 
0152   // get reference position for attenuating hits and contributions
0153   if (!m_cfg.attenuationReferencePositionName.empty()) {
0154     m_attenuationReferencePosition =
0155         m_geo.detector()->constant<double>(m_cfg.attenuationReferencePositionName) *
0156         edm4eic::unit::mm / dd4hep::mm;
0157   }
0158 }
0159 
0160 // Group contributions by (primary particle, cell ID), apply optional attenuation, and optionally merge into superhits
0161 void SimCalorimeterHitProcessor::process(const SimCalorimeterHitProcessor::Input& input,
0162                                          const SimCalorimeterHitProcessor::Output& output) const {
0163 
0164   const auto [in_hits]              = input;
0165   auto [out_hits, out_hit_contribs] = output;
0166 
0167   // Map for staging output information. We have 2 levels of structure:
0168   //   - top level: (MCParticle, Merged Hit CellID)
0169   //   - second level: (Merged Contributions)
0170   // Ideally we would want immediately create our output objects and modify the
0171   // contributions when needed. That could reduce the following code to a single loop
0172   // (instead of 2 consecutive loops). However, this is not possible as we may have to merge
0173   // (hence modify) contributions which is not supported for PodIO VectorMembers. Using
0174   // reasonable contribution merging, at least the intermediary structure should be
0175   // quite a bit smaller than the original hit collection.
0176   using HitIndex = std::tuple<edm4hep::MCParticle, uint64_t /* cellID */>;
0177   std::unordered_map<HitIndex,
0178                      std::unordered_map<uint64_t /* cellID */, HitContributionAccumulator>>
0179       hit_map;
0180 
0181   for (const auto& ih : *in_hits) {
0182     // the cell ID of the new superhit we are making
0183     const uint64_t newhit_cellID =
0184         (ih.getCellID() & m_hit_id_mask.value() & m_contribution_id_mask.value());
0185     // the cell ID of this particular contribution (we are using contributions to store
0186     // the hits making up this "superhit" with more segmentation)
0187     const uint64_t newcontrib_cellID = (ih.getCellID() & m_hit_id_mask.value());
0188     // Optional attenuation
0189     const double attFactor =
0190         m_attenuationReferencePosition ? get_attenuation(ih.getPosition().z) : 1.;
0191     // Use primary particle (traced back through parents) to group contributions
0192     for (const auto& contrib : ih.getContributions()) {
0193       edm4hep::MCParticle primary = lookup_primary(contrib);
0194       auto& hit_accum             = hit_map[{primary, newhit_cellID}][newcontrib_cellID];
0195       hit_accum.add(contrib.getEnergy() * attFactor, contrib.getTime(), ih.getPosition());
0196     }
0197   }
0198 
0199   // We now have our data structured as we want it, next we need to visit all hits again
0200   // and create our output structures
0201   for (const auto& [hit_idx, contribs] : hit_map) {
0202 
0203     auto out_hit = out_hits->create();
0204 
0205     const auto& [particle, cellID] = hit_idx;
0206     HitContributionAccumulator new_hit;
0207     for (const auto& [contrib_idx, contrib] : contribs) {
0208       // Aggregate contributions to for the global hit
0209       new_hit.add(contrib.getEnergy(), contrib.getMinTime(), contrib.getAvgPosition());
0210       // Now store the contribution itself
0211       auto out_hit_contrib = out_hit_contribs->create();
0212       out_hit_contrib.setPDG(particle.getPDG());
0213       out_hit_contrib.setEnergy(contrib.getEnergy());
0214       out_hit_contrib.setTime(contrib.getMinTime());
0215       out_hit_contrib.setStepPosition(contrib.getAvgPosition());
0216       out_hit_contrib.setParticle(particle);
0217       out_hit.addToContributions(out_hit_contrib);
0218     }
0219     out_hit.setCellID(cellID);
0220     out_hit.setEnergy(new_hit.getEnergy());
0221     out_hit.setPosition(new_hit.getAvgPosition());
0222   }
0223 }
0224 
0225 double SimCalorimeterHitProcessor::get_attenuation(double zpos) const {
0226   double length = std::abs(m_attenuationReferencePosition.value() - zpos);
0227   double factor =
0228       m_cfg.attenuationParameters[0] * std::exp(-length / m_cfg.attenuationParameters[1]) +
0229       (1 - m_cfg.attenuationParameters[0]) * std::exp(-length / m_cfg.attenuationParameters[2]);
0230   return factor;
0231 }
0232 } // namespace eicrecon