refactor: unify LinkNavigator truth-link handling

src/algorithms/calorimetry/CalorimeterClusterRecoCoG.cc

@@ -13,14 +13,16 @@
 #include <boost/range/adaptor/map.hpp>
 #include <edm4eic/CalorimeterHitCollection.h>
 #include <edm4eic/Cov3f.h>
+#include <edm4hep/CaloHitContribution.h>
+#include <edm4hep/MCParticle.h>
 #include <edm4hep/RawCalorimeterHit.h>
 #include <edm4hep/SimCalorimeterHitCollection.h>
 #include <edm4hep/Vector3f.h>
 #include <edm4hep/utils/vector_utils.h>
+#include <gsl/pointers>
 #include <podio/LinkNavigator.h>
 #include <podio/ObjectID.h>
 #include <podio/RelationRange.h>
-#include <podio/detail/Link.h>
 #include <algorithm>
 #include <cctype>
 #include <cstddef>
@@ -32,6 +34,8 @@
 
 #include "CalorimeterClusterRecoCoG.h"
 #include "algorithms/calorimetry/CalorimeterClusterRecoCoGConfig.h"
+#include "algorithms/interfaces/CompareObjectID.h"
+#include "algorithms/interfaces/LinkTruthUtils.h"
 
 namespace eicrecon {
 
@@ -57,16 +61,12 @@ void CalorimeterClusterRecoCoG::process(const CalorimeterClusterRecoCoG::Input&
   auto [clusters, links, associations]              = output;
 
   // Check if truth associations are possible
-  const bool do_assoc = mchitlinks != nullptr && !mchitlinks->empty();
+  const truth::EventLinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection> link_nav(mchitlinks);
+  const bool do_assoc = link_nav.enabled();
   if (!do_assoc) {
     debug("Provided MCRecoCalorimeterHitLink collection is empty. No truth associations "
           "will be performed.");
   }
-  // Build fast lookup once per event using podio::LinkNavigator
-  std::optional<podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>> link_nav;
-  if (do_assoc) {
-    link_nav.emplace(*mchitlinks);
-  }
 
   for (const auto& pcl : *proto) {
     // skip protoclusters with no hits
@@ -87,7 +87,7 @@ void CalorimeterClusterRecoCoG::process(const CalorimeterClusterRecoCoG::Input&
 
     // If sim hits are available, associate cluster with MCParticle
     if (do_assoc) {
-      associate(cl, mchitassociations, *link_nav, links, associations);
+      associate(cl, mchitassociations, link_nav, links, associations);
     }
   }
 }
@@ -177,7 +177,7 @@ CalorimeterClusterRecoCoG::reconstruct(const edm4eic::ProtoCluster& pcl) const {
 void CalorimeterClusterRecoCoG::associate(
     const edm4eic::Cluster& cl,
     [[maybe_unused]] const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
-    const podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
+    const truth::EventLinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
     edm4eic::MCRecoClusterParticleLinkCollection* links,
     edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const {
   // --------------------------------------------------------------------------
@@ -193,16 +193,8 @@ void CalorimeterClusterRecoCoG::associate(
    *     of contributed energy over total sim hit energy.
    */
 
-  // lambda to compare MCParticles
-  auto compare = [](const edm4hep::MCParticle& lhs, const edm4hep::MCParticle& rhs) {
-    if (lhs.getObjectID().collectionID == rhs.getObjectID().collectionID) {
-      return (lhs.getObjectID().index < rhs.getObjectID().index);
-    }
-    return (lhs.getObjectID().collectionID < rhs.getObjectID().collectionID);
-  };
-
   // bookkeeping maps for associated primaries
-  std::map<edm4hep::MCParticle, double, decltype(compare)> mapMCParToContrib(compare);
+  std::map<edm4hep::MCParticle, double, CompareObjectID<edm4hep::MCParticle>> mapMCParToContrib;
 
   // --------------------------------------------------------------------------
   // 1. get associated sim hits and sum energy
@@ -211,7 +203,7 @@ void CalorimeterClusterRecoCoG::associate(
   for (auto clhit : cl.getHits()) {
 
     // Get linked sim hits using LinkNavigator
-    const auto vecAssocSimHits = link_nav.getLinked(clhit.getRawHit());
+    const auto vecAssocSimHits = link_nav.linked(clhit.getRawHit());
 
     for (const auto& [simHit, weight] : vecAssocSimHits) {
       eSimHitSum += simHit.getEnergy();
@@ -228,7 +220,7 @@ void CalorimeterClusterRecoCoG::associate(
         // --------------------------------------------------------------------
         // grab primary responsible for contribution & increment relevant sum
         // --------------------------------------------------------------------
-        edm4hep::MCParticle primary = get_primary(contrib);
+        edm4hep::MCParticle primary = truth::primaryFrom(contrib);
         mapMCParToContrib[primary] += contrib.getEnergy();
 
         trace("Identified primary: id = {}, pid = {}, total energy = {}, contributed = {}",
@@ -246,17 +238,10 @@ void CalorimeterClusterRecoCoG::associate(
     // calculate weight
     const double weight = contribution / eSimHitSum;
 
-    // create link
-    auto link = links->create();
-    link.setWeight(weight);
-    link.setFrom(cl);
-    link.setTo(part);
-
-    // set association
-    auto assoc = assocs->create();
-    assoc.setWeight(weight);
-    assoc.setRec(cl);
-    assoc.setSim(part);
+    truth::addWeightedRelation(
+        cl, part, static_cast<float>(weight),
+        gsl::not_null<edm4eic::MCRecoClusterParticleLinkCollection*>{links},
+        gsl::not_null<edm4eic::MCRecoClusterParticleAssociationCollection*>{assocs});
 
     debug("Associated cluster #{} to MC Particle #{} (pid = {}, status = {}, energy = {}) with "
           "weight ({})",
@@ -265,22 +250,4 @@ 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) {
-      break;
-    }
-    primary = primary.getParents(0);
-  }
-  return primary;
-}
-
 } // namespace eicrecon

src/algorithms/calorimetry/CalorimeterClusterRecoCoG.h

@@ -17,9 +17,6 @@
 #include <edm4eic/MCRecoClusterParticleAssociationCollection.h>
 #include <edm4eic/MCRecoClusterParticleLinkCollection.h>
 #include <edm4eic/ProtoClusterCollection.h>
-#include <edm4hep/CaloHitContribution.h>
-#include <edm4hep/MCParticle.h>
-#include <podio/LinkNavigator.h>
 #include <algorithm>
 #include <cmath>
 #include <functional>
@@ -31,6 +28,7 @@
 #include <utility>
 
 #include "CalorimeterClusterRecoCoGConfig.h"
+#include "algorithms/interfaces/LinkTruthUtils.h"
 #include "algorithms/interfaces/WithPodConfig.h"
 
 static double constWeight(double /*E*/, double /*tE*/, double /*p*/, int /*type*/) { return 1.0; }
@@ -83,12 +81,12 @@ class CalorimeterClusterRecoCoG : public CalorimeterClusterRecoCoGAlgorithm,
 
 private:
   std::optional<edm4eic::MutableCluster> reconstruct(const edm4eic::ProtoCluster& pcl) const;
-  void associate(const edm4eic::Cluster& cl,
-                 const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
-                 const podio::LinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
-                 edm4eic::MCRecoClusterParticleLinkCollection* links,
-                 edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const;
-  static edm4hep::MCParticle get_primary(const edm4hep::CaloHitContribution& contrib);
+  void
+  associate(const edm4eic::Cluster& cl,
+            const edm4eic::MCRecoCalorimeterHitAssociationCollection* mchitassociations,
+            const truth::EventLinkNavigator<edm4eic::MCRecoCalorimeterHitLinkCollection>& link_nav,
+            edm4eic::MCRecoClusterParticleLinkCollection* links,
+            edm4eic::MCRecoClusterParticleAssociationCollection* assocs) const;
 };
 
 } // namespace eicrecon

src/algorithms/calorimetry/CalorimeterClusterShape.cc

@@ -10,9 +10,8 @@
 #include <edm4hep/MCParticle.h>
 #include <edm4hep/Vector3f.h>
 #include <edm4hep/utils/vector_utils.h>
-#include <podio/RelationRange.h>
-#include <podio/detail/Link.h>
-#include <podio/detail/LinkCollectionImpl.h>
+#include <gsl/pointers>
+#include <podio/LinkNavigator.h>
 #include <Eigen/Core>
 #include <Eigen/Eigenvalues>
 #include <Eigen/Householder> // IWYU pragma: keep
@@ -27,6 +26,7 @@
 #include <vector>
 
 #include "algorithms/calorimetry/CalorimeterClusterShapeConfig.h"
+#include "algorithms/interfaces/LinkTruthUtils.h"
 
 namespace eicrecon {
 
@@ -62,9 +62,11 @@ void CalorimeterClusterShape::process(const CalorimeterClusterShape::Input& inpu
                                       const CalorimeterClusterShape::Output& output) const {
 
   // grab inputs/outputs
-  const auto [in_clusters, in_associations]        = input;
+  const auto [in_clusters, in_links]               = input;
   auto [out_clusters, out_links, out_associations] = output;
 
+  const truth::EventLinkNavigator<edm4eic::MCRecoClusterParticleLinkCollection> link_nav(in_links);
+
   // exit if no clusters in collection
   if (in_clusters->empty()) {
     debug("No clusters in input collection.");
@@ -226,21 +228,16 @@ void CalorimeterClusterShape::process(const CalorimeterClusterShape::Input& inpu
     out_clusters->push_back(out_clust);
 
     // ----------------------------------------------------------------------
-    // if provided, copy associations
+    // if provided, copy links and associations
     // ----------------------------------------------------------------------
-    for (auto in_assoc : *in_associations) {
-      if (in_assoc.getRec() == in_clust) {
-        auto mc_par   = in_assoc.getSim();
-        auto out_link = out_links->create();
-        out_link.setFrom(out_clust);
-        out_link.setTo(mc_par);
-        out_link.setWeight(in_assoc.getWeight());
-        auto out_assoc = out_associations->create();
-        out_assoc.setRec(out_clust);
-        out_assoc.setSim(mc_par);
-        out_assoc.setWeight(in_assoc.getWeight());
+    if (link_nav.enabled()) {
+      for (const auto& [mc_par, weight] : link_nav.linked(in_clust)) {
+        truth::addWeightedRelation(
+            out_clust, mc_par, weight,
+            gsl::not_null<edm4eic::MCRecoClusterParticleLinkCollection*>{out_links},
+            gsl::not_null<edm4eic::MCRecoClusterParticleAssociationCollection*>{out_associations});
       }
-    } // end input association loop
+    } // end input link loop
   } // end input cluster loop
   debug("Completed processing input clusters");
 

src/algorithms/calorimetry/CalorimeterClusterShape.h

@@ -26,7 +26,7 @@ namespace eicrecon {
 // --------------------------------------------------------------------------
 using CalorimeterClusterShapeAlgorithm = algorithms::Algorithm<
     algorithms::Input<edm4eic::ClusterCollection,
-                      std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>>,
+                      std::optional<edm4eic::MCRecoClusterParticleLinkCollection>>,
     algorithms::Output<edm4eic::ClusterCollection,
                        std::optional<edm4eic::MCRecoClusterParticleLinkCollection>,
                        std::optional<edm4eic::MCRecoClusterParticleAssociationCollection>>>;
@@ -46,7 +46,7 @@ class CalorimeterClusterShape : public CalorimeterClusterShapeAlgorithm,
   CalorimeterClusterShape(std::string_view name)
       : CalorimeterClusterShapeAlgorithm{
             name,
-            {"inputClusters", "inputMCClusterAssociations"},
+            {"inputClusters", "inputMCClusterLinks"},
             {"outputClusters", "outputMCClusterLinks", "outputMCClusterAssociations"},
             "Computes cluster shape parameters"} {}