Redefined get_primary function to return first stable particle

src/algorithms/calorimetry/CalorimeterClusterRecoCoG.cc

@@ -21,12 +21,14 @@
 #include <podio/ObjectID.h>
 #include <podio/RelationRange.h>
 #include <podio/detail/Link.h>
+#include <cstdlib>
 #include <algorithm>
 #include <cctype>
 #include <cstddef>
 #include <limits>
 #include <map>
 #include <optional>
+#include <ranges>
 #include <tuple>
 #include <vector>
 
@@ -266,21 +268,45 @@ void CalorimeterClusterRecoCoG::associate(
 }
 
 edm4hep::MCParticle
-CalorimeterClusterRecoCoG::get_primary(const edm4hep::CaloHitContribution& contrib) {
-  // get contributing particle
-  const auto contributor = contrib.getParticle();
-
-  // walk back through parents to find primary
-  //   - TODO finalize primary selection. This
-  //     can be improved!!
-  edm4hep::MCParticle primary = contributor;
-  while (primary.parents_size() > 0) {
-    if (primary.getGeneratorStatus() != 0) {
+CalorimeterClusterRecoCoG::get_primary(const edm4hep::CaloHitContribution& contrib) const {
+  edm4hep::MCParticle current = contrib.getParticle();
+  if (!current.isAvailable()) {
+    return current;
+  }
+
+  const edm4hep::MCParticle original = current;
+  std::vector<edm4hep::MCParticle> chain;
+  chain.push_back(current);
+  while (current.getGeneratorStatus() == 0 && current.parents_size() > 0) {
+    const auto parent = current.getParents(0);
+
+    if (!parent.isAvailable()) {
       break;
     }
-    primary = primary.getParents(0);
+
+    current = parent;
+    chain.push_back(current);
   }
-  return primary;
+
+  const auto is_prompt_decay_particle = [this](const edm4hep::MCParticle& particle) {
+    return std::ranges::find(m_cfg.promptDecayPDGs, std::abs(particle.getPDG())) !=
+           m_cfg.promptDecayPDGs.end();
+  };
+
+  for (auto iterator = chain.rbegin(); iterator != chain.rend(); ++iterator) {
+    if (!iterator->isAvailable()) {
+      continue;
+    }
+    const bool isPrompt = is_prompt_decay_particle(*iterator);
+
+    if (isPrompt) {
+      continue;
+    }
+
+    return *iterator;
+  }
+
+  return original;
 }
 
 } // namespace eicrecon

src/algorithms/calorimetry/CalorimeterClusterRecoCoG.h

@@ -88,7 +88,7 @@ class CalorimeterClusterRecoCoG : public CalorimeterClusterRecoCoGAlgorithm,
                  const podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
                  edm4eic::MCRecoClusterParticleLinkCollection* links,
                  edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const;
-  static edm4hep::MCParticle get_primary(const edm4hep::CaloHitContribution& contrib);
+  edm4hep::MCParticle get_primary(const edm4hep::CaloHitContribution& contrib) const;
 };
 
 } // namespace eicrecon

src/algorithms/calorimetry/CalorimeterClusterRecoCoGConfig.h

@@ -27,6 +27,7 @@ struct CalorimeterClusterRecoCoGConfig {
   // the eta of the contributing hits. This is useful to avoid edge effects
   // for endcaps.
   bool enableEtaBounds = false;
+  std::vector<int> promptDecayPDGs{111, 221, 331, 310, 3122};
 };
 
 } // namespace eicrecon

src/algorithms/calorimetry/ImagingClusterReco.cc

@@ -18,6 +18,7 @@
 #include <podio/RelationRange.h>
 #include <podio/detail/Link.h>
 #include <podio/detail/LinkCollectionImpl.h>
+#include <cstdlib>
 #include <Eigen/Core>
 #include <Eigen/Householder> // IWYU pragma: keep
 #include <Eigen/Jacobi>
@@ -28,7 +29,9 @@
 #include <memory>
 #include <new>
 #include <optional>
+#include <ranges>
 #include <tuple>
+#include <vector>
 
 #include "algorithms/calorimetry/ClusterTypes.h"
 #include "algorithms/calorimetry/ImagingClusterReco.h"
@@ -329,20 +332,43 @@ void ImagingClusterReco::associate_mc_particles(
 
 edm4hep::MCParticle
 ImagingClusterReco::get_primary(const edm4hep::CaloHitContribution& contrib) const {
-  // get contributing particle
-  const auto contributor = contrib.getParticle();
-
-  // walk back through parents to find primary
-  //   - TODO finalize primary selection. This
-  //     can be improved!!
-  edm4hep::MCParticle primary = contributor;
-  while (primary.parents_size() > 0) {
-    if (primary.getGeneratorStatus() != 0) {
+  edm4hep::MCParticle current = contrib.getParticle();
+  if (!current.isAvailable()) {
+    return current;
+  }
+
+  const edm4hep::MCParticle original = current;
+  std::vector<edm4hep::MCParticle> chain;
+  chain.push_back(current);
+  while (current.getGeneratorStatus() == 0 && current.parents_size() > 0) {
+    const auto parent = current.getParents(0);
+
+    if (!parent.isAvailable()) {
       break;
     }
-    primary = primary.getParents(0);
+
+    current = parent;
+    chain.push_back(current);
+  }
+
+  const auto is_prompt_decay_particle = [this](const edm4hep::MCParticle& particle) {
+    return std::ranges::find(m_cfg.promptDecayPDGs, std::abs(particle.getPDG())) !=
+           m_cfg.promptDecayPDGs.end();
+  };
+
+  for (auto iterator = chain.rbegin(); iterator != chain.rend(); ++iterator) {
+    if (!iterator->isAvailable()) {
+      continue;
+    }
+    const bool isPrompt = is_prompt_decay_particle(*iterator);
+
+    if (isPrompt) {
+      continue;
+    }
+
+    return *iterator;
   }
-  return primary;
+  return original;
 }
 
 } // namespace eicrecon

src/algorithms/calorimetry/ImagingClusterRecoConfig.h

@@ -2,12 +2,14 @@
 // Copyright (C) 2022 Chao Peng, Sylvester Joosten, Whitney Armstrong
 
 #pragma once
+#include <vector>
 
 namespace eicrecon {
 
 struct ImagingClusterRecoConfig {
 
   int trackStopLayer = 9;
+  std::vector<int> promptDecayPDGs{111, 221, 331, 310, 3122};
 };
 
 } // namespace eicrecon

src/algorithms/calorimetry/SimCalorimeterHitProcessor.cc

@@ -17,9 +17,11 @@
 #include <podio/RelationRange.h>
 #include <cmath>
 #include <cstddef>
+#include <cstdlib>
 #include <functional>
 #include <gsl/pointers>
 #include <limits>
+#include <ranges>
 #include <stdexcept>
 #include <tuple>
 #include <unordered_map>
@@ -52,18 +54,39 @@ template <> struct hash<std::tuple<edm4hep::MCParticle, uint64_t, int>> {
 namespace {
 // Lookup primary MCParticle @TODO this should be a shared utility function in the edm4xxx
 // libraries
-edm4hep::MCParticle lookup_primary(const edm4hep::CaloHitContribution& contrib) {
-  const auto contributor = contrib.getParticle();
+edm4hep::MCParticle lookup_primary(const edm4hep::CaloHitContribution& contrib,
+                                   const std::vector<int>& promptDecayPDGs) {
+  edm4hep::MCParticle current = contrib.getParticle();
 
-  edm4hep::MCParticle primary = contributor;
-  while (primary.parents_size() > 0) {
-    if (primary.getGeneratorStatus() != 0) {
+  if (!current.isAvailable()) {
+    return current;
+  }
+
+  const edm4hep::MCParticle original = current;
+  std::vector<edm4hep::MCParticle> chain{current};
+  while (current.getGeneratorStatus() == 0 && current.parents_size() > 0) {
+    const auto parent = current.getParents(0);
+    if (!parent.isAvailable()) {
       break;
     }
-    primary = primary.getParents(0);
+    current = parent;
+    chain.push_back(current);
+  }
+  const auto is_prompt = [&promptDecayPDGs](const edm4hep::MCParticle& particle) {
+    return std::ranges::find(promptDecayPDGs, std::abs(particle.getPDG())) != promptDecayPDGs.end();
+  };
+  for (auto iterator = chain.rbegin(); iterator != chain.rend(); ++iterator) {
+    if (!iterator->isAvailable()) {
+      continue;
+    }
+    if (is_prompt(*iterator)) {
+      continue;
+    }
+    return *iterator;
   }
-  return primary;
+  return original;
 }
+
 class HitContributionAccumulator {
 private:
   float m_energy{0};
@@ -175,7 +198,7 @@ void SimCalorimeterHitProcessor::process(const SimCalorimeterHitProcessor::Input
         m_attenuationReferencePosition ? get_attenuation(ih.getPosition().z) : 1.;
     // Use primary particle (traced back through parents) to group contributions
     for (const auto& contrib : ih.getContributions()) {
-      edm4hep::MCParticle primary = lookup_primary(contrib);
+      edm4hep::MCParticle primary = lookup_primary(contrib, m_cfg.promptDecayPDGs);
       const double propagationTime =
           m_attenuationReferencePosition
               ? std::abs(m_attenuationReferencePosition.value() - ih.getPosition().z) *

src/algorithms/calorimetry/SimCalorimeterHitProcessorConfig.h

@@ -30,6 +30,7 @@ struct SimCalorimeterHitProcessorConfig {
   double fixedTimeDelay{};
   // time window for grouping contributions
   double timeWindow{100 * edm4eic::unit::ns};
+  std::vector<int> promptDecayPDGs{111, 221, 331, 310, 3122};
 };
 
 } // namespace eicrecon