Convert PID meta information in standalone conversion

key4hep · Apr 12, 2024 · 3b946d3 · 3b946d3
1 parent 66898a3
commit 3b946d3
Show file tree

Hide file tree

Showing 4 changed files with 50 additions and 1 deletion.
diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.h
@@ -35,6 +35,7 @@ namespace edm4hep {
 #endif
 #include "edm4hep/TrackerHitPlaneCollection.h"
 #include "edm4hep/VertexCollection.h"
+#include "edm4hep/utils/ParticleIDUtils.h"
 
 // LCIO
 #include <EVENT/CalorimeterHit.h>
@@ -57,6 +58,7 @@ namespace edm4hep {
 #include <EVENT/LCIntVec.h>
 #include <EVENT/LCFloatVec.h>
 #include <UTIL/LCIterator.h>
+#include <UTIL/PIDHandler.h>
 #include <lcio.h>
 
 #include "podio/Frame.h"
@@ -171,6 +173,21 @@ namespace LCIO2EDM4hepConv {
 
   inline edm4hep::Vector3f Vector3fFrom(const EVENT::FloatVec& v) { return edm4hep::Vector3f(v[0], v[1], v[2]); }
 
+  /**
+   * Get the name of a ParticleID collection from the name of the reco
+   * collection (from which it is created) and the PID algorithm name.
+   */
+  inline std::string getPIDCollName(const std::string& recoCollName, const std::string& algoName)
+  {
+    return recoCollName + "_PID_" + algoName;
+  }
+
+  /**
+   * Get the meta information for all particle id collections that are available
+   * from the PIDHandler
+   */
+  std::vector<edm4hep::utils::ParticleIDMeta> getPIDMetaInfo(const UTIL::PIDHandler& pidHandler);
+
   /**
    * Convert a TrackState
    */

diff --git a/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp b/k4EDM4hep2LcioConv/include/k4EDM4hep2LcioConv/k4Lcio2EDM4hepConv.ipp
@@ -174,7 +174,7 @@ namespace LCIO2EDM4hepConv {
     results.reserve(particleIDs.size() + 1);
     results.emplace_back(name, std::move(dest));
     for (auto& [id, coll] : particleIDs) {
-      results.emplace_back(name + "_PID_" + pidHandler.getAlgorithmName(id), std::move(coll));
+      results.emplace_back(getPIDCollName(name, pidHandler.getAlgorithmName(id)), std::move(coll));
     }
     return results;
   }

diff --git a/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp b/k4EDM4hep2LcioConv/src/k4Lcio2EDM4hepConv.cpp
@@ -71,6 +71,16 @@ namespace LCIO2EDM4hepConv {
     return headerColl;
   }
 
+  std::vector<edm4hep::utils::ParticleIDMeta> getPIDMetaInfo(const UTIL::PIDHandler& pidHandler)
+  {
+    std::vector<edm4hep::utils::ParticleIDMeta> pidInfos {};
+    for (const auto id : pidHandler.getAlgorithmIDs()) {
+      pidInfos.emplace_back(pidHandler.getAlgorithmName(id), id, pidHandler.getParameterNames(id));
+    }
+
+    return pidInfos;
+  }
+
   podio::Frame convertEvent(EVENT::LCEvent* evt, const std::vector<std::string>& collsToConvert)
   {
     auto typeMapping = LcioEdmTypeMapping {};

diff --git a/standalone/lcio2edm4hep.cpp b/standalone/lcio2edm4hep.cpp
@@ -4,6 +4,8 @@
 #include <IOIMPL/LCFactory.h>
 #include <UTIL/CheckCollections.h>
 
+#include <edm4hep/utils/ParticleIDUtils.h>
+
 #include "podio/podioVersion.h"
 #if PODIO_BUILD_VERSION >= PODIO_VERSION(0, 99, 0)
 #include "podio/ROOTWriter.h"
@@ -166,6 +168,8 @@ int main(int argc, char* argv[])
 
   podio::ROOTWriter writer(args.outputFile);
 
+  podio::Frame metadata {};
+
   for (int j = 0; j < lcreader->getNumberOfRuns(); ++j) {
     if (j % 1 == 0) {
       std::cout << "processing RunHeader: " << j << std::endl;
@@ -187,9 +191,27 @@ int main(int argc, char* argv[])
       colPatcher.patchCollections(evt);
     }
     const auto edmEvent = LCIO2EDM4hepConv::convertEvent(evt, collsToConvert);
+
+    // For the first event we also convert some meta information for the
+    // ParticleID handling
+    if (i == 0) {
+      for (const auto& name : *evt->getCollectionNames()) {
+        auto coll = evt->getCollection(name);
+        if (coll->getTypeName() == "ReconstructedParticle") {
+          auto pidHandler = UTIL::PIDHandler(coll);
+          for (const auto& pidInfo : LCIO2EDM4hepConv::getPIDMetaInfo(pidHandler)) {
+            edm4hep::utils::PIDHandler::setAlgoInfo(
+              metadata, LCIO2EDM4hepConv::getPIDCollName(name, pidInfo.algoName), pidInfo);
+          }
+        }
+      }
+    }
+
     writer.writeFrame(edmEvent, "events");
   }
 
+  writer.writeFrame(metadata, podio::Category::Metadata);
+
   writer.finish();
 
   return 0;