Conversion of ParticleIDs of the Cluster Collection

key4hep · Jun 2, 2023 · 44e8ae0 · 44e8ae0
1 parent 366411d
commit 44e8ae0
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 7 deletions.
diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h
@@ -271,11 +271,16 @@ namespace LCIO2EDM4hepConv {
   /**
    * Convert a Cluster collection and return the resulting collection.
    * Simultaneously populates the mapping from LCIO to EDM4hep objects.
+   *
+   * NOTE: Also populates a ParticleID collection, as those are persisted as
+   * part of the Cluster collection in LCIO. The name of this collection is
+   * <name>_particleIDs
    */
-  std::unique_ptr<edm4hep::ClusterCollection> convertClusters(
+  std::vector<CollNamePair> convertClusters(
     const std::string& name,
     EVENT::LCCollection* LCCollection,
-    TypeMapT<const lcio::Cluster*, edm4hep::MutableCluster>& clusterMap);
+    TypeMapT<const lcio::Cluster*, edm4hep::MutableCluster>& clusterMap,
+    TypeMapT<const lcio::ParticleID*, edm4hep::MutableParticleID>& particleIDMap);
 
   /**
    * Create an EventHeaderCollection and fills it with the Metadata.

diff --git a/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp b/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp
@@ -457,11 +457,13 @@ namespace LCIO2EDM4hepConv {
     return dest;
   }
 
-  std::unique_ptr<edm4hep::ClusterCollection> convertClusters(
+  std::vector<CollNamePair> convertClusters(
     const std::string& name,
     EVENT::LCCollection* LCCollection,
-    TypeMapT<const lcio::Cluster*, edm4hep::MutableCluster>& clusterMap)
+    TypeMapT<const lcio::Cluster*, edm4hep::MutableCluster>& clusterMap,
+    TypeMapT<const lcio::ParticleID*, edm4hep::MutableParticleID>& particleIDMap)
   {
+    auto particleIDs = std::make_unique<edm4hep::ParticleIDCollection>();
     auto dest = std::make_unique<edm4hep::ClusterCollection>();
 
     for (unsigned i = 0, N = LCCollection->getNumberOfElements(); i < N; ++i) {
@@ -485,9 +487,26 @@ namespace LCIO2EDM4hepConv {
         std::cerr << "EDM4hep element  " << existingId << " did not get inserted. It belongs to the " << name
                   << " collection" << std::endl;
       }
-    }
 
-    return dest;
+      // Need to convert the particle IDs here, since they are part of the cluster
+      // collection in LCIO
+      for (const auto lcioPid : rval->getParticleIDs()) {
+        auto pid = convertParticleID(lcioPid);
+        const auto [pidIt, pidInserted] = particleIDMap.emplace(lcioPid, pid);
+        if (pidInserted) {
+          lval.addToParticleIDs(pid);
+          particleIDs->push_back(pid);
+        }
+        else {
+          lval.addToParticleIDs(pidIt->second);
+        }
+      }
+    }
+    std::vector<CollNamePair> results;
+    results.reserve(2);
+    results.emplace_back(name, std::move(dest));
+    results.emplace_back(name + "_particleIDs", std::move(particleIDs));
+    return results;
   }
 
   std::vector<CollNamePair>
@@ -508,7 +527,7 @@ namespace LCIO2EDM4hepConv {
       retColls.emplace_back(name, convertTracks(name, LCCollection, typeMapping.tracks));
     }
     else if (type == "Cluster") {
-      retColls.emplace_back(name, convertClusters(name, LCCollection, typeMapping.clusters));
+      return convertClusters(name, LCCollection, typeMapping.clusters, typeMapping.particleIDs);
     }
     else if (type == "SimCalorimeterHit") {
       retColls.emplace_back(name, convertSimCalorimeterHits(name, LCCollection, typeMapping.simCaloHits));